feat(customizer): build RL training image to support DPO+GRPO - #901
Conversation
|
dc44a1d to
8434602
Compare
8434602 to
8c0c013
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR relocates NeMo-RL image definitions into ChangesNeMo-RL image pipeline
Sequence Diagram(s)sequenceDiagram
participant DockerBake
participant RLBase
participant RLTraining
participant SmokeTest
DockerBake->>RLBase: build pinned nmp-rl-base
RLBase->>RLTraining: provide nmp-rl-base context
DockerBake->>RLTraining: build training runtime
DockerBake->>SmokeTest: build smoke-test stage with marker
SmokeTest->>SmokeTest: validate actor venv imports and cache links
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docker/rl/README.md (1)
79-79: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRemove the unmatched code fence.
This opens a fenced block at EOF without a matching opener.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/rl/README.md` at line 79, Remove the unmatched code fence at the end of the README so the document does not end with an unterminated fenced block.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docker/rl/Dockerfile.nmp-rl-base`:
- Around line 155-157: Pin the OpenSandbox and Tenacity runtime dependencies
installed by the Dockerfile command instead of using lower bounds. Use exact
versions sourced from the RL dependency set or a pinned constraints file, and
update the uv pip install step so repeatable builds no longer resolve different
package versions.
- Around line 179-181: Remove the `|| true` fallback from the RUN step invoking
`cve-cleanup.sh` so cleanup failures fail the image build. Update
`cve-cleanup.sh` to handle expected missing files or packages idempotently while
propagating unexpected command errors.
- Around line 18-42: Update the nemo-rl source stage to consume the advertised
named build context instead of always fetching via remote ADD. Configure the
corresponding Bake target to provide the default Git source for nemo-rl while
preserving the NEMO_RL_REPO and NEMO_RL_REF overrides, or remove the documented
local-context override if it cannot be supported.
- Around line 195-221: After the permission setup in the nmp-rl-base stage,
configure a verified non-root default user for the published image instead of
leaving it as root. Add the user selection after the chmod command so the
existing setup remains root-owned, and ensure it matches the UID 1000 runtime
user expected by the training image.
In `@docker/rl/README.md`:
- Around line 1-8: Add a prerequisites section at the top of the README before
the build graph, documenting the required setup and dependencies. Append a ##
Next Steps section at the end with links to the Bake target, plugin guide, and
smoke test, preserving the existing design-decision content between these
sections.
- Around line 1-5: Update the product-name references in the README heading and
introductory text to use the repository’s configured documentation substitutions
instead of hardcoded names. Preserve the existing meaning and formatting while
applying substitutions consistently to NeMo-RL, NeMo Platform, and NeMo-Gym
references as supported by the Sphinx configuration.
---
Outside diff comments:
In `@docker/rl/README.md`:
- Line 79: Remove the unmatched code fence at the end of the README so the
document does not end with an unterminated fenced block.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dccaf471-9bdd-46e8-a6d9-efdd24e2d5ee
📒 Files selected for processing (9)
docker-bake.hcldocker/Dockerfile.nmp-rl-basedocker/Dockerfile.nmp-rl-trainingdocker/rl/Dockerfile.nmp-rl-basedocker/rl/Dockerfile.nmp-rl-trainingdocker/rl/README.mdplugins/nemo-rl/README.mdtests/smoke_gpu/conftest.pytests/smoke_gpu/test_rl_training.py
💤 Files with no reviewable changes (2)
- docker/Dockerfile.nmp-rl-base
- docker/Dockerfile.nmp-rl-training
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
docker/rl/Dockerfile.nmp-rl-base (2)
33-34: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winPin
NEMO_RL_REFto an immutable commit.
nmp/customizeris a mutable branch, so identical builds can consume different source. Default to a full commit SHA.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/rl/Dockerfile.nmp-rl-base` around lines 33 - 34, Update the NEMO_RL_REF build argument in the Dockerfile to default to the full immutable commit SHA for the intended revision, replacing the mutable nmp/customizer branch while leaving NEMO_RL_REPO unchanged.
74-94: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winVerify toolchain downloads before execution.
The CMake archive and
uvinstaller execute without checksum or signature verification. Pin and validate their SHA-256 digests first.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/rl/Dockerfile.nmp-rl-base` around lines 74 - 94, Update the CMake download and uv installation flow to verify SHA-256 digests before extracting or executing either toolchain artifact. Add build arguments or equivalent pinned digest values for the CMake archive and uv installer, validate each downloaded file with the existing checksum tooling, and only continue when verification succeeds; preserve the current installation paths and version-based URLs.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docker/base/Dockerfile.python-wheels`:
- Around line 348-364: Replace the floating TE_REF branch reference with the
full commit SHA 42b840051647eef89761a16dfdff87e82bb253ab in the Dockerfile’s
TransformerEngine checkout, and update the corresponding TransformerEngine
reference in docker-bake.hcl to the same SHA.
In `@docker/rl/Dockerfile.nmp-rl-base`:
- Around line 328-333: Update the Git cleanup RUN command near the .git removal
step to remove the `2>/dev/null || true` fallback, allowing `find` or `rm`
failures to fail the image build while retaining the existing `.git` removal
behavior.
- Around line 221-224: The OpenSandbox and multimedia dependency installs
currently target only /opt/nemo_rl_venv, while actors use /opt/ray_venvs. Update
the actor-prefetch or worker-venv setup near the existing opensandbox and
av/opencv_python_headless/decord2 installation logic so /opt/ray_venvs receives
the same packages and wheel replacements.
In `@docker/rl/README.md`:
- Around line 60-61: Update the sentence in the README to use the possessive
“its” instead of “it's” and change “is” to “are” for subject-verb agreement,
preserving the referenced symbols and meaning.
- Line 40: Update all five fenced code blocks in the README, including the
blocks near the referenced sections, to specify an appropriate language
identifier such as text for the ASCII diagrams. Preserve the diagram contents
and formatting.
---
Outside diff comments:
In `@docker/rl/Dockerfile.nmp-rl-base`:
- Around line 33-34: Update the NEMO_RL_REF build argument in the Dockerfile to
default to the full immutable commit SHA for the intended revision, replacing
the mutable nmp/customizer branch while leaving NEMO_RL_REPO unchanged.
- Around line 74-94: Update the CMake download and uv installation flow to
verify SHA-256 digests before extracting or executing either toolchain artifact.
Add build arguments or equivalent pinned digest values for the CMake archive and
uv installer, validate each downloaded file with the existing checksum tooling,
and only continue when verification succeeds; preserve the current installation
paths and version-based URLs.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a3dfd466-83d7-43a3-befc-62e1ed80bdc0
📒 Files selected for processing (5)
docker-bake.hcldocker/base/Dockerfile.python-wheelsdocker/rl/Dockerfile.nmp-rl-basedocker/rl/Dockerfile.nmp-rl-trainingdocker/rl/README.md
🚧 Files skipped from review as they are similar to previous changes (2)
- docker/rl/Dockerfile.nmp-rl-training
- docker-bake.hcl
97edbf0 to
51db768
Compare
gabwow
left a comment
There was a problem hiding this comment.
Some option changes and questions. The smoke test and .git delete need properly addressed (either correcting me or fixing) before approval
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docker/rl/Dockerfile.nmp-rl-base (1)
401-407: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winExclude
.gitbefore copying a local source context.
NO_VCS_VERSION=1only skips version probing. A documented localnemo-rlcontext can still copy.gitinto the layer at line 258, exposing history/remotes. Exclude VCS entries during the sourceCOPY; do not remove them in a later layer. This also makestest_no_git_directories_shippedreliable for local builds.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/rl/Dockerfile.nmp-rl-base` around lines 401 - 407, Update the local nemo-rl source COPY instruction near the existing source copy to exclude .git directories before they enter the image layer, using the Docker build’s established ignore mechanism. Do not rely on deleting .git in a later layer, and preserve the NO_VCS_VERSION environment setting separately.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/smoke_gpu/test_rl_training.py`:
- Around line 37-57: Extend WORKER_VENV_IMPORTS with entries for the async Vllm
worker and SyncRolloutActor venvs, including each venv’s required import
packages. Ensure all five Docker-prefetched actor venvs are represented so smoke
tests validate every prefetch filter case.
---
Outside diff comments:
In `@docker/rl/Dockerfile.nmp-rl-base`:
- Around line 401-407: Update the local nemo-rl source COPY instruction near the
existing source copy to exclude .git directories before they enter the image
layer, using the Docker build’s established ignore mechanism. Do not rely on
deleting .git in a later layer, and preserve the NO_VCS_VERSION environment
setting separately.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f771cacf-0932-40e3-83fc-1dc64d337262
📒 Files selected for processing (3)
docker/rl/Dockerfile.nmp-rl-basedocker/rl/README.mdtests/smoke_gpu/test_rl_training.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docker/rl/README.md
gabwow
left a comment
There was a problem hiding this comment.
LGTM -- as a fast follow can we document the hardware restriction for mamba DPO and GRPO? I'm not sure how to surface this on our skills. If you have ideas, fold it into the docs change, otherwise let's just create a ticket to refine
Signed-off-by: anubhutiv <anubhutiv@nvidia.com>
08ae0fe to
be54770
Compare
Summary by CodeRabbit
Testing:
Built the image and can run the DPO job fine with image chnages in this PR and DPO restoration in this PR - #943