feat(customizer): add RL backend in customizer plugin - #479
Merged
Conversation
anubhutivyas
force-pushed
the
aalgo-292-dpo-support/anubhutiv
branch
from
June 25, 2026 23:17
cc4e23d to
bae7dad
Compare
Contributor
|
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:
📝 WalkthroughWalkthroughAdds NeMo-RL build targets, plugin submit wiring, service schemas, dataset and prompt compilation, training runtime support, and file I/O/model-entity task entrypoints. ChangesNeMo-RL support stack
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docker/Dockerfile.nmp-rl-base (1)
12-45: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winDon’t publish
nmp-rl-baseas root.
docker-bake.hcltags this stage directly, so this is a shipped image, not just an internal builder. Add a non-rootUSERhere and let child images switch back to root only for the install/chown steps that need it.🤖 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/Dockerfile.nmp-rl-base` around lines 12 - 45, The nmp-rl-base stage is being published directly, so it should not remain running as root. In the Dockerfile.nmp-rl-base stage, add a non-root USER before the image is finalized, and only switch back to root in downstream child images for the install or chown steps that require it; keep the existing VIRTUAL_ENV, chmod, and vllm removal logic in nmp-rl-base, but ensure the final shipped stage uses a non-root runtime user.Source: Linters/SAST tools
🟠 Major comments (36)
services/rl/src/nmp/rl/tasks/model_entity/run.py-154-163 (1)
154-163: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winLet transient fileset lookup failures reach the retry wrapper.
This
except Exceptionconverts API timeouts/500s fromfilesets.retrieve()intoModelEntityCreationError, so@retryoncreate_model_entity()never retries them. Only translate permanent lookup failures here.🤖 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 `@services/rl/src/nmp/rl/tasks/model_entity/run.py` around lines 154 - 163, The broad exception handling around fileset validation in create_model_entity() is wrapping transient SDK/API failures into ModelEntityCreationError, preventing the retry decorator from retrying them. Update the try/except around self.sdk.files.filesets.retrieve() so only permanent “not found” or access-denied cases are translated to ModelEntityCreationError, and let timeouts/5xx/other transient exceptions propagate out of the fileset validation block.services/rl/src/nmp/rl/tasks/file_io/run.py-141-181 (1)
141-181: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRecreate progress state on each retry.
DownloadStats/UploadStatsandCompositeCallbackare built once, then reused across tenacity retries. After a partial attempt,BaseSingleFileCallback.close()has already incremented counts/bytes, so the retry double-counts progress and can leave final totals above the real transfer. Move stats/callback construction inside the retried path or reset them per attempt.Also applies to: 218-251
🤖 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 `@services/rl/src/nmp/rl/tasks/file_io/run.py` around lines 141 - 181, The download progress state is reused across retries, so a partial attempt can leave DownloadStats and the callback chain carrying stale counts into the next try. Update the retried flow in run.py around the download path that builds CompositeCallback, TqdmPerFileDownloadCallback, and FileDownloadProgressCallback so these objects are created fresh per attempt or explicitly reset before each retry. Apply the same per-attempt reset pattern to the matching upload path referenced by the comment so progress totals stay accurate after tenacity retries.services/rl/src/nmp/rl/tasks/model_entity/run.py-424-428 (1)
424-428: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winDon’t log the raw model-entity config.
config.model_dump_json()will dump inlinedeployment_config.additional_envsverbatim. That can leak deployment secrets into task logs.🤖 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 `@services/rl/src/nmp/rl/tasks/model_entity/run.py` around lines 424 - 428, The startup logging in run() is dumping the full model-entity config, which can expose deployment secrets via deployment_config.additional_envs. Update the logger.info call that prints config.model_dump_json(indent=2) so it logs a redacted or summarized config instead, and ensure any sensitive nested fields are excluded or masked before logging.services/rl/src/nmp/rl/tasks/model_entity/run.py-148-169 (1)
148-169: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse the configured workspace for full-model outputs.
ModelEntityTaskConfigcarries a target workspace, but this path hardcodesself.job_ctx.workspaceand passes that intomodels.create()/models.update(). Cross-workspace submissions will create or update the wrong model entity.Also applies to: 226-260
🤖 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 `@services/rl/src/nmp/rl/tasks/model_entity/run.py` around lines 148 - 169, The full-model creation/update path is using the job context workspace instead of the configured target workspace, so cross-workspace runs write to the wrong place. Update the model entity flow in model_entity/run.py to use ModelEntityTaskConfig.workspace for the full-model output path, including the call chain through _create_or_update_full_entity and any models.create/models.update logic, while keeping the fileset workspace fallback separate.services/rl/src/nmp/rl/tasks/model_entity/run.py-236-242 (1)
236-242: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winHonor
config.trust_remote_code.The full-entity request always copies
base_me.trust_remote_code, so the task config flag never affects full or merged model entities. That can persist the wrong trust policy.🤖 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 `@services/rl/src/nmp/rl/tasks/model_entity/run.py` around lines 236 - 242, The full-entity request body is always using base_me.trust_remote_code, so the task-level trust policy from config is ignored. Update the request construction in the model entity run flow to use config.trust_remote_code for full/merged model entity requests, and make sure the request_body assembly in this section consistently honors the config flag instead of copying the base model value.services/rl/src/nmp/rl/schemas.py-26-39 (1)
26-39: 🗄️ Data Integrity & Integration | 🟠 MajorForbid extra fields in the nested training schemas.
ParallelismParamsand_TrainingBasestill accept unknown keys, so typos liketraining.learningrateorparallelism.num_nodeare silently dropped and the job falls back to defaults. Addmodel_config = ConfigDict(extra="forbid")to both models.🤖 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 `@services/rl/src/nmp/rl/schemas.py` around lines 26 - 39, The nested training schemas currently accept unknown keys, so typos in submitted config fields are silently ignored and defaults are used instead. Update both ParallelismParams and _TrainingBase in schemas.py to forbid extra fields by adding model_config = ConfigDict(extra="forbid"), ensuring invalid keys under training and parallelism are rejected during validation.services/rl/src/nmp/rl/entities/values.py-19-21 (1)
19-21: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winRemove unsupported
PPOfrom the enum.
TrainingStepConfig.training.training_typeaccepts this value, but the backend only compilesDPOandGRPOand raises on anything else. That turns validation into a runtime failure in the training container.Suggested fix
class TrainingType(str, Enum): """RL training algorithm. DPO is wired; the rest are headroom.""" DPO = "dpo" GRPO = "grpo" - PPO = "ppo"🤖 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 `@services/rl/src/nmp/rl/entities/values.py` around lines 19 - 21, The TrainingType enum in values.py still exposes PPO even though the backend only supports DPO and GRPO, so remove the unsupported PPO member from the enum. Update the TrainingType definition and any related validation or mappings that reference PPO so TrainingStepConfig.training.training_type only accepts the supported values and cannot reach a runtime failure in the training container.services/rl/src/nmp/rl/app/jobs/compiler.py-263-263 (1)
263-263: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winDon't log the full submitted spec at INFO.
Line 263 serializes the entire request, and the same
integrationsobject is later mined for secret envs. That creates a normal-path secret leak into job logs.Suggested fix
- logger.info("Compiling NeMo-RL DPO job to PlatformJobSpec: %s", job_spec.model_dump_json(indent=2)) + logger.info( + "Compiling NeMo-RL DPO job to PlatformJobSpec", + extra={"model": job_spec.model, "output_model": job_spec.output.name}, + )🤖 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 `@services/rl/src/nmp/rl/app/jobs/compiler.py` at line 263, The INFO log in compiler.py is serializing the full submitted job spec, which can leak secret values from the later-processed integrations payload. Update the logging in the job compilation flow around the compiler entry point to avoid dumping job_spec.model_dump_json(...) at INFO; instead log only non-sensitive high-level context (for example job identifiers or summary fields) and keep any detailed spec output out of normal-path logs.services/rl/src/nmp/rl/tasks/training/errors/parser.py-76-79 (1)
76-79: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winParse non-
*Errorexception names too.This regex never extracts
TimeoutExpired, so subprocess timeouts cannot hit theTrainingTimeoutErrorrule inerror_rules.yamland fall back to a generic runtime failure. Match generic exception class names and strip any module prefix.🤖 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 `@services/rl/src/nmp/rl/tasks/training/errors/parser.py` around lines 76 - 79, The exception parser in _EXCEPTION_RE only matches names ending in Error or Exception, so it misses generic exception types like TimeoutExpired. Update the regex in parser.py to capture any exception class name, and make sure the parsing logic strips any module prefix before applying the training error rules so TrainingTimeoutError can match subprocess timeout failures.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py-44-47 (1)
44-47: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDo not advertise GRPO from a stub.
NemoRLBackend.compile_config()routes GRPO jobs here, but this path still prints the full config/job context and then raisesNotImplementedError. Every GRPO submission fails here. Either implement it or reject GRPO before backend selection.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py` around lines 44 - 47, GRPO is still routed into a stub in NemoRLBackend.compile_config() via the GRPO config path, which prints the full training_config/job_ctx and then fails. Remove the stub behavior in the GRPO config compiler and either implement the actual GRPO config compilation in the relevant grpo_config handling or make backend selection reject GRPO before it reaches this path; also eliminate the debug prints from the stub path.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py-126-135 (1)
126-135: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winRestore the process env after each run.
These writes mutate global state, and
MLFLOW_URIis only set when present, never cleared. A reused worker can send a later job to the previous job’s tracking backend or run with stale settings. Build a per-run env or save/restore these keys infinally.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py` around lines 126 - 135, The env setup in the Nemo RL backend mutates global process state and leaves stale values behind between runs; update the run setup around the existing environment writes in backend.py to save and restore the affected keys, or construct a per-run env passed into the execution path. Make sure the logic around BASE_LOG_DIR, GPUS_PER_NODE, and especially MLFLOW_URI in the backend entrypoint is restored in a finally block so a reused worker cannot inherit prior run settings.services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml-415-418 (1)
415-418: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winNarrow the generic OOM matcher.
contains: "out of memory"will also match host/Ray/system OOMs, so those failures get classified asCudaErrorand users get GPU-specific remediation. Keep the CUDA-specific rules and drop or remap this catch-all.🤖 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 `@services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml` around lines 415 - 418, The generic OOM matcher in error_rules.yaml is too broad because the contains: "out of memory" rule can catch host, Ray, or system OOMs and misclassify them as CudaError. Narrow this rule by removing it or changing it to a CUDA-specific matcher, and keep the existing CUDA-focused rules so only true GPU OOMs map to CudaError; the relevant symbols are the OOM keyword rule under the training error rules and the CudaError mapping.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py-157-165 (1)
157-165: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winPut the original signal handlers back.
This installs process-wide SIGINT/SIGTERM handlers and never restores them. After
run_with_driver()returns, later signals in the same process still hit this stale closure and can terminate the wrong driver or raiseSystemExitunexpectedly.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py` around lines 157 - 165, The cleanup signal handler in bootstrap.run_with_driver setup is replacing the process-wide SIGINT/SIGTERM handlers and never restoring them, leaving a stale closure behind after the driver exits. Update the cleanup logic around the cleanup function and the signal.signal registrations to save the original handlers before installing them, then restore those original handlers after run_with_driver(str(driver_path), driver_args) completes (preferably in a finally block), so later signals do not call the old closure.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt-1-8 (1)
1-8: 🗄️ Data Integrity & Integration | 🟠 MajorThese markers don't pin base-image versions.
sys_platform == 'never'drops the requirement entirely, so it doesn't protectray,starlette,cryptography,mlflow, orprometheus-clientfrom being replaced if another dependency pulls them in.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt` around lines 1 - 8, The no_override_requirements markers are not preserving base-image packages because sys_platform == 'never' removes the requirement instead of constraining it. Update the entries in no_override_requirements.txt so each preserved package (prometheus-client, starlette, cryptography, mlflow, ray) is explicitly pinned or otherwise constrained in a way that prevents pip from upgrading/replacing the base-image version, using the same package names already listed in the file.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py-246-248 (1)
246-248: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winNamespace or clear
ENDEDper run.A stale
log_dir / "ENDED"from a previous retry makes new head/worker startup return immediately. Use a task/attempt-scoped marker path, or clear it with startup coordination before any rank checks it.Also applies to: 357-360, 392-395
🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py` around lines 246 - 248, The ENDED marker in ended_file is currently shared across retries because it always returns log_dir/"ENDED", so a stale file can short-circuit new runs. Update the marker handling in ray_bootstrap.py to make the path task/attempt-scoped (for example via the existing run/task identifiers used by the bootstrap flow) or ensure the marker is cleared during startup coordination before any rank checks it; apply the same fix to the related startup/check paths referenced by ended_file and the retry coordination logic.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py-55-104 (1)
55-104: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftDo not ship an executable driver that only raises
NotImplementedError.If this module is reachable via backend selection or
python -m, GRPO jobs fail immediately. Gate it out of runtime wiring or implement the driver before exposing it.I can help generate the implementation plan or open an issue for the missing GRPO driver.
🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py` around lines 55 - 104, The GRPO driver module is currently a stub and will fail at runtime because `load_config`, `get_environment`, `run_grpo_training`, and `main` all raise `NotImplementedError`. Either fully implement these entry points in `grpo_driver.py` or remove/gate the module from any backend selection and `python -m` wiring so it cannot be executed until complete. Make sure the runtime path that imports or dispatches to `main` only exposes a working driver, and use the existing function names to locate the missing implementation.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py-55-64 (1)
55-64: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winDo not dump full config or job context to stdout.
Hydra overrides, resolved configs, and
NMPJobContextcan contain internal URLs, identifiers, or credentials. Log redacted fields only, preferably at debug level.Also applies to: 93-94
🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py` around lines 55 - 64, The startup logging in dpo_driver should not print raw Hydra overrides or the fully resolved config to stdout. Update the logic around parse_hydra_overrides, OmegaConf.to_container, and the final pprint output so only redacted or minimal fields are logged, ideally through a debug logger instead of print. Also review the related job-context logging referenced in the same flow and ensure NMPJobContext values are sanitized before being emitted.services/rl/src/nmp/rl/tasks/training/distributed.py-91-99 (1)
91-99: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSerialize stale-barrier cleanup before workers continue.
Workers can create or wait on current-run markers before rank 0 executes
shutil.rmtree, so coordinator cleanup can delete live markers and deadlock barriers. Add a startup-ready marker, or use a per-attempt barrier directory instead of deleting a shared path.🤖 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 `@services/rl/src/nmp/rl/tasks/training/distributed.py` around lines 91 - 99, The stale-barrier cleanup in the distributed startup path can race with workers creating or waiting on current-run markers, so `shutil.rmtree` may delete live state. Update the initialization flow in `distributed.py` around `barrier_dir` handling in the `world_size > 1` block so cleanup is serialized before any worker proceeds, either by adding a startup-ready marker handshake or by switching `barrier_dir` to a per-attempt path instead of deleting a shared directory.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py-450-450 (1)
450-450: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winPut deadlines on Ray CLI subprocesses.
These calls can block indefinitely, making retry, worker-wait, and cleanup timeouts ineffective. Add
timeout=and handlesubprocess.TimeoutExpired; make cleanup unable to keep the process alive after its timeout.Also applies to: 485-485, 529-536, 676-680, 697-703, 707-710
🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py` at line 450, The Ray CLI subprocess calls in ray_bootstrap.py can block forever and bypass the existing retry/cleanup timeouts. Update the subprocess.run invocations in the relevant bootstrap/cleanup paths (including the call sites around result handling and the Ray cleanup helpers) to pass an explicit timeout, and add subprocess.TimeoutExpired handling so the code can abort and continue cleanup safely. Make sure the cleanup logic cannot leave the child process running after its timeout by terminating/killing it and waiting for it to exit, using the existing bootstrap and cleanup functions as the entry points to patch.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py-69-71 (1)
69-71: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDo not silently drop split overrides.
train_split/val_splitare public args, but the HuggingFace fallback ignores them. Pass them through if the base dataset supports them, or reject them whentrain_data_pathis absent.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py` around lines 69 - 71, The HuggingFace fallback in HelpSteer3Dataset currently ignores the public train_split and val_split args when train_data_path is not provided. Update the fallback path in the dataset initializer to either pass these split overrides through to the base dataset constructor if it supports them, or explicitly validate and reject unsupported overrides before calling super().__init__().services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py-71-73 (1)
71-73: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDo not silently drop split overrides.
train_split/val_splitare public args, but the HuggingFace fallback ignores them. Pass them through if the base dataset supports them, or reject them whentrain_data_pathis absent.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py` around lines 71 - 73, The HuggingFace fallback in the Tulu3 preference dataset constructor is ignoring the public train_split and val_split arguments by calling the base initializer without them. Update the fallback path around the super().__init__ call in the Tulu3 dataset class to either forward train_split and val_split to the base dataset if supported, or explicitly reject those overrides when train_data_path is not provided instead of silently dropping them.services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja-20-24 (1)
20-24: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSerialize tool metadata as JSON instead of concatenating it.
tnameortdesccontaining quotes/newlines produces invalid JSON in the prompt.Proposed fix
- {%- set tparams = t.function.parameters | tojson %} {{- "Use the function '" + tname + "' to '" + tdesc + "':\n" }} - {{- '{"name": "' + tname + '", "description": "' + tdesc + '", "parameters": ' + tparams + '}\n\n' }} + {{- ({"name": tname, "description": tdesc, "parameters": t.function.parameters} | tojson) + "\n\n" }}🤖 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 `@services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja` around lines 20 - 24, The tool metadata rendering in the llama-3.1-instruct template is building JSON by string concatenation, which breaks when tname or tdesc contains quotes/newlines. Update the template logic around the tname, tdesc, and tparams block to serialize the tool payload as proper JSON instead of manually assembling strings, and keep the existing function metadata fields intact in the rendered prompt.services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja-10-37 (1)
10-37: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winGuard undefined tool variables.
In Jinja, an undefined value is not
none, so this branch can emit the tool preamble even whentools/tool_choicewere never provided.Proposed fix
-{%- if tools is not none and tool_choice is not none %} +{%- if tools is defined and tools is not none and tool_choice is defined and tool_choice is not none %}🤖 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 `@services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja` around lines 10 - 37, The tool-preamble branch in llama-3.1-instruct.jinja only checks for none, so it can still render when tools or tool_choice are undefined. Update the conditional guarding the system/user preamble to explicitly handle undefined values before entering the tools loop, using the existing tools and tool_choice symbols, so the template only emits function-calling instructions when both inputs are actually provided.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py-357-372 (1)
357-372: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDetect dataset class per split before wiring NeMo-RL loaders.
Validation only proves each row matches the broad DPO union. Here the loader class is inferred from
prepared.train_fileonce and then reused for validation too, so mixed-format train shards or a validation file in another supported DPO format compile cleanly and then hit the wrong NeMo-RL loader at runtime. Detect the schema for each prepared file and reject mismatches, or emit split-specificdataset_namevalues if NeMo-RL supports that.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py` around lines 357 - 372, The dataset loader is currently inferred once from prepared.train_file and reused for validation, which can assign the wrong NeMo-RL class to a split. Update the DPO config logic around detect_dpo_schema_name and _dataset_spec to determine dataset_name separately for each prepared file (train and validation), and verify the split schemas are compatible before building the loader config. If NeMo-RL allows it, pass split-specific dataset_name values; otherwise reject mismatched formats early with a clear error.services/rl/src/nmp/rl/tasks/training/datasets/preparation.py-366-373 (1)
366-373: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winStream merges instead of loading each shard at once.
inp.read()pulls the whole dataset shard into memory before writing it back out. Large training files can OOM the worker during config preparation. Copy incrementally and only track whether the last line ended with\n.Suggested fix
with open(output_file, "w", encoding="utf-8") as out: for f in files: with open(f, "r", encoding="utf-8") as inp: - content = inp.read() - out.write(content) - # Ensure newline between files - if content and not content.endswith("\n"): + last_line: str | None = None + for last_line in inp: + out.write(last_line) + if last_line and not last_line.endswith("\n"): out.write("\n")🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/preparation.py` around lines 366 - 373, The merge logic in preparation.py currently uses the dataset shard read path in the write loop, which loads each file fully into memory before appending it to the output. Update the file-copying flow in the shard merge section to stream data incrementally instead of calling inp.read(), and use the existing output handling to track whether the last chunk ended with a newline so the separator behavior stays correct. Keep the change localized to the merge loop that iterates over files and writes to output_file.services/rl/src/nmp/rl/tasks/training/datasets/preparation.py-406-408 (1)
406-408: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winGuard tiny datasets before auto-splitting.
A 1-row file produces
val_size == 1, an empty train split, andcompile_dpo_config()later fails on the generated empty train file. Clamp validation to at mostlen(lines) - 1and fail early when the file is too small to split.Suggested fix
with open(train_file, "r", encoding="utf-8") as f: lines = [line for line in f if line.strip()] + + if len(lines) < 2: + raise DatasetFormatError("Need at least 2 samples to create a validation split") # Shuffle for reproducibility (important for multi-node!) # Uses global seed if not explicitly provided random.seed(seed) random.shuffle(lines) - val_size = max(1, int(len(lines) * val_ratio)) + val_size = min(max(1, int(len(lines) * val_ratio)), len(lines) - 1)🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/preparation.py` around lines 406 - 408, Guard the auto-split logic in preparation.py so tiny datasets do not produce an empty train file. In the code that sets val_size and builds val_lines/train_lines, clamp the validation count to at most len(lines) - 1 and add an early failure path when len(lines) is too small to split safely. Use the existing split logic around val_size, val_lines, and train_lines to keep the train output non-empty before compile_dpo_config() consumes it.services/rl/src/nmp/rl/tasks/training/datasets/schemas.py-292-305 (1)
292-305: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject empty
tool_calls.
tool_calls=[]passes this validator, butservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinjaLines 28-41 treat any non-Nonevalue as a tool-call block and never close it when the list is empty. That yields malformed prompts.🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/schemas.py` around lines 292 - 305, The SFTChatMessage validator in check_has_content_or_thinking_or_tool_calls currently accepts tool_calls=[] because it only checks for None; update it to reject empty tool call lists so only non-empty tool_calls are allowed. Keep the existing content/thinking rules intact, and make sure the validation logic in SFTChatMessage treats tool_calls the same way the nemotron-3.1.jinja template does when deciding whether a tool-call block should exist.services/rl/src/nmp/rl/tasks/training/datasets/schemas.py-366-369 (1)
366-369: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRequire at least one chat message.
messages=[]validates here, butservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinjaLines 1-6 andservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinjaLines 2-7 indexmessages[0]immediately. Empty rows pass validation and then crash during template rendering.🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/schemas.py` around lines 366 - 369, Require at least one chat message in the schema used by the Nemotron training dataset. Update the messages field in the relevant model/class in schemas.py so empty lists are rejected during validation, matching the assumptions in the nemotron-3.1.jinja and nemotron-3.3.jinja templates that access messages[0]. Keep the fix localized to the schema definition for the conversation object so invalid rows fail early instead of crashing during template rendering.services/rl/src/nmp/rl/tasks/training/datasets/schemas.py-268-271 (1)
268-271: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
thinkingis accepted without any renderer.This schema allows messages with
thinkingand nocontent, butnemotron-3.1.jinjaLines 42-46,nemotron-3.3.jinjaLines 12-17, andphi-4.jinjaLines 6-10 all rendermessage['content']only. Valid rows will render incorrectly or fail.Also applies to: 292-305
🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/schemas.py` around lines 268 - 271, The message schema currently allows `thinking` without any rendering support, so update the dataset/message handling to either reject rows where `thinking` is present but `content` is missing or explicitly map `thinking` into the rendered prompt path. Make the fix in the message schema/validation flow and in the Jinja renderers (`nemotron-3.1.jinja`, `nemotron-3.3.jinja`, `phi-4.jinja`) so they use the same source field consistently rather than only `message['content']`; use the `thinking` and message schema definitions in `schemas.py` to locate the affected logic.services/rl/src/nmp/rl/tasks/training/datasets/schemas.py-351-351 (1)
351-351: 🎯 Functional Correctness | 🟠 MajorFix the chat schema name/tag mismatch.
get_sft_dataset_discriminator()returnsSFTChatDatasetItemSchemaformessages, but the class and union tag areSFTPChatDatasetItemSchema, so chat rows won’t resolve to that branch. Update the class/tag names to match.🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/schemas.py` at line 351, The chat dataset discriminator and schema class name are mismatched, so `get_sft_dataset_discriminator()` cannot resolve `messages` rows to the intended branch. Update the `SFTPChatDatasetItemSchema` class and its union/tag usage in `schemas.py` so the name and discriminator key match exactly with `SFTChatDatasetItemSchema`, and make sure any references in the dataset schema union use the corrected identifier consistently.services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja-17-20 (1)
17-20: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse real newlines here.
Line 17 and Line 20 emit
\\n\\n, while Line 8 uses\n\n. That writes literal backslashes into the prompt and changes the expected chat format.Suggested fix
- {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' + content | trim + '<|eot_id|>' }} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + content | trim + '<|eot_id|>' }} @@ - {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}🤖 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 `@services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja` around lines 17 - 20, In the Nemotron 3.3 Jinja template, the message rendering block is emitting literal backslash characters because the assistant/user header template uses escaped newline text instead of actual newlines. Update the prompt construction around the message loop and the add_generation_prompt branch so it matches the existing newline style used earlier in the template, and ensure the rendered chat format from the template stays consistent across the header/content separators.services/rl/src/nmp/rl/tasks/training/utils.py-69-77 (1)
69-77: 🩺 Stability & Availability | 🟠 MajorGuard against zero GPUs.
torch.cuda.device_count()can return0when CUDA is unavailable, so this branch can emit--nproc_per_node 0. Treat<= 0as missing and fall back to1.🤖 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 `@services/rl/src/nmp/rl/tasks/training/utils.py` around lines 69 - 77, The GPU count fallback in the training utils can currently pass through a zero value from torch.cuda.device_count(), which leads to an invalid nproc_per_node setting. Update the logic in the GPU detection block around gpus_per_node so that any value less than or equal to 0 is treated as unavailable, and fall back to "1" just like the exception path. Keep the fix localized to the existing os.environ.get("GPUS_PER_NODE") and torch.cuda.device_count() handling.services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py-42-49 (1)
42-49: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winValidate
steps_per_epochandlog_intervalin__init__.Lines 100 and 105 divide/modulo by constructor inputs with no guard. A zero or negative value crashes the first metrics update. Reject invalid values before training starts.
Suggested fix
def __init__( self, steps_per_epoch: int, job_ctx: NMPJobContext | None = None, log_interval: int = 10, max_steps: int | None = None, num_epochs: int | None = None, ): @@ + if steps_per_epoch <= 0: + raise ValueError("steps_per_epoch must be > 0") + if log_interval <= 0: + raise ValueError("log_interval must be > 0") + self._job_ctx = job_ctx or NMPJobContext.from_env()Also applies to: 97-105
🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py` around lines 42 - 49, Validate the constructor inputs in NemoRLLogger.__init__: reject zero or negative values for steps_per_epoch and log_interval before storing them, since methods like _log_metrics and the step/update logic later perform division and modulo using these fields. Add an early guard in __init__ that raises a clear error for invalid values so training fails fast instead of crashing on the first metrics update.services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja-19-25 (1)
19-25: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSerialize the tool schema with one
tojsoncall.Line 24 hand-builds JSON with raw
tnameandtdesc. Any quote or newline in those fields makes the tool definition invalid and corrupts the prompt format. Serialize a dict instead of concatenating JSON fragments.Suggested fix
{%- for t in tools %} {%- set tname = t.function.name %} {%- set tdesc = t.function.description %} - {%- set tparams = t.function.parameters | tojson %} {{- "Use the function '" + tname + "' to '" + tdesc + "':\n" }} - {{- '{"name": "' + tname + '", "description": "' + tdesc + '", "parameters": ' + tparams + '}\n\n' }} + {{- {"name": tname, "description": tdesc, "parameters": t.function.parameters} | tojson + '\n\n' }} {%- endfor %}🤖 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 `@services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja` around lines 19 - 25, The tool schema rendering in the llama-3.2-instruct.jinja template is manually concatenating JSON fragments with raw tname and tdesc, which can break on quotes or newlines; update the loop over tools to serialize a single dictionary with one tojson call instead of building the JSON string by hand, while keeping the existing t.function.name, t.function.description, and t.function.parameters values in the output.services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja-19-25 (1)
19-25: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSerialize the tool schema with one
tojsoncall.Line 24 hand-builds JSON with raw
tnameandtdesc. Any quote or newline in those fields makes the tool definition invalid and corrupts the prompt format. Serialize a dict instead of concatenating JSON fragments.Suggested fix
{%- for t in tools %} {%- set tname = t.function.name %} {%- set tdesc = t.function.description %} - {%- set tparams = t.function.parameters | tojson %} {{- "Use the function '" + tname + "' to '" + tdesc + "':\n" }} - {{- '{"name": "' + tname + '", "description": "' + tdesc + '", "parameters": ' + tparams + '}\n\n' }} + {{- {"name": tname, "description": tdesc, "parameters": t.function.parameters} | tojson + '\n\n' }} {%- endfor %}🤖 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 `@services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja` around lines 19 - 25, The tool schema serialization in the llama-3.3-instruct.jinja template is being hand-built with raw tname and tdesc, which can break JSON when those values contain quotes or newlines. Update the loop over tools to construct a single dictionary-like payload for each tool and serialize that whole object with one tojson call, using the existing t.function.name, t.function.description, and t.function.parameters fields. Keep the surrounding prompt text generation in the same template, but remove the concatenated JSON fragment so the tool definition stays valid and robust.services/rl/src/nmp/rl/tasks/training/errors/converter.py-89-102 (1)
89-102: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winKeep
create_error_details()non-throwing.Line 90 calls
get_error_converter()outside the fallback guard. If rule loading fails here, the original training exception escapes the error path andTrainingRunner.run()falls back to its stale default result instead of structured error details. Wrap converter acquisition in the sametryand use theInternalErrorfallback from that branch too.Suggested fix
def create_error_details(exception: Exception) -> ErrorDetails: @@ - converter = get_error_converter() try: + converter = get_error_converter() converter.raise_converted_or_default(exception) except CustomizerTrainingError as converted: return converted.to_error_details() except Exception as e: # noqa: BLE001 - intentional last-resort guard to guarantee dict return # Unexpected exception type - wrap in InternalError to ensure we always return a dict logger.warning(f"Unexpected exception type from converter: {type(e).__name__}: {e}") exc = InternalError( message=f"An internal error occurred. ({type(exception).__name__}: {exception})", - detail=str(exception), + detail=f"{type(exception).__name__}: {exception}", ) return exc.to_error_details()🤖 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 `@services/rl/src/nmp/rl/tasks/training/errors/converter.py` around lines 89 - 102, Keep create_error_details() non-throwing by moving get_error_converter() inside the same try block as converter.raise_converted_or_default() in converter.py. If converter acquisition fails, handle it in the existing except Exception branch and return InternalError.to_error_details() so TrainingRunner.run() always gets structured error details instead of leaking the original exception. Use the existing create_error_details(), get_error_converter(), and InternalError flow as the fix point.
🟡 Minor comments (2)
plugins/nemo-rl/src/nemo_rl_plugin/schema.py-36-41 (1)
36-41: 🗄️ Data Integrity & Integration | 🟡 MinorCap
OutputRequest.nameat 255 chars
OutputResponsealready limits bothnameandfilesetto 255, so an oversized explicit output name should be rejected here instead of during transform.Suggested fix
class OutputRequest(BaseModel): """Submitter-facing output preferences. ``name`` is auto-derived if omitted.""" model_config = ConfigDict(extra="forbid") - name: str | None = None + name: str | None = Field(default=None, max_length=255)🤖 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 `@plugins/nemo-rl/src/nemo_rl_plugin/schema.py` around lines 36 - 41, OutputRequest.name is not bounded, so an explicit oversized output name can slip through until later transform logic. Add a 255-character max constraint to the name field in OutputRequest (the model in schema.py) so validation fails immediately, matching the limits already enforced by OutputResponse for name and fileset. Keep the change localized to the OutputRequest model definition.services/rl/src/nmp/rl/tasks/training/runner.py-81-81 (1)
81-81: 📐 Maintainability & Code Quality | 🟡 MinorUse
Selffor__enter__. Inservices/rl/src/nmp/rl/tasks/training/runner.py:81, replace the string return type withSelfand import it fromtyping.🤖 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 `@services/rl/src/nmp/rl/tasks/training/runner.py` at line 81, The __enter__ method in TrainingRunner is still using a string annotation instead of the modern Self type. Update the TrainingRunner.__enter__ signature to return Self, and add Self to the typing import list in the same module so the annotation is resolved correctly.Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 118f7c09-f7be-40f2-bcb5-727036cd8217
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (75)
docker-bake.hcldocker/Dockerfile.nmp-rl-basedocker/Dockerfile.nmp-rl-tasksdocker/Dockerfile.nmp-rl-trainingdocker/rl/Dockerfile.platform-workspacedocker/rl/pyproject.workspace.tomlpackages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.pyplugins/nemo-rl/README.mdplugins/nemo-rl/pyproject.tomlplugins/nemo-rl/src/nemo_rl_plugin/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/main.pyplugins/nemo-rl/src/nemo_rl_plugin/config.pyplugins/nemo-rl/src/nemo_rl_plugin/contributor.pyplugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.pyplugins/nemo-rl/src/nemo_rl_plugin/schema.pyplugins/nemo-rl/src/nemo_rl_plugin/sdk/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.pyplugins/nemo-rl/src/nemo_rl_plugin/transform.pypyproject.tomlservices/rl/README.mdservices/rl/pyproject.tomlservices/rl/src/nmp/rl/app/constants.pyservices/rl/src/nmp/rl/app/jobs/compiler.pyservices/rl/src/nmp/rl/app/jobs/training/schemas.pyservices/rl/src/nmp/rl/compile.pyservices/rl/src/nmp/rl/config.pyservices/rl/src/nmp/rl/entities/values.pyservices/rl/src/nmp/rl/images.pyservices/rl/src/nmp/rl/schemas.pyservices/rl/src/nmp/rl/tasks/file_io/__init__.pyservices/rl/src/nmp/rl/tasks/file_io/__main__.pyservices/rl/src/nmp/rl/tasks/file_io/callbacks.pyservices/rl/src/nmp/rl/tasks/file_io/run.pyservices/rl/src/nmp/rl/tasks/model_entity/__init__.pyservices/rl/src/nmp/rl/tasks/model_entity/__main__.pyservices/rl/src/nmp/rl/tasks/model_entity/run.pyservices/rl/src/nmp/rl/tasks/training/__main__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txtservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.pyservices/rl/src/nmp/rl/tasks/training/chat_templates.pyservices/rl/src/nmp/rl/tasks/training/datasets/preparation.pyservices/rl/src/nmp/rl/tasks/training/datasets/schemas.pyservices/rl/src/nmp/rl/tasks/training/datasets/validation.pyservices/rl/src/nmp/rl/tasks/training/distributed.pyservices/rl/src/nmp/rl/tasks/training/errors/converter.pyservices/rl/src/nmp/rl/tasks/training/errors/error_rules.yamlservices/rl/src/nmp/rl/tasks/training/errors/exceptions.pyservices/rl/src/nmp/rl/tasks/training/errors/parser.pyservices/rl/src/nmp/rl/tasks/training/integrations.pyservices/rl/src/nmp/rl/tasks/training/protocol.pyservices/rl/src/nmp/rl/tasks/training/runner.pyservices/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinjaservices/rl/src/nmp/rl/tasks/training/templates/phi-4.jinjaservices/rl/src/nmp/rl/tasks/training/utils.pyservices/rl/tests/test_compiler.pyservices/rl/tests/test_dpo_config.pyservices/rl/tests/test_schemas.py
anubhutivyas
force-pushed
the
aalgo-292-dpo-support/anubhutiv
branch
from
June 29, 2026 08:20
bae7dad to
0e09e28
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 13
🧹 Nitpick comments (1)
services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py (1)
17-17: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDrop
from __future__ import annotationsinservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py. The package targets Python 3.11+, and this module has no forward refs or string annotations, so the import is unnecessary.🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py` at line 17, Remove the unnecessary future import from ray_bootstrap.py by deleting the top-level from __future__ import annotations statement; this module does not use forward references or stringified type hints, and the codebase targets Python 3.11+, so no other changes should be needed.Source: Coding guidelines
🤖 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/Dockerfile.nmp-rl-base`:
- Around line 12-44: The nmp-rl-base stage is published as a final image but
currently never switches off root, so add a non-root user context before the
stage is published. Update the Dockerfile in the nmp-rl-base stage to switch
ownership/permissions as needed and set a non-root USER near the end of the
stage, using the existing runtime assumption that training runs as USER 1000.
In `@services/rl/src/nmp/rl/tasks/model_entity/run.py`:
- Around line 119-123: The model_entity parsing in run.py is too permissive
because the current split-and-index logic accepts extra path segments and
silently drops them. Update the parsing around the model_entity split so it only
allows exactly one segment for name or exactly two segments for workspace/name,
and reject any input with more than two parts before assigning me_workspace and
me_name.
- Line 378: Stop logging raw task and deployment config objects because they may
include secrets in job logs. Update the logging in the run flow around the
Deployment config output and the related task/config logs near the same section
to redact sensitive fields like deployment_config.additional_envs, or replace
the dumps with identifiers and summary fields only. Use the existing run/logger
calls in run.py to keep the context but ensure no full config payloads are
emitted.
- Around line 155-163: The fileset validation in run.py is catching every
exception and wrapping transient SDK failures into ModelEntityCreationError,
which prevents create_model_entity() from retrying. Update the fileset lookup in
the model entity creation flow to only translate genuine “not found” or access
errors, and let timeout/5xx/transient SDK exceptions propagate so the existing
`@retry` can handle them. Use the existing fileset workspace/name validation block
around self.sdk.files.filesets.retrieve() and keep the success/error logging,
but narrow the exception handling to specific non-retryable cases.
In `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py`:
- Line 235: The scheduler setup in DPO config is using total_steps instead of
the actual capped max_steps, which can stretch warmup/decay past the real
training run. Update the scheduler construction in the DPO config path that
calls _build_scheduler_config so it is based on the effective max_steps value
used for training, and make the same adjustment in the related config/build
sites at the other referenced locations to keep all scheduler definitions
consistent.
- Line 149: The final validation toggle in the DPO config is allowing the
end-of-training validation pass to be skipped even when checkpoint ranking
depends on validation loss. Update the logic in the DPO config builder so that
the value assigned to val_at_end is forced on whenever the checkpoint-ranker
path requires val:validation-default_loss, keeping the behavior consistent with
the earlier validation settings in the same config. Use the DPO config
construction and the val_at_end setting as the main symbols to locate and adjust
the conditional.
In `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py`:
- Line 157: The ENDED marker is currently shared across runs because the default
log_dir in the Ray bootstrap path points to a fixed logs directory, so stale
markers can affect new or unrelated workers. Update the startup/cleanup flow in
ray_bootstrap.py so the marker path is namespaced per run using the existing
run-specific identifiers from the bootstrap context (for example in the setup
around the log_dir field and the ENDED marker handling near the referenced
lifecycle methods), and make sure cleanup removes the run-scoped marker
consistently.
- Line 450: The Ray CLI subprocess calls in ray_bootstrap can hang forever
because subprocess.run is invoked without a timeout and _cleanup_with_timeout()
only limits the wrapper thread. Update the relevant subprocess.run usages in the
Ray stop/status/memory cleanup paths (including the call sites around the ray
bootstrap helpers) to pass a real timeout and handle the timeout exception so a
stuck ray command cannot block shutdown. Keep the fix localized to the existing
cleanup/command-execution flow and apply it consistently to each subprocess.run
call referenced in the diff.
In `@services/rl/src/nmp/rl/tasks/training/chat_templates.py`:
- Around line 169-195: The helper that updates tokenizer_config.json only writes
chat_template and never backfills pad_token, which contradicts the documented
fallback behavior. In apply_chat_template_to_tokenizer_config, after loading the
config and before writing it back, ensure pad_token is set when missing by using
eos_token as the fallback, while preserving any existing pad_token value. Keep
the update logic localized to the same tokenizer_config handling block.
In `@services/rl/src/nmp/rl/tasks/training/datasets/preparation.py`:
- Around line 39-56: The dataset discovery heuristics in TRAIN_PATTERNS and
VAL_PATTERNS are too broad because downstream code like detect_dataset_schema(),
_merge_files(), and _create_val_split() only handle JSONL line-by-line. Restrict
discovery in preparation.py to JSONL files and JSONL-oriented directories only,
or add real JSON parsing before accepting .json files; update any related
validation/merge logic so normal JSON arrays/objects are not treated as
line-delimited input.
- Around line 181-190: The non-empty line counting in `_count_jsonl_samples` is
treating every `subprocess.CalledProcessError` from `grep` as an empty file,
which hides real failures. Update the `grep` handling to distinguish the
expected no-match case from actual grep errors: only return 0 when the command
indicates no matches, and otherwise fall back to `_count_jsonl_samples_python`
or re-raise for the `subprocess.check_output` path in `preparation.py`.
In `@services/rl/src/nmp/rl/tasks/training/datasets/validation.py`:
- Around line 197-204: Schema detection and JSONL validation treat blank lines
inconsistently in the validation helpers. Update the logic around the first-row
read in the relevant helpers (including the code near the existing json.loads
path and the related schema-detection/validation functions) to skip
empty/whitespace-only lines and use the first non-empty row instead, and raise a
DatasetFormatError if no non-empty line exists. Ensure the behavior is
consistent wherever the file is parsed so leading blank lines do not break
detection and blank-only files do not pass validation.
- Line 203: The debug/error logging in the dataset validation flow is leaking
raw training payloads by including line and object contents on parse/schema
failures. Update the logging around the validation helpers that emit these
messages (including the JSON parse failure and the schema validation failure
paths) so they only log the file path, row number, and exception details, and
remove any references to line, obj, or other raw customer data from
logger.debug/logger.error calls.
---
Nitpick comments:
In `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py`:
- Line 17: Remove the unnecessary future import from ray_bootstrap.py by
deleting the top-level from __future__ import annotations statement; this module
does not use forward references or stringified type hints, and the codebase
targets Python 3.11+, so no other changes should be needed.
🪄 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: 1944ca97-554b-4279-b64d-f7b374ec31f4
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (75)
docker-bake.hcldocker/Dockerfile.nmp-rl-basedocker/Dockerfile.nmp-rl-tasksdocker/Dockerfile.nmp-rl-trainingdocker/rl/Dockerfile.platform-workspacedocker/rl/pyproject.workspace.tomlpackages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.pyplugins/nemo-rl/README.mdplugins/nemo-rl/pyproject.tomlplugins/nemo-rl/src/nemo_rl_plugin/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/main.pyplugins/nemo-rl/src/nemo_rl_plugin/config.pyplugins/nemo-rl/src/nemo_rl_plugin/contributor.pyplugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.pyplugins/nemo-rl/src/nemo_rl_plugin/schema.pyplugins/nemo-rl/src/nemo_rl_plugin/sdk/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.pyplugins/nemo-rl/src/nemo_rl_plugin/transform.pypyproject.tomlservices/rl/README.mdservices/rl/pyproject.tomlservices/rl/src/nmp/rl/app/constants.pyservices/rl/src/nmp/rl/app/jobs/compiler.pyservices/rl/src/nmp/rl/app/jobs/training/schemas.pyservices/rl/src/nmp/rl/compile.pyservices/rl/src/nmp/rl/config.pyservices/rl/src/nmp/rl/entities/values.pyservices/rl/src/nmp/rl/images.pyservices/rl/src/nmp/rl/schemas.pyservices/rl/src/nmp/rl/tasks/file_io/__init__.pyservices/rl/src/nmp/rl/tasks/file_io/__main__.pyservices/rl/src/nmp/rl/tasks/file_io/callbacks.pyservices/rl/src/nmp/rl/tasks/file_io/run.pyservices/rl/src/nmp/rl/tasks/model_entity/__init__.pyservices/rl/src/nmp/rl/tasks/model_entity/__main__.pyservices/rl/src/nmp/rl/tasks/model_entity/run.pyservices/rl/src/nmp/rl/tasks/training/__main__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txtservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.pyservices/rl/src/nmp/rl/tasks/training/chat_templates.pyservices/rl/src/nmp/rl/tasks/training/datasets/preparation.pyservices/rl/src/nmp/rl/tasks/training/datasets/schemas.pyservices/rl/src/nmp/rl/tasks/training/datasets/validation.pyservices/rl/src/nmp/rl/tasks/training/distributed.pyservices/rl/src/nmp/rl/tasks/training/errors/converter.pyservices/rl/src/nmp/rl/tasks/training/errors/error_rules.yamlservices/rl/src/nmp/rl/tasks/training/errors/exceptions.pyservices/rl/src/nmp/rl/tasks/training/errors/parser.pyservices/rl/src/nmp/rl/tasks/training/integrations.pyservices/rl/src/nmp/rl/tasks/training/protocol.pyservices/rl/src/nmp/rl/tasks/training/runner.pyservices/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinjaservices/rl/src/nmp/rl/tasks/training/templates/phi-4.jinjaservices/rl/src/nmp/rl/tasks/training/utils.pyservices/rl/tests/test_compiler.pyservices/rl/tests/test_dpo_config.pyservices/rl/tests/test_schemas.py
💤 Files with no reviewable changes (17)
- services/rl/src/nmp/rl/tasks/training/templates/phi-4.jinja
- services/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinja
- services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja
- services/rl/tests/test_schemas.py
- services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja
- services/rl/src/nmp/rl/tasks/training/errors/parser.py
- services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinja
- services/rl/src/nmp/rl/tasks/training/errors/exceptions.py
- services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml
- services/rl/src/nmp/rl/tasks/training/utils.py
- services/rl/src/nmp/rl/tasks/training/protocol.py
- services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja
- services/rl/src/nmp/rl/tasks/training/runner.py
- services/rl/tests/test_dpo_config.py
- services/rl/src/nmp/rl/tasks/training/integrations.py
- services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja
- services/rl/tests/test_compiler.py
✅ Files skipped from review due to trivial changes (9)
- plugins/nemo-rl/src/nemo_rl_plugin/sdk/init.py
- plugins/nemo-rl/src/nemo_rl_plugin/init.py
- services/rl/pyproject.toml
- services/rl/src/nmp/rl/tasks/model_entity/init.py
- services/rl/src/nmp/rl/tasks/file_io/init.py
- plugins/nemo-rl/src/nemo_rl_plugin/cli/init.py
- services/rl/README.md
- plugins/nemo-rl/README.md
- docker/rl/pyproject.workspace.toml
🚧 Files skipped from review as they are similar to previous changes (38)
- services/rl/src/nmp/rl/tasks/training/main.py
- services/rl/src/nmp/rl/tasks/model_entity/main.py
- services/rl/src/nmp/rl/tasks/file_io/main.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/init.py
- plugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py
- plugins/nemo-rl/src/nemo_rl_plugin/config.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt
- plugins/nemo-rl/src/nemo_rl_plugin/transform.py
- plugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.py
- services/rl/src/nmp/rl/compile.py
- plugins/nemo-rl/src/nemo_rl_plugin/contributor.py
- plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py
- services/rl/src/nmp/rl/config.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py
- plugins/nemo-rl/src/nemo_rl_plugin/cli/main.py
- services/rl/src/nmp/rl/entities/values.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py
- services/rl/src/nmp/rl/app/jobs/training/schemas.py
- plugins/nemo-rl/src/nemo_rl_plugin/schema.py
- pyproject.toml
- services/rl/src/nmp/rl/images.py
- packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py
- docker-bake.hcl
- services/rl/src/nmp/rl/tasks/training/errors/converter.py
- services/rl/src/nmp/rl/app/constants.py
- services/rl/src/nmp/rl/app/jobs/compiler.py
- plugins/nemo-rl/pyproject.toml
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py
- services/rl/src/nmp/rl/schemas.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py
- services/rl/src/nmp/rl/tasks/training/distributed.py
- services/rl/src/nmp/rl/tasks/training/datasets/schemas.py
- services/rl/src/nmp/rl/tasks/file_io/run.py
- services/rl/src/nmp/rl/tasks/file_io/callbacks.py
albcui
requested changes
Jun 29, 2026
anubhutivyas
force-pushed
the
aalgo-292-dpo-support/anubhutiv
branch
from
June 30, 2026 07:12
0e09e28 to
9624583
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (1)
services/rl/src/nmp/rl/tasks/training/distributed.py (1)
73-74: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a concrete return annotation.
-> "DistributedContext"breaks the repo rule against string-based Python type hints. As per coding guidelines,**/*.py: Always prefer concrete type hints over string-based ones in Python code.Diff
+from __future__ import annotations @@ - def from_env(cls, barrier_dir: Path) -> "DistributedContext": + def from_env(cls, barrier_dir: Path) -> DistributedContext:🤖 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 `@services/rl/src/nmp/rl/tasks/training/distributed.py` around lines 73 - 74, The from_env classmethod on DistributedContext uses a string-based return annotation, which violates the Python typing guideline. Update the signature of DistributedContext.from_env to use a concrete return type annotation instead of a quoted forward reference, keeping the method behavior unchanged.Source: Coding guidelines
🤖 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 `@services/rl/src/nmp/rl/entities/values.py`:
- Around line 16-25: Remove GRPO from TrainingType so the validation layer stays
aligned with the backend; the current TrainingType enum in values.py should
remain DPO-only until NemoRLBackend.compile_config() and _get_driver_path()
support GRPO. If GRPO must stay in the enum for future work, add schema-level
rejection or equivalent validation before scheduling so jobs cannot be accepted
and then fail later.
In `@services/rl/src/nmp/rl/tasks/model_entity/run.py`:
- Around line 104-106: The early return in run() is too broad because it accepts
any truthy target.spec, but NIM profile selection requires spec.family and
spec.base_num_parameters to be present. Update the spec-populated check in the
target handling logic to wait until both required fields are set before logging
and returning, and keep the existing flow in run() and related target/spec
handling unchanged otherwise.
In `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py`:
- Around line 63-66: The path comparison in the checkpoint validation is
comparing raw strings in the `convert_dcp_to_hf()` flow, which can fail for
equivalent relative/absolute paths. Normalize both `output_path` and
`saved_hf_checkpoint_path` to resolved absolute paths before the equality check
in the checkpoint code so the `ValueError` is only raised when they truly
differ.
In
`@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py`:
- Around line 31-40: The HuggingFace fallback in Tulu3Dataset.__init__ is
ignoring the split argument when data_path is None, so non-train configurations
can silently use the parent default. Update the fallback path in
Tulu3Dataset.__init__ to forward split along with the existing kwargs when
calling super().__init__, and make sure the local-file branch keeps its current
split handling consistent.
In `@services/rl/src/nmp/rl/tasks/training/datasets/preparation.py`:
- Around line 405-421: The auto-split path in preparation.py still loads every
non-empty line into memory before shuffling, which can OOM on large JSONL
inputs; update the training-file split logic around the existing split/check
code to avoid materializing all rows at once. Use a streaming or incremental
approach in the same dataset-preparation flow (near the current train/validation
split handling and random.shuffle usage) so validation splitting works without
reading the full file into a list.
In `@services/rl/src/nmp/rl/tasks/training/datasets/validation.py`:
- Around line 240-263: The JSONL validation in validate_dataset currently
accepts any row that matches the DPO union, which allows mixed schema files to
pass and later be misclassified by detect_dpo_schema_name. Update the validation
flow in validation.py so it enforces a single concrete schema across all rows,
using the first valid row’s schema as the expected type and rejecting any
subsequent row that matches a different DPO schema. Keep the logic aligned with
detect_dpo_schema_name and the dataset validation helpers in
_validate_json_object/schema_factory so mixed-format files fail fast with a
clear DatasetFormatError.
- Around line 232-235: The validation path in the dataset schema lookup
currently skips unsupported dataset types by returning early, which hides typos
and training-type mismatches. Update the unsupported-type branch in the
validation logic around schema_factory and SCHEMAS so it raises an error instead
of silently returning, and make sure the check clearly distinguishes valid
training types from values like detect_dpo_schema_name() output.
In `@services/rl/src/nmp/rl/tasks/training/distributed.py`:
- Around line 97-114: The stale-marker cleanup in distributed barrier setup is
racy because only the coordinator removes old markers while workers may still
read or create files in the same directory. Update the barrier handling in
distributed.py so each run attempt uses its own barrier directory, or add a
coordinator-initialized sentinel that workers must wait for before creating or
checking any .ready markers. Keep the cleanup and synchronization logic anchored
around the distributed barrier setup flow so workers never touch files until the
coordinator has finished preparing the directory.
- Around line 204-205: The timeout handling in the barrier wait paths is
treating an explicit 0 as falsy and replacing it with the default, which should
not happen. Update the logic in the relevant distributed training methods
(including the barrier wait code around the marker setup and the other wait path
referenced in the review) to use an explicit is None check on timeout before
falling back to self._barrier_timeout. Keep the behavior in the methods that use
timeout and _barrier_timeout so a caller-supplied zero is honored unchanged.
---
Nitpick comments:
In `@services/rl/src/nmp/rl/tasks/training/distributed.py`:
- Around line 73-74: The from_env classmethod on DistributedContext uses a
string-based return annotation, which violates the Python typing guideline.
Update the signature of DistributedContext.from_env to use a concrete return
type annotation instead of a quoted forward reference, keeping the method
behavior unchanged.
🪄 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: 29fd351b-d968-477c-ba2b-d4e68a05b313
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (75)
docker-bake.hcldocker/Dockerfile.nmp-rl-basedocker/Dockerfile.nmp-rl-tasksdocker/Dockerfile.nmp-rl-trainingdocker/rl/Dockerfile.platform-workspacedocker/rl/pyproject.workspace.tomlpackages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.pyplugins/nemo-rl/README.mdplugins/nemo-rl/pyproject.tomlplugins/nemo-rl/src/nemo_rl_plugin/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.pyplugins/nemo-rl/src/nemo_rl_plugin/cli/main.pyplugins/nemo-rl/src/nemo_rl_plugin/config.pyplugins/nemo-rl/src/nemo_rl_plugin/contributor.pyplugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.pyplugins/nemo-rl/src/nemo_rl_plugin/schema.pyplugins/nemo-rl/src/nemo_rl_plugin/sdk/__init__.pyplugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.pyplugins/nemo-rl/src/nemo_rl_plugin/transform.pypyproject.tomlservices/rl/README.mdservices/rl/pyproject.tomlservices/rl/src/nmp/rl/app/constants.pyservices/rl/src/nmp/rl/app/jobs/compiler.pyservices/rl/src/nmp/rl/app/jobs/training/schemas.pyservices/rl/src/nmp/rl/compile.pyservices/rl/src/nmp/rl/config.pyservices/rl/src/nmp/rl/entities/values.pyservices/rl/src/nmp/rl/images.pyservices/rl/src/nmp/rl/schemas.pyservices/rl/src/nmp/rl/tasks/file_io/__init__.pyservices/rl/src/nmp/rl/tasks/file_io/__main__.pyservices/rl/src/nmp/rl/tasks/file_io/callbacks.pyservices/rl/src/nmp/rl/tasks/file_io/run.pyservices/rl/src/nmp/rl/tasks/model_entity/__init__.pyservices/rl/src/nmp/rl/tasks/model_entity/__main__.pyservices/rl/src/nmp/rl/tasks/model_entity/run.pyservices/rl/src/nmp/rl/tasks/training/__main__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txtservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.pyservices/rl/src/nmp/rl/tasks/training/chat_templates.pyservices/rl/src/nmp/rl/tasks/training/datasets/preparation.pyservices/rl/src/nmp/rl/tasks/training/datasets/schemas.pyservices/rl/src/nmp/rl/tasks/training/datasets/validation.pyservices/rl/src/nmp/rl/tasks/training/distributed.pyservices/rl/src/nmp/rl/tasks/training/errors/converter.pyservices/rl/src/nmp/rl/tasks/training/errors/error_rules.yamlservices/rl/src/nmp/rl/tasks/training/errors/exceptions.pyservices/rl/src/nmp/rl/tasks/training/errors/parser.pyservices/rl/src/nmp/rl/tasks/training/integrations.pyservices/rl/src/nmp/rl/tasks/training/protocol.pyservices/rl/src/nmp/rl/tasks/training/runner.pyservices/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinjaservices/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinjaservices/rl/src/nmp/rl/tasks/training/templates/phi-4.jinjaservices/rl/src/nmp/rl/tasks/training/utils.pyservices/rl/tests/test_compiler.pyservices/rl/tests/test_dpo_config.pyservices/rl/tests/test_schemas.py
✅ Files skipped from review due to trivial changes (10)
- services/rl/src/nmp/rl/tasks/model_entity/init.py
- plugins/nemo-rl/src/nemo_rl_plugin/init.py
- plugins/nemo-rl/src/nemo_rl_plugin/sdk/init.py
- plugins/nemo-rl/src/nemo_rl_plugin/cli/init.py
- plugins/nemo-rl/README.md
- plugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.py
- services/rl/src/nmp/rl/tasks/file_io/init.py
- docker/rl/pyproject.workspace.toml
- services/rl/README.md
- services/rl/src/nmp/rl/app/constants.py
🚧 Files skipped from review as they are similar to previous changes (47)
- plugins/nemo-rl/src/nemo_rl_plugin/config.py
- services/rl/src/nmp/rl/tasks/training/templates/phi-4.jinja
- services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja
- services/rl/src/nmp/rl/tasks/model_entity/main.py
- plugins/nemo-rl/src/nemo_rl_plugin/contributor.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt
- services/rl/src/nmp/rl/config.py
- services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja
- plugins/nemo-rl/src/nemo_rl_plugin/cli/main.py
- services/rl/src/nmp/rl/tasks/file_io/main.py
- services/rl/src/nmp/rl/tasks/training/main.py
- services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja
- plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py
- services/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinja
- services/rl/pyproject.toml
- services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinja
- packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py
- services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja
- services/rl/src/nmp/rl/tasks/training/errors/converter.py
- services/rl/src/nmp/rl/tasks/training/protocol.py
- plugins/nemo-rl/pyproject.toml
- services/rl/src/nmp/rl/tasks/training/utils.py
- plugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.py
- plugins/nemo-rl/src/nemo_rl_plugin/transform.py
- services/rl/src/nmp/rl/images.py
- services/rl/src/nmp/rl/compile.py
- docker-bake.hcl
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py
- pyproject.toml
- plugins/nemo-rl/src/nemo_rl_plugin/schema.py
- services/rl/src/nmp/rl/app/jobs/training/schemas.py
- services/rl/src/nmp/rl/tasks/training/errors/parser.py
- services/rl/src/nmp/rl/app/jobs/compiler.py
- services/rl/src/nmp/rl/tasks/training/integrations.py
- services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml
- services/rl/src/nmp/rl/tasks/training/errors/exceptions.py
- services/rl/tests/test_schemas.py
- services/rl/tests/test_dpo_config.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py
- services/rl/src/nmp/rl/schemas.py
- services/rl/src/nmp/rl/tasks/training/datasets/schemas.py
- services/rl/src/nmp/rl/tasks/file_io/run.py
- services/rl/src/nmp/rl/tasks/file_io/callbacks.py
- services/rl/tests/test_compiler.py
anubhutivyas
force-pushed
the
aalgo-292-dpo-support/anubhutiv
branch
from
June 30, 2026 16:06
9624583 to
b14c3d8
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py`:
- Around line 727-737: The Ray shutdown path in `_stop_ray` can exceed the outer
`cleanup(timeout=30)` budget because the inner `subprocess.run(...,
timeout=grace_period + 30)` plus the post-call sleep can keep the non-daemon
cleanup thread alive too long. Reduce the inner stop timeout in `_stop_ray` (and
any related follow-up delay) so it always fits within the caller’s cleanup
timeout, and keep the total stop-and-sleep duration comfortably below
`cleanup()`’s limit.
In `@services/rl/src/nmp/rl/tasks/training/datasets/preparation.py`:
- Around line 418-421: The split shuffling in the dataset preparation logic is
mutating the process-wide RNG state via random.seed(seed), which can leak into
later training code. Update the shuffle in the relevant preparation function to
use a local RNG instance instead of the global random module, while preserving
determinism for the split. Keep the change scoped to the shuffle path in
preparation.py so the existing split behavior remains reproducible without
affecting other code.
🪄 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: 7d6f8f2d-d5b4-4778-9375-e72bf2c5151e
📒 Files selected for processing (25)
services/rl/src/nmp/rl/app/jobs/compiler.pyservices/rl/src/nmp/rl/entities/values.pyservices/rl/src/nmp/rl/schemas.pyservices/rl/src/nmp/rl/tasks/file_io/run.pyservices/rl/src/nmp/rl/tasks/model_entity/run.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txtservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.pyservices/rl/src/nmp/rl/tasks/training/chat_templates.pyservices/rl/src/nmp/rl/tasks/training/datasets/preparation.pyservices/rl/src/nmp/rl/tasks/training/datasets/schemas.pyservices/rl/src/nmp/rl/tasks/training/datasets/validation.pyservices/rl/src/nmp/rl/tasks/training/distributed.pyservices/rl/src/nmp/rl/tasks/training/errors/converter.pyservices/rl/src/nmp/rl/tasks/training/errors/error_rules.yamlservices/rl/src/nmp/rl/tasks/training/errors/parser.pyservices/rl/src/nmp/rl/tasks/training/runner.pyservices/rl/src/nmp/rl/tasks/training/utils.py
🚧 Files skipped from review as they are similar to previous changes (23)
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py
- services/rl/src/nmp/rl/entities/values.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py
- services/rl/src/nmp/rl/tasks/training/utils.py
- services/rl/src/nmp/rl/tasks/training/errors/converter.py
- services/rl/src/nmp/rl/tasks/training/chat_templates.py
- services/rl/src/nmp/rl/tasks/training/errors/parser.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/init.py
- services/rl/src/nmp/rl/schemas.py
- services/rl/src/nmp/rl/tasks/training/distributed.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py
- services/rl/src/nmp/rl/tasks/training/runner.py
- services/rl/src/nmp/rl/tasks/training/datasets/validation.py
- services/rl/src/nmp/rl/app/jobs/compiler.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py
- services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml
- services/rl/src/nmp/rl/tasks/file_io/run.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py
- services/rl/src/nmp/rl/tasks/model_entity/run.py
- services/rl/src/nmp/rl/tasks/training/datasets/schemas.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@services/rl/src/nmp/rl/tasks/training/datasets/preparation.py`:
- Around line 439-443: The auto-split file handling in preparation.py truncates
the source when `train_file` and `output_train` are the same path because
`split_train_val` opens the write target before streaming from `src`. Update the
split logic around the file-open block so `train_file` is never opened for
writing until after the source has been fully read, or write through a separate
temporary path and then replace `output_train` afterward. Keep the fix localized
to the auto-split flow that uses `train_file`, `output_train`, and `output_val`
so merged training data is preserved.
🪄 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: a05ae04b-2f8b-44fb-b711-080f5226efc8
📒 Files selected for processing (12)
docker/Dockerfile.nmp-rl-basedocker/Dockerfile.nmp-rl-trainingservices/rl/src/nmp/rl/tasks/model_entity/run.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.pyservices/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.pyservices/rl/src/nmp/rl/tasks/training/chat_templates.pyservices/rl/src/nmp/rl/tasks/training/datasets/preparation.pyservices/rl/src/nmp/rl/tasks/training/datasets/validation.pyservices/rl/src/nmp/rl/tasks/training/distributed.pyservices/rl/tests/test_dpo_config.py
🚧 Files skipped from review as they are similar to previous changes (8)
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py
- docker/Dockerfile.nmp-rl-base
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py
- docker/Dockerfile.nmp-rl-training
- services/rl/tests/test_dpo_config.py
- services/rl/src/nmp/rl/tasks/training/distributed.py
- services/rl/src/nmp/rl/tasks/model_entity/run.py
- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py
albcui
reviewed
Jun 30, 2026
anubhutivyas
force-pushed
the
aalgo-292-dpo-support/anubhutiv
branch
from
June 30, 2026 23:58
0db93a2 to
80cd9ab
Compare
Signed-off-by: anubhutiv <[email protected]>
anubhutivyas
force-pushed
the
aalgo-292-dpo-support/anubhutiv
branch
from
July 1, 2026 00:02
80cd9ab to
d9dc55f
Compare
albcui
self-requested a review
July 2, 2026 16:33
albcui
approved these changes
Jul 2, 2026
arpitsardhana
pushed a commit
that referenced
this pull request
Jul 9, 2026
Signed-off-by: anubhutiv <[email protected]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Linear ticket - https://linear.app/nvidia/issue/AALGO-292/add-dpo-support-for-customizer-plugin
Summary by CodeRabbit
New Features
Bug Fixes
Tests