diff --git a/docker-bake.hcl b/docker-bake.hcl index 6343ea078d..2c662977bd 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -327,6 +327,67 @@ group "nmp-unsloth" { ] } +group "nmp-rl" { + targets = [ + "nmp-rl-base-builder", + "nmp-rl-tasks", + "nmp-rl-training", + ] +} + +# Pruned workspace slice for nmp-rl images (keep in sync with +# docker/rl/pyproject.workspace.toml + Dockerfile.platform-workspace members). +target "rl-platform-workspace" { + target = "platform-workspace" + context = "." + dockerfile = "docker/rl/Dockerfile.platform-workspace" + output = ["type=cacheonly"] + platforms = get_platforms() +} + +# Heavy base: NGC torch/CUDA + NeMo-RL v0.4.0 + Ray. +target "nmp-rl-base-builder" { + target = "nmp-rl-base" + context = "." + dockerfile = "docker/Dockerfile.nmp-rl-base" + cache-to = maybe_registry_cache_to("nmp-rl-base") + cache-from = maybe_registry_cache_from("nmp-rl-base") + tags = base_tags("nmp-rl-base") + output = image_output() + platforms = get_platforms() +} + +# GPU DPO training image: base + platform glue. Bootstraps Ray at runtime. +target "nmp-rl-training" { + target = "runtime" + context = "." + dockerfile = "docker/Dockerfile.nmp-rl-training" + contexts = { + platform-workspace = "target:rl-platform-workspace" + nmp-rl-base = "target:nmp-rl-base-builder" + } + cache-to = maybe_registry_cache_to("nmp-rl-training") + cache-from = maybe_registry_cache_from("nmp-rl-training") + tags = sha_and_maybe_latest_tags("nmp-rl-training") + output = image_output() + platforms = get_platforms() +} + +# Lighter CPU image for the file_io / model_entity steps (no NeMo-RL/Ray). +target "nmp-rl-tasks" { + target = "runtime" + context = "." + dockerfile = "docker/Dockerfile.nmp-rl-tasks" + contexts = { + platform-workspace = "target:rl-platform-workspace" + } + cache-to = maybe_registry_cache_to("nmp-rl-tasks") + cache-from = maybe_registry_cache_from("nmp-rl-tasks") + tags = sha_and_maybe_latest_tags("nmp-rl-tasks") + output = image_output() + platforms = get_platforms() +} + # Base images for consolidated containers target "nmp-python-base" { target = python_base_target() diff --git a/docker/Dockerfile.nmp-rl-base b/docker/Dockerfile.nmp-rl-base new file mode 100644 index 0000000000..191f01da84 --- /dev/null +++ b/docker/Dockerfile.nmp-rl-base @@ -0,0 +1,56 @@ +# syntax=docker/dockerfile:1 +# nmp-rl base - NVIDIA NeMo-RL v0.6.0 NGC container + CVE hardening. +# +# Bases on the PUBLISHED NGC NeMo-RL container (amd64 + arm64 at v0.6.0), which +# bundles the full RL stack on Python 3.13: PyTorch 2.10, Ray 2.54.0, vLLM +# 0.17.1, Megatron-Core 0.18, Transformers 5.3. +# +# Publish target: nmp-rl-base. + +# Referenced by tag rather than a @sha256 digest pin: NGC publishes separate +# per-architecture tags/digests for nemo-rl, so a single digest doesn't cleanly +# cover the multi-arch (amd64 + arm64) base we build on. The :v0.6.0 tag is the +# stable, arch-agnostic reference. +ARG NEMO_RL_IMAGE=nvcr.io/nvidia/nemo-rl:v0.6.0 + +FROM ${NEMO_RL_IMAGE} AS nmp-rl-base + +# NeMo-RL v0.6.0 ships its venv at /opt/nemo_rl_venv (Python 3.13). +# VERIFY this path against the actual image if NeMo-RL relocates it in a future +# tag — the training image's /opt/venv symlink and entrypoint depend on it. +ENV VIRTUAL_ENV=/opt/nemo_rl_venv \ + HF_HUB_ENABLE_HF_TRANSFER=1 \ + OTEL_PYTHON_EXCLUDED_URLS="health" + +# The NGC base venv is built as root with locked-down cache dirs; relax them so +# editable installs in the training image can write. +RUN chmod 755 /root /root/.cache /root/.cache/uv /root/.local /root/.local/share /root/.local/share/uv 2>/dev/null || true + +# NeMo-RL is installed EDITABLE in this base: its source tree lives at +# /opt/nemo-rl (root-owned, non-traversable) and the editable import finder +# os.stat()s files under it at import time. The training image drops to +# USER 1000, so `from nemo_rl... import ...` hits the finder and fails with +# PermissionError on /opt/nemo-rl/nemo_rl/__init__.py. Make the source tree +# world-readable + dir-traversable (a+rX) so the non-root runtime can import it. +# Read-only is enough — nothing writes back into the editable tree at runtime. +RUN chmod -R a+rX /opt/nemo-rl + +# CVE remediation: remove vLLM from the main venv and any pre-built Ray worker +# venvs (NeMo-RL rebuilds those at runtime). vLLM is unused by DPO/SFT/LoRA. +# Add it back when adding GRPO. +RUN set -e; \ + rm -rf ${VIRTUAL_ENV}/lib/python*/site-packages/vllm \ + ${VIRTUAL_ENV}/lib/python*/site-packages/vllm-*.dist-info; \ + for venv_dir in /opt/ray_venvs/*/; do \ + [ -d "$venv_dir/lib" ] || continue; \ + rm -rf "$venv_dir"/lib/python*/site-packages/vllm \ + "$venv_dir"/lib/python*/site-packages/vllm-*.dist-info; \ + done + +# Drop to a non-root user by default so DIRECT consumers of this published base do +# not run with container root. The build steps above intentionally run as root (the +# NGC base default — a USER directive only changes the final/default user, not the +# RUNs preceding it). Derived images that need root for their own build steps +# (e.g. Dockerfile.nmp-rl-training) re-assert `USER root` and then drop back to a +# non-root UID at the end. UID 1000 is the NGC base's `ubuntu` user. +USER 1000:1000 diff --git a/docker/Dockerfile.nmp-rl-tasks b/docker/Dockerfile.nmp-rl-tasks new file mode 100644 index 0000000000..a9e2270c86 --- /dev/null +++ b/docker/Dockerfile.nmp-rl-tasks @@ -0,0 +1,73 @@ +# syntax=docker/dockerfile:1 +# nmp-rl tasks - CPU file_io / model_entity steps for an nmp-rl DPO job. +# +# Deliberately does NOT build on nmp-rl-base: the download/upload/model-entity +# steps only need the platform glue (SDK + customization-common), not NeMo-RL / +# Ray / vLLM. Basing on the NGC image keeps it consistent with the platform's +# CUDA userspace while staying far lighter than the training image. + +ARG SMOKE_MARKER=smoke_nmp_rl_tasks +ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.02-py3 + +FROM ${PYTORCH_BASE} AS base + +WORKDIR /app + +COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv +ENV PATH="/bin:${PATH}" + +ENV VIRTUAL_ENV=/opt/venv \ + UV_PROJECT_ENVIRONMENT=/opt/venv \ + UV_LINK_MODE=copy \ + UV_COMPILE_BYTECODE=1 \ + HF_HUB_ENABLE_HF_TRANSFER=1 \ + OTEL_PYTHON_EXCLUDED_URLS="health" +ENV PATH="/opt/venv/bin:/root/.local/bin:${PATH}" + +RUN uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages + +FROM base AS runtime + +ARG USERNAME=ubuntu +ARG USER_UID=1000 +ARG USER_GID=1000 + +COPY --from=platform-workspace / /app +WORKDIR /app + +RUN mkdir -p /home/${USERNAME}/.cache && \ + chown -R ${USER_UID}:${USER_GID} /home/${USERNAME} /app/services/rl + +# Only the glue + nmp-rl (compile/tasks side). No NeMo-RL extra → no Ray/vLLM. +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \ + -e /app/sdk/python/nemo-platform \ + -e /app/packages/nemo_platform_plugin \ + -e /app/packages/nmp_common \ + -e /app/packages/nmp_customization_common \ + -e /app/services/rl + +ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" +# Default CMD is a harmless help invocation; the platform overrides `command` +# per step (e.g. -m nmp.rl.tasks.file_io / nmp.rl.tasks.model_entity). +ENTRYPOINT ["/opt/venv/bin/python"] +CMD ["-m", "nmp.rl.tasks.file_io", "--help"] + +USER ${USER_UID}:${USER_GID} + +# NOTE: this smoke-test stage is intentionally NOT wired into docker-bake.hcl +# (unlike nmp-automodel-{tasks,training}-smoke-test). There are no RL smoke tests +# yet — tests/smoke_gpu/ carries no tests marked `smoke_nmp_rl_tasks`, and that +# marker isn't registered in its conftest.py. Wiring a bake target now would run +# `pytest -m smoke_nmp_rl_tasks` against zero collected tests, which exits 5 +# ("no tests ran") and fails the bake. Add RL import smoke tests + register the +# marker first, then add the bake target (see nmp-automodel-tasks-smoke-test for +# the pattern). The stage is kept so it's ready to wire up once tests exist. +FROM runtime AS smoke-test +ARG SMOKE_MARKER +USER 0 +COPY tests/smoke_gpu/ /smoke_test/ +RUN uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --reinstall pytest && \ + ${VIRTUAL_ENV}/bin/pytest /smoke_test/ -m ${SMOKE_MARKER} -v + +FROM runtime diff --git a/docker/Dockerfile.nmp-rl-training b/docker/Dockerfile.nmp-rl-training new file mode 100644 index 0000000000..a7dd4a0db0 --- /dev/null +++ b/docker/Dockerfile.nmp-rl-training @@ -0,0 +1,105 @@ +# syntax=docker/dockerfile:1 +# nmp-rl training - GPU DPO step (NeMo-RL v0.6.0 + Ray + nmp-rl package). +# +# Built on nmp-rl-base (NGC NeMo-RL v0.6.0 venv at /opt/nemo_rl_venv, Python +# 3.13). Adds the lightweight platform glue editably into that venv, using the +# ported no_override_requirements.txt so the install does not clobber NeMo-RL's +# pinned ML stack (ray / torch / mlflow / cryptography / starlette / +# prometheus-client). The training step runs `python -m nmp.rl.tasks.training`, +# which bootstraps a Ray cluster and runs the DPO driver against NeMo-RL. + +ARG SMOKE_MARKER=smoke_nmp_rl_training + +# Supplied by bake (target:nmp-rl-base-builder). +FROM nmp-rl-base AS rl-base + +FROM rl-base AS runtime + +# nmp-rl-base now defaults to a non-root UID; re-assert root for the editable +# installs, chowns and symlink below. The stage drops back to USER ${USER_UID} at +# the end, so the published training image still runs non-root. +USER root + +COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv +ENV PATH="/bin:${PATH}" + +ARG USERNAME=ubuntu +ARG USER_UID=1000 +ARG USER_GID=1000 + +COPY --from=platform-workspace / /app +WORKDIR /app + +RUN mkdir -p /home/${USERNAME}/.cache && \ + chown -R ${USER_UID}:${USER_GID} /home/${USERNAME} /app/services/rl + +# NMP is pinned to Python 3.11, but the NeMo-RL base venv is 3.13. The glue is +# pure-Python and installs into the 3.13 venv fine. Symlink /opt/venv -> +# /opt/nemo_rl_venv so the single RL_PYTHON_ENTRYPOINT (/opt/venv/bin/python) +# the compiler stamps onto every step resolves in this image too (the CPU tasks +# image has a real /opt/venv). +RUN ln -sfn /opt/nemo_rl_venv /opt/venv + +# The NeMo-RL base venv ships some packages with missing dist-info METADATA (a +# known property of the image — the package code is present, only uv's metadata +# file is gone). Normal `uv pip install` reads ALL installed metadata during +# dependency resolution and dies ("Failed to read anyio==... METADATA"). So we +# install with --no-deps, which skips resolution entirely and never touches the +# corrupted metadata. The heavy ML stack already lives in the base venv, +#so --no-deps leaves it intact. +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --no-deps \ + -e /app/sdk/python/nemo-platform \ + -e /app/packages/nemo_platform_plugin \ + -e /app/packages/nmp_common \ + -e /app/packages/nmp_customization_common \ + -e /app/services/rl + +# Light PURE-PYTHON glue deps the NeMo-RL base lacks (or ships with corrupted +# dist-info). Three rules: +# 1. --reinstall: force a clean (re)install of each target. uv otherwise reads +# the installed copy's METADATA for its "already satisfied?" check, which +# fails on the base's missing-METADATA dist-info (e.g. idna, anyio). A +# version pin alone is NOT enough — when the pin matches the installed +# (corrupted) version, uv still tries to read it. --reinstall skips that +# check entirely and installs fresh from the index (also repairing metadata). +# 2. --no-deps: never resolve/read OTHER packages' metadata. +# 3. ONLY non-ML, version-stable packages. Base ML deps (pydantic, anyio, +# httpx, typing-extensions, …) are present + version-sensitive to the +# torch/transformers stack — we must NOT reinstall them, so they are omitted. +# If the training step hits a runtime ImportError for another pure-python +# package, add it here (pinned), NOT one of the ML deps above. +# Discovered empirically via scripts/gpu-dpo-smoke/discover_deps.py against the +# v0.6.0 base — the ONLY pure-python deps the training entrypoint imports that are +# missing from /opt/nemo_rl_venv. base58 + lark are NeMo-RL's OWN deps (absent +# from the main venv); distro is the SDK's; pydantic-settings is our glue's. +# Re-run discover_deps.py after a base-image bump to refresh this list. +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --no-deps --reinstall \ + "base58==2.1.1" \ + "lark==1.3.1" \ + "distro==1.9.0" \ + "pydantic-settings==2.14.2" + +ENV PATH="${VIRTUAL_ENV}/bin:${PATH}" +ENTRYPOINT ["/opt/venv/bin/python"] +CMD ["-m", "nmp.rl.tasks.training", "--help"] + +USER ${USER_UID}:${USER_GID} + +# NOTE: this smoke-test stage is intentionally NOT wired into docker-bake.hcl +# (unlike nmp-automodel-{tasks,training}-smoke-test). There are no RL smoke tests +# yet — tests/smoke_gpu/ carries no tests marked `smoke_nmp_rl_training`, and that +# marker isn't registered in its conftest.py. Wiring a bake target now would run +# `pytest -m smoke_nmp_rl_training` against zero collected tests, which exits 5 +# ("no tests ran") and fails the bake. Add RL import smoke tests + register the +# marker first, then add the bake target (see nmp-automodel-training-smoke-test +# for the pattern). The stage is kept so it's ready to wire up once tests exist. +FROM runtime AS smoke-test +ARG SMOKE_MARKER +USER 0 +COPY tests/smoke_gpu/ /smoke_test/ +RUN uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache --reinstall pytest && \ + ${VIRTUAL_ENV}/bin/pytest /smoke_test/ -m ${SMOKE_MARKER} -v + +FROM runtime diff --git a/docker/rl/Dockerfile.platform-workspace b/docker/rl/Dockerfile.platform-workspace new file mode 100644 index 0000000000..3a6f85c71f --- /dev/null +++ b/docker/rl/Dockerfile.platform-workspace @@ -0,0 +1,18 @@ +# syntax=docker/dockerfile:1 +# Minimal Platform workspace slice for nmp-rl container installs. +# Used as a named build context (platform-workspace). +# Keep in sync with docker/rl/pyproject.workspace.toml members. + +FROM scratch AS platform-workspace +# Reduced workspace file for this partial source tree. +COPY docker/rl/pyproject.workspace.toml pyproject.toml +# nemo-platform-sdk's hatch build force-includes docs/ from the repo root; the +# openapi symlink must resolve at build time, so copy both trees. +COPY docs docs +COPY openapi openapi +COPY packages/nmp_build_tools packages/nmp_build_tools +COPY packages/nmp_common packages/nmp_common +COPY packages/nmp_customization_common packages/nmp_customization_common +COPY packages/nemo_platform_plugin packages/nemo_platform_plugin +COPY sdk/python/nemo-platform sdk/python/nemo-platform +COPY services/rl services/rl diff --git a/docker/rl/pyproject.workspace.toml b/docker/rl/pyproject.workspace.toml new file mode 100644 index 0000000000..3e1321d55b --- /dev/null +++ b/docker/rl/pyproject.workspace.toml @@ -0,0 +1,29 @@ +# Minimal uv workspace for nmp-rl container image builds only. +# Keeps uv validation scoped to the partial source tree copied into the image. +# Keep in sync with docker/rl/Dockerfile.platform-workspace members. + +[project] +name = "nemo-platform-rl-image" +version = "0.0.0" +requires-python = ">=3.11,<3.14" + +[tool.uv] +required-version = ">=0.9.14,<0.10.0" + +[tool.uv.workspace] +members = [ + "packages/nmp_build_tools", + "sdk/python/nemo-platform", + "packages/nemo_platform_plugin", + "packages/nmp_common", + "packages/nmp_customization_common", + "services/rl", +] + +[tool.uv.sources] +nmp-build-tools = { workspace = true } +nemo-platform-sdk = { workspace = true } +nemo-platform-plugin = { workspace = true } +nmp-common = { workspace = true } +nmp-customization-common = { workspace = true } +nmp-rl = { workspace = true } diff --git a/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py b/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py index b56c555802..28343c4e95 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/contributor/jobs.py @@ -39,6 +39,24 @@ def require_docker_runtime(backend_label: str) -> None: ) +def require_distributed_runtime(backend_label: str) -> None: + """Refuse to compile when the platform isn't a remote Kubernetes cluster. + + Sibling to :func:`require_docker_runtime` for backends that provision a Ray + cluster (e.g. NeMo-RL DPO). Unlike the SFT backends, these have no local + single-node Docker fallback: they need the platform's Kubernetes/Volcano + scheduler to place GPU pods and inject the distributed env + (``RANK``/``WORLD_SIZE``/``MASTER_ADDR``). Surface the misconfiguration before + the Jobs API rejects the spec. + """ + platform_config = NemoPlatformConfig.get() + if platform_config.runtime != Runtime.KUBERNETES: + raise PlatformJobCompilationError( + f"{backend_label} training requires platform.runtime: kubernetes — it provisions a Ray " + "cluster on the remote GPU cluster and has no local Docker fallback.", + ) + + class BaseSubmitJob(NemoJob): """Shared submit-only job scaffold. diff --git a/plugins/nemo-rl/README.md b/plugins/nemo-rl/README.md new file mode 100644 index 0000000000..701ed26d85 --- /dev/null +++ b/plugins/nemo-rl/README.md @@ -0,0 +1,75 @@ +# nemo-rl-plugin + +NeMo-RL customization contributor for the NeMo Platform. Adds **DPO** training +on a Ray cluster (via [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL) +v0.6.0) as the `rl` backend under `/apis/customization`. + +Thin contributor layer only — the heavy compile glue and container tasks live in +[`services/rl`](../../services/rl) (`nmp-rl`). + +## Surfaces + +- **CLI:** `nemo customization rl submit -w ` (submit-only; + `run` is disabled — there is no local execution). +- **REST:** `POST /apis/customization/v2/workspaces/{workspace}/rl/jobs` +- **SDK:** `client.customization.rl.jobs.create(...)` + +## Constraints + +- **Remote Kubernetes only** — gated via `require_distributed_runtime`. There is + no local Docker fallback (unlike automodel/unsloth). +- **Single-node multi-GPU and multi-node** both supported (`parallelism.num_nodes`). + Multi-node requires `NMP_RL_MULTINODE_SHARED_STORAGE_PATH`. +- **DPO is full-weight** (no PEFT). GRPO/PPO are headroom (`TrainingMethod` is a + single-member union today). + +## Job spec + +`model` and `dataset` are string refs; the method lives under `training` with +`type: "dpo"`. The `dataset` fileset holds **both** `training.jsonl` and +`validation.jsonl` as `{prompt, chosen, rejected}` preference rows. + +```json +{ + "model": "default/qwen3-0.6b", + "dataset": "default/dpo-data", + "training": { + "type": "dpo", + "epochs": 1, + "learning_rate": 5e-6, + "max_seq_length": 1024, + "batch_size": 32, + "micro_batch_size": 1, + "ref_policy_kl_penalty": 0.05, + "parallelism": { "num_nodes": 1, "num_gpus_per_node": 1 } + }, + "output": { "name": "qwen3-0.6b-dpo" } +} +``` + +Configurable `training` knobs (full reference: the skill's +`references/hyperparameters.md` § NeMo-RL (DPO)): the optimizer/schedule/batch +fields, `parallelism`, `optimizer_type`, `adam_eps`, `activation_checkpointing`, +`keep_top_k`, `val_at_end`, and the DPO-specific `ref_policy_kl_penalty`, +`preference_loss_weight`, `sft_loss_weight`, `preference_average_log_probs`, +`sft_average_log_probs`, `max_grad_norm`. `RlJobInput` (`schema.py`) is the +authoritative input shape; `nemo customization rl explain` prints it live. + +## Compiled job (4 steps) + +`submit` → `RlJobInput` → transform → `RlJobOutput` → compiled `PlatformJobSpec`: + +1. **download** — model fileset + preference dataset → PVC (CPU, `nmp-rl-tasks`) +2. **dpo-training** — Ray DPO step (GPU, `nmp-rl-training`); single-node `gpu` or + multi-node `gpu_distributed` executor, selected by `parallelism.num_nodes` +3. **upload** — trained checkpoint → output fileset (CPU) +4. **model-entity** — register the full-weight output `ModelEntity` + +## Related + +- **Skill:** the `nemo-customizer` skill documents the end-to-end DPO workflow + (`plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/`). +- **Design:** [`docs/customizer/nemo-rl-dpo-plugin-design.md`](../../docs/customizer/nemo-rl-dpo-plugin-design.md). +- **GPU e2e smoke test:** [`scripts/gpu-dpo-smoke/`](../../scripts/gpu-dpo-smoke). +- **Images:** [`docker/Dockerfile.nmp-rl-base`](../../docker/Dockerfile.nmp-rl-base), + `Dockerfile.nmp-rl-training`, `Dockerfile.nmp-rl-tasks`. diff --git a/plugins/nemo-rl/pyproject.toml b/plugins/nemo-rl/pyproject.toml new file mode 100644 index 0000000000..5ed5122379 --- /dev/null +++ b/plugins/nemo-rl/pyproject.toml @@ -0,0 +1,57 @@ +[project] +name = "nemo-rl-plugin" +version = "0.1.0" +description = "NeMo-RL DPO customization contributor for NeMo Platform (Ray on Kubernetes)." +readme = "README.md" +requires-python = ">=3.11,<3.14" +dependencies = [ + "nemo-platform-plugin", + "nemo-platform", + "nmp-rl", + "nmp-customization-common", + "pydantic>=2.10.6", + "pydantic-settings>=2.6.1", + "typer>=0.12.5", +] + +# Heavy ML deps (nemo-rl, ray, vllm) are baked into the nmp-rl-training container +# image, not installed by the plugin. The plugin process only needs the +# lightweight compile-side imports. + +[project.entry-points."nemo.customization.contributors"] +rl = "nemo_rl_plugin.contributor:RlContributor" + +[project.entry-points."nemo.jobs"] +"customization.rl.jobs" = "nemo_rl_plugin.jobs.jobs:RlJob" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/nemo_rl_plugin"] + +[tool.uv.sources] +nemo-platform-plugin = { workspace = true } +nemo-platform = { workspace = true } +nemo-customizer-plugin = { workspace = true } +nmp-rl = { workspace = true } +nmp-customization-common = { workspace = true } + +[dependency-groups] +dev = [ + "pytest>=8.3.4", + "pytest-asyncio>=0.25.3", + "ruff>=0.11.8", + "fastapi>=0.115.0", + "httpx>=0.27.0", + "nemo-customizer-plugin", +] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +pythonpath = ["src"] +testpaths = ["tests"] + +[tool.pyright] +extraPaths = ["src"] diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/__init__.py b/plugins/nemo-rl/src/nemo_rl_plugin/__init__.py new file mode 100644 index 0000000000..cc713c109e --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NeMo-RL customization contributor — DPO training on a Ray cluster (Kubernetes).""" diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/cli/__init__.py b/plugins/nemo-rl/src/nemo_rl_plugin/cli/__init__.py new file mode 100644 index 0000000000..e5725ea5a4 --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/cli/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.py b/plugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.py new file mode 100644 index 0000000000..8a4fd34959 --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/cli/inputs.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI overrides for the NeMo-RL contributor. + +The override machinery is shared in :mod:`nmp.customization_common.cli.overrides`; +this module supplies the RL specifics: the ``RlJobInput`` schema (via +``load_job_json``), the ``JOB_JSON`` help text, and the run-disabled message. +""" + +import json +from pathlib import Path + +import typer +from nmp.customization_common.cli.overrides import apply_job_cli_overrides + +from nemo_rl_plugin.schema import RlJobInput + +_JOB_JSON_HELP = "Path to NeMo-RL job JSON (RlJobInput schema)." +_RUN_DISABLED_MESSAGE = ( + "NeMo-RL does not support local run (it provisions a Ray cluster on the remote Kubernetes cluster). " + "Submit to the platform API instead:\n" + " nemo customization rl submit -w " +) + + +def load_job_json(path: Path) -> str: + """Load and validate job JSON; return canonical JSON string for ``--spec``.""" + data = json.loads(path.read_text()) + validated = RlJobInput.model_validate(data) + return validated.model_dump_json() + + +def apply_rl_job_cli_overrides(group: typer.Typer) -> None: + """Flat ``rl`` CLI: ``submit JOB.json``; ``run`` is disabled.""" + apply_job_cli_overrides( + group, + load_job_json=load_job_json, + job_json_help=_JOB_JSON_HELP, + run_disabled_message=_RUN_DISABLED_MESSAGE, + ) diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/cli/main.py b/plugins/nemo-rl/src/nemo_rl_plugin/cli/main.py new file mode 100644 index 0000000000..105d4a35b7 --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/cli/main.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI hooks for the NeMo-RL customization contributor. + +The plugin's CLI surface is auto-mounted by the customization hub via +:meth:`RlContributor.get_cli`. This class provides the ``add_job_commands`` +integration hook for any caller that builds the CLI through that helper instead. +""" + +from __future__ import annotations + +import typer +from nemo_platform_plugin.job import NemoJob + +from nemo_rl_plugin.cli.inputs import apply_rl_job_cli_overrides +from nemo_rl_plugin.jobs.jobs import RlJob + + +class RlContributorCLI: + """Passed to ``add_job_commands`` to override run/submit with job-file args.""" + + def update_job_cli(self, job_cls: type[NemoJob], group: typer.Typer) -> None: + if job_cls is RlJob: + apply_rl_job_cli_overrides(group) diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/config.py b/plugins/nemo-rl/src/nemo_rl_plugin/config.py new file mode 100644 index 0000000000..cd592e60d0 --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/config.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Plugin configuration for NeMo-RL DPO training.""" + +from __future__ import annotations + +from nmp.customization_common.contributor.config import BaseTrainingPluginConfig, generate_job_id +from pydantic_settings import SettingsConfigDict + + +class RlPluginConfig(BaseTrainingPluginConfig): + """Environment-driven NeMo-RL plugin settings (``NMP_RL_`` prefix).""" + + model_config = SettingsConfigDict(env_prefix="NMP_RL_", extra="ignore") + + +def get_config() -> RlPluginConfig: + return RlPluginConfig() + + +def generate_rl_id() -> str: + """Generate a job name when the submitter omits ``name``.""" + return generate_job_id("rl") diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/contributor.py b/plugins/nemo-rl/src/nemo_rl_plugin/contributor.py new file mode 100644 index 0000000000..7a7b92c24b --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/contributor.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NeMo-RL customization contributor. + +Registered under ``nemo.customization.contributors`` (key ``rl``). The +customization router hub (``nemo-customizer-plugin``) discovers this class and +merges its routes/CLI/authz/SDK. Shared shape lives in +:class:`nmp.customization_common.contributor.base.BaseContributor`. +""" + +from __future__ import annotations + +from typing import ClassVar + +import typer +from nemo_platform_plugin.customization_contributor import CustomizationContributorSDKResources +from nmp.customization_common.contributor.base import BaseContributor + +from nemo_rl_plugin.config import RlPluginConfig, generate_rl_id, get_config +from nemo_rl_plugin.jobs.jobs import RlJob + + +class RlContributor(BaseContributor): + """Registers NeMo-RL routes/CLI under the customization router (DPO, Kubernetes only).""" + + name: ClassVar[str] = "rl" + job_cls: ClassVar[type] = RlJob + cli_help: ClassVar[str] = "NeMo-RL preference training (DPO) on a Ray cluster. Remote Kubernetes only." + jobs_router_description: ClassVar[str] = "NeMo-RL DPO training jobs (Ray on Kubernetes)." + + generate_job_name = staticmethod(generate_rl_id) + + def _get_config(self) -> RlPluginConfig: + return get_config() + + def apply_cli_overrides(self, app: typer.Typer) -> None: + from nemo_rl_plugin.cli.inputs import apply_rl_job_cli_overrides + + apply_rl_job_cli_overrides(app) + + def get_sdk_resources(self) -> CustomizationContributorSDKResources: + from nemo_rl_plugin.sdk.resources import AsyncRlCustomization, RlCustomization + + return CustomizationContributorSDKResources( + sync_resource=RlCustomization, + async_resource=AsyncRlCustomization, + ) diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py b/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py new file mode 100644 index 0000000000..bfb5cd11e7 --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/jobs/jobs.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NeMo-RL remote-submit DPO training job (NemoJob). + +Submit-only — executes as a 4-step ``PlatformJobSpec`` (download → DPO train → +upload → model-entity) on the platform's Kubernetes GPU cluster, where the +training step provisions a Ray cluster. + +Shared scaffold (``to_spec``) lives in +:class:`nmp.customization_common.contributor.jobs.BaseSubmitJob`. ``compile`` +stays here: it gates on the Kubernetes runtime (no local Docker fallback) and +resolves the execution profile. +""" + +from __future__ import annotations + +from typing import ClassVar, cast + +from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.jobs.api_factory import PlatformJobSpec +from nemo_rl_plugin.schema import RlJobInput +from nemo_rl_plugin.transform import transform_input_to_output +from nmp.customization_common.contributor.jobs import BaseSubmitJob, require_distributed_runtime +from nmp.rl.compile import platform_job_config_compiler +from nmp.rl.schemas import RlJobOutput +from pydantic import BaseModel + + +class RlJob(BaseSubmitJob): + """NeMo-RL DPO training job under the customization router (submit-only).""" + + name: ClassVar[str] = "rl.jobs" + description: ClassVar[str] = "NeMo-RL DPO training jobs on the platform Kubernetes GPU cluster (Ray)." + job_collection_path: ClassVar[str | None] = "/rl/jobs" + input_spec_schema: ClassVar[type[BaseModel] | None] = RlJobInput + spec_schema: ClassVar[type[BaseModel] | None] = RlJobOutput + docker_runtime_label: ClassVar[str] = "NeMo-RL" + + @classmethod + async def _transform(cls, job_input: BaseModel, workspace: str, async_sdk: AsyncNeMoPlatform) -> RlJobOutput: + return await transform_input_to_output(cast(RlJobInput, job_input), workspace, async_sdk) + + @classmethod + async def compile( + cls, + workspace: str, + spec: BaseModel, + entity_client: object, + job_name: str | None, + async_sdk: object, + profile: str | None = None, + options: dict | None = None, + ) -> PlatformJobSpec: + """Compile a validated :class:`RlJobOutput` into a 4-step Ray DPO job. + + Gates on ``platform.runtime: kubernetes`` — NeMo-RL provisions a Ray + cluster and has no local Docker fallback. An explicit + ``training.execution_profile`` (or the ``profile`` arg) wins; when both + are unset the compiler picks the topology-appropriate default + (single-node ``gpu`` vs multi-node ``gpu_distributed``). + """ + del entity_client, options + require_distributed_runtime(cls.docker_runtime_label) + canonical = spec if isinstance(spec, RlJobOutput) else RlJobOutput.model_validate(spec.model_dump()) + canonical.validate_for_training() + + # Leave ``None`` when unset so the compiler can default per topology. + execution_profile = canonical.training.execution_profile or profile + + return await platform_job_config_compiler( + workspace=workspace, + spec=canonical, + sdk=cast(AsyncNeMoPlatform, async_sdk), + job_name=job_name, + profile=execution_profile, + ) diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/schema.py b/plugins/nemo-rl/src/nemo_rl_plugin/schema.py new file mode 100644 index 0000000000..2cc4fa07db --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/schema.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Submitter-facing NeMo-RL schemas. + +The **canonical** types (``RlJobOutput``, ``DPOTraining``, ``OutputResponse``) +live in :mod:`nmp.rl.schemas` and are re-exported here for concise imports. Only +the thin input shape (``RlJobInput`` + ``OutputRequest``) is defined here; the +plugin's :func:`~nemo_rl_plugin.transform.transform_input_to_output` resolves it +into the canonical output. +""" + +from __future__ import annotations + +from nemo_platform_plugin.integrations import IntegrationsSpec +from nmp.rl.schemas import ( + DPOTraining, + OutputResponse, + ParallelismParams, + RlJobOutput, + TrainingMethod, +) +from pydantic import BaseModel, ConfigDict, Field + +__all__ = [ + "DPOTraining", + "OutputRequest", + "OutputResponse", + "ParallelismParams", + "RlJobInput", + "RlJobOutput", + "TrainingMethod", +] + + +class OutputRequest(BaseModel): + """Submitter-facing output preferences. ``name`` is auto-derived if omitted.""" + + model_config = ConfigDict(extra="forbid") + + name: str | None = None + + +class RlJobInput(BaseModel): + """POST body / CLI JSON for ``nemo customization rl submit``.""" + + model_config = ConfigDict(extra="forbid", protected_namespaces=()) + + name: str | None = None + model: str = Field(description="Model entity reference ('name' or 'workspace/name').") + dataset: str = Field( + description="Preference dataset fileset reference. Must contain training.jsonl + validation.jsonl.", + ) + training: TrainingMethod = Field(description="DPO training method and hyperparameters.") + integrations: IntegrationsSpec | None = None + output: OutputRequest | None = None diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/sdk/__init__.py b/plugins/nemo-rl/src/nemo_rl_plugin/sdk/__init__.py new file mode 100644 index 0000000000..2d688d793c --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/sdk/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NeMo-RL contributor SDK (mounted under ``client.customization`` by nemo-customizer).""" + +from nemo_rl_plugin.sdk.resources import ( + AsyncRlCustomization, + AsyncRlJobsResource, + RlCustomization, + RlJobsResource, +) + +__all__ = [ + "AsyncRlCustomization", + "AsyncRlJobsResource", + "RlCustomization", + "RlJobsResource", +] diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.py b/plugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.py new file mode 100644 index 0000000000..0d32993d3c --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/sdk/resources.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NeMo-RL contributor SDK resources (composed by ``nemo-customizer-plugin``). + +Thin shim over the shared :func:`nmp.customization_common.sdk.client.make_customization_sdk` +factory. ``RlCustomization`` / ``AsyncRlCustomization`` are imported by string by +the SDK hub and must not move. +""" + +from nmp.customization_common.sdk.client import make_customization_sdk + +RlCustomization, AsyncRlCustomization = make_customization_sdk("rl") + +# Jobs-resource classes re-exported for ``sdk/__init__.py`` and backward compatibility. +RlJobsResource = RlCustomization.jobs_resource_cls +AsyncRlJobsResource = AsyncRlCustomization.jobs_resource_cls diff --git a/plugins/nemo-rl/src/nemo_rl_plugin/transform.py b/plugins/nemo-rl/src/nemo_rl_plugin/transform.py new file mode 100644 index 0000000000..ba66c20177 --- /dev/null +++ b/plugins/nemo-rl/src/nemo_rl_plugin/transform.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Validate the platform refs (model entity + preference dataset fileset) +against the live SDK, resolve output naming, and return the +canonical :class:`~nmp.rl.schemas.RlJobOutput`. Only platform refs are +accepted — the container pipeline expects a real fileset to download from. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from nmp.customization_common.contributor.transform import generated_output_name +from nmp.customization_common.schemas.values import OutputNameType +from nmp.customization_common.service.platform_client import check_dataset_access, fetch_model_entity +from nmp.rl.schemas import OutputResponse, RlJobOutput + +from nemo_rl_plugin.schema import OutputRequest, RlJobInput + +if TYPE_CHECKING: + from nemo_platform import AsyncNeMoPlatform + + +async def transform_input_to_output( + input_spec: RlJobInput, + workspace: str, + sdk: "AsyncNeMoPlatform", +) -> RlJobOutput: + """Enrich submitter input into a canonical :class:`RlJobOutput`. + + Raises: + ValueError: When the model entity or dataset fileset cannot be resolved. + PermissionError: When access to the model or dataset is denied. + """ + # Strict refs: both calls error if the entity / fileset is missing. + model_entity = await fetch_model_entity(input_spec.model, workspace, sdk) + await check_dataset_access(sdk, input_spec.dataset, workspace) + + is_embedding = bool(model_entity.spec and getattr(model_entity.spec, "is_embedding_model", False)) + if is_embedding: + raise ValueError( + "DPO is not supported for embedding models. Use a causal LM model entity instead.", + ) + + output_request = input_spec.output or OutputRequest() + out_name = output_request.name or generated_output_name(input_spec.model, input_spec.dataset, workspace) + + output = OutputResponse( + name=out_name, + type=OutputNameType.MODEL, # DPO is full-weight + fileset=out_name, + ) + + return RlJobOutput( + name=input_spec.name, + model=input_spec.model, + dataset=input_spec.dataset, + training=input_spec.training, + integrations=input_spec.integrations, + output=output, + ) diff --git a/pyproject.toml b/pyproject.toml index 5c05c3d74f..1de30a21ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,6 +177,7 @@ enabled-plugins = [ "nemo-customizer-plugin", "nemo-automodel-plugin", "nemo-unsloth-plugin", + "nemo-rl-plugin", ] # Legacy runtime needed specifically for task images that still invoke @@ -375,8 +376,10 @@ nemo-agents-example-calculator = { workspace = true } nemo-customizer-plugin = { workspace = true } nemo-automodel-plugin = { workspace = true } nemo-unsloth-plugin = { workspace = true } +nemo-rl-plugin = { workspace = true } nmp-automodel = { workspace = true } nmp-unsloth = { workspace = true } +nmp-rl = { workspace = true } nmp-customization-common = { workspace = true } @@ -429,8 +432,10 @@ members = [ "plugins/nemo-customizer", "plugins/nemo-automodel", "plugins/nemo-unsloth", + "plugins/nemo-rl", "services/automodel", "services/unsloth", + "services/rl", ] @@ -566,10 +571,13 @@ exclude = [ "packages/garak_api/", - # GPU-container training drivers import torch/peft/nemo_automodel/unsloth at runtime only. + # GPU-container training drivers import torch/peft/nemo_automodel/unsloth/nemo_rl at runtime only. "services/automodel/src/nmp/automodel/tasks/training/backends/", "services/automodel/src/nmp/automodel/tasks/training/utils.py", "services/unsloth/src/nmp/unsloth/tasks/training/backends/", + "services/rl/src/nmp/rl/tasks/training/backends/", + "services/rl/src/nmp/rl/tasks/training/utils.py", + "services/rl/src/nmp/rl/tasks/training/runner.py", # Helm chart test payloads run inside purpose-built Kubernetes containers. # The NCCL worker imports torch at runtime, which is not part of the root # workspace type-checking environment. diff --git a/services/rl/README.md b/services/rl/README.md new file mode 100644 index 0000000000..4c5d85569e --- /dev/null +++ b/services/rl/README.md @@ -0,0 +1,31 @@ +# nmp-rl + +NeMo-RL task package for the NeMo Platform Customizer. Provides the compile glue +and the container-side tasks for **Direct Preference Optimization (DPO)** run on +a Ray cluster via [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL). The +training image bases on the published NGC container +`nvcr.io/nvidia/nemo-rl:v0.6.0` (amd64 + arm64, Python 3.13). + +No HTTP server. The thin contributor layer lives in +[`plugins/nemo-rl`](../../plugins/nemo-rl); this package holds: + +- `nmp.rl.schemas` — canonical `RlJobOutput` / `DPOTraining`. +- `nmp.rl.compile` / `nmp.rl.app.jobs.compiler` — the 4-step `PlatformJobSpec` + (download → DPO train → upload → model-entity). The training step's executor + is chosen by `parallelism.num_nodes` (single-node `gpu` vs multi-node + `gpu_distributed`). +- `nmp.rl.tasks.*` — container entrypoints (`file_io`, `model_entity`, + `training`). The training task bootstraps a Ray cluster and runs the DPO + driver against the NeMo-RL library. + +## Scope + +- **Remote Kubernetes only.** There is no local Docker fallback; `compile()` + requires `platform.runtime: kubernetes` (via + `require_distributed_runtime`). +- **Single-node multi-GPU and multi-node** are both supported. Multi-node + (`num_nodes > 1`) additionally requires a shared filesystem + (`NMP_RL_MULTINODE_SHARED_STORAGE_PATH`) for Ray's cross-node coordination; + `compile()` fails fast otherwise. +- **DPO is full-weight only** (PEFT unsupported). GRPO/PPO are reserved as + headroom in the schema/driver layout. diff --git a/services/rl/pyproject.toml b/services/rl/pyproject.toml new file mode 100644 index 0000000000..e426ae506f --- /dev/null +++ b/services/rl/pyproject.toml @@ -0,0 +1,58 @@ +[project] +name = "nmp-rl" +version = "0.1.0" +description = "NeMo-RL task package — Ray DPO driver and compile glue. No HTTP server." +readme = "README.md" +requires-python = ">=3.11,<3.14" +dependencies = [ + "nmp-common", + "nmp-customization-common", + "nemo-platform-plugin", + "nemo-platform-sdk", + "pydantic>=2.10.6", + "pydantic-settings>=2.6.1", + "httpx>=0.27.0", + "tenacity>=8.5.0", +] + +# Heavy ML deps (nemo-rl, ray, megatron, etc.) are NOT declared here: they are +# provided by the published NGC NeMo-RL container (v0.6.0) that +# docker/Dockerfile.nmp-rl-base builds on, with the platform glue installed on +# top — exactly like the reference customizer. The package stays importable +# without them (compile.py + schemas), so plugin discovery and the compiler pay +# no install cost. +[project.optional-dependencies] +integrations = [ + "wandb>=0.25.1", + "mlflow-skinny", +] + +[project.scripts] +# Container entrypoints. Names match the automodel/unsloth pattern for parity. +nmp-rl-training = "nmp.rl.tasks.training.__main__:main" +nmp-rl-file-io = "nmp.rl.tasks.file_io.run:run" +nmp-rl-model-entity = "nmp.rl.tasks.model_entity.__main__:run" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/nmp"] + +[tool.uv.sources] +nmp-common = { workspace = true } +nmp-customization-common = { workspace = true } +nemo-platform-plugin = { workspace = true } +nemo-platform-sdk = { workspace = true } + +[dependency-groups] +dev = [ + "pytest>=8.3.4", + "pytest-asyncio>=0.25.3", +] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +pythonpath = ["src"] +testpaths = ["tests"] diff --git a/services/rl/src/nmp/rl/app/constants.py b/services/rl/src/nmp/rl/app/constants.py new file mode 100644 index 0000000000..43656eaf0c --- /dev/null +++ b/services/rl/src/nmp/rl/app/constants.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Constants for the nmp-rl container job pipeline. + +Shared container-path/env constants come from +:mod:`nmp.customization_common.service.constants`; this module only adds the +nmp-rl ``SERVICE_NAME``, the training-output/workspace paths the runner uses, +and the ``BASE_LOG_DIR`` env name the Ray bootstrap reads for cross-node +coordination. +""" + +from nmp.customization_common.service.constants import ( + DEFAULT_DATASET_PATH, + DEFAULT_MODEL_PATH, + DEFAULT_OUTPUT_MODEL_PATH, + DEFAULT_VALIDATION_DATASET_PATH, + NMP_FILES_URL_ENVVAR, + NMP_JOBS_URL_ENVVAR, +) + +__all__ = [ + "BASE_LOG_DIR_ENVVAR", + "DEFAULT_DATASET_PATH", + "DEFAULT_MODEL_PATH", + "DEFAULT_OUTPUT_MODEL_PATH", + "DEFAULT_SEED", + "DEFAULT_TRAINING_OUTPUT_PATH", + "DEFAULT_TRAINING_RESULT_FILE_NAME", + "DEFAULT_VALIDATION_DATASET_PATH", + "NMP_FILES_URL_ENVVAR", + "NMP_JOBS_URL_ENVVAR", + "SERVICE_NAME", +] + +SERVICE_NAME = "rl" + +DEFAULT_SEED = 42 + +# File name the training runner writes the serialized TrainingResult to under +# the workspace path; downstream steps read it back. +DEFAULT_TRAINING_RESULT_FILE_NAME = "rl_training_result.json" + +# Workspace scratch dir the runner writes the compiled YAML, checkpoints, and +# training_result.json into. Single-node uses local scratch; multi-node points +# BASE_LOG_DIR at shared storage (see RlConfig.multinode_shared_storage_path). +DEFAULT_TRAINING_OUTPUT_PATH = "/var/run/scratch/job/training" + +# Env var the Ray bootstrap reads to locate the shared dir for the ENDED marker +# and barrier files across nodes. +BASE_LOG_DIR_ENVVAR = "BASE_LOG_DIR" diff --git a/services/rl/src/nmp/rl/app/jobs/compiler.py b/services/rl/src/nmp/rl/app/jobs/compiler.py new file mode 100644 index 0000000000..feda73b43d --- /dev/null +++ b/services/rl/src/nmp/rl/app/jobs/compiler.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Job compiler — transforms ``RlJobOutput`` into a 4-step ``PlatformJobSpec``. + +Steps mirror unsloth/automodel: + +1. file_io download — pull model fileset + preference dataset to the PVC +2. training — Ray DPO step (single-node GPU or multi-node distributed) +3. file_io upload — push the trained checkpoint to a new fileset +4. model_entity — create the output ``ModelEntity`` referencing it + +The training step's executor is selected by ``parallelism.num_nodes``: +``num_nodes == 1`` means a single-node ``GPUExecutionProviderSpec``; +``num_nodes > 1`` means a multi-node ``DistributedGPUExecutionProviderSpec``. +Multi-node additionally requires a shared filesystem for Ray's +cross-node ENDED/barrier coordination, enforced here with a fail-fast. +""" + +from __future__ import annotations + +import logging + +from nemo_platform import AsyncNeMoPlatform +from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.integrations import IntegrationsSpec +from nemo_platform_plugin.jobs.api_factory import ( + ContainerSpec, + CPUExecutionProviderSpec, + DistributedGPUExecutionProviderSpec, + EnvironmentVariable, + GPUExecutionProviderSpec, + PlatformJobSpec, + PlatformJobStep, + ResourcesLimitsSpec, + ResourcesRequestsSpec, + ResourcesSpec, +) +from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError +from nmp.common.jobs.constants import DEFAULT_JOB_STORAGE_PATH, PERSISTENT_JOB_STORAGE_PATH_ENVVAR +from nmp.customization_common.integrations import ( + collect_integration_secret_envs, + warn_incomplete_integrations, +) +from nmp.customization_common.schemas.file_io import ( + DownloadItem, + FileIOTaskConfig, + FileSetRef, + UploadItem, +) +from nmp.customization_common.schemas.model_entity import ModelEntityTaskConfig +from nmp.customization_common.service.platform_client import fetch_model_entity +from nmp.rl.app.constants import ( + BASE_LOG_DIR_ENVVAR, + DEFAULT_DATASET_PATH, + DEFAULT_MODEL_PATH, + DEFAULT_OUTPUT_MODEL_PATH, +) +from nmp.rl.app.jobs.training.schemas import ( + DPOConfig, + MLflowConfig, + ModelConfig, + TrainingBackend, + TrainingStepConfig, + WandBConfig, +) +from nmp.rl.config import config +from nmp.rl.entities.values import FinetuningType, TrainingType +from nmp.rl.images import RL_PYTHON_ENTRYPOINT, get_tasks_image, get_training_image +from nmp.rl.schemas import DPOTraining, RlJobOutput + +logger = logging.getLogger(__name__) + + +def _get_cpu_resources() -> ResourcesSpec: + return ResourcesSpec( + limits=ResourcesLimitsSpec( + cpu=config.default_job_resource_cpu_limit, + memory=config.default_job_resource_memory_limit, + ), + requests=ResourcesRequestsSpec( + cpu=config.default_job_resource_cpu_request, + memory=config.default_job_resource_memory_request, + ), + ) + + +def _base_environment() -> list[EnvironmentVariable]: + return [EnvironmentVariable(name=PERSISTENT_JOB_STORAGE_PATH_ENVVAR, value=DEFAULT_JOB_STORAGE_PATH)] + + +def _require_fileset(name: str | None, *, label: str) -> str: + if not name or not str(name).strip(): + raise PlatformJobCompilationError( + f"{label} has no fileset attached. Attach a platform FileSet (workspace/name) before training.", + ) + return str(name) + + +def _build_download_config(job_spec: RlJobOutput, me: ModelEntity, *, workspace: str) -> FileIOTaskConfig: + model_fileset = _require_fileset(me.fileset, label=f"Model '{me.workspace}/{me.name}'") + # The model ref is already workspace-qualified (from me.fileset), but the + # dataset ref comes straight from the submitted spec and may be a bare name + # ("dpo-data"). The file_io download step's SDK list() requires an explicit + # workspace, so qualify it with the job's workspace when unset. + dataset_ref = FileSetRef.model_validate(job_spec.dataset) + if dataset_ref.workspace is None: + dataset_ref = FileSetRef(workspace=workspace, name=dataset_ref.name) + return FileIOTaskConfig( + download=[ + DownloadItem(src=FileSetRef.model_validate(model_fileset), dest=DEFAULT_MODEL_PATH), + DownloadItem(src=dataset_ref, dest=DEFAULT_DATASET_PATH), + ], + ) + + +def _build_upload_config(output_fileset_name: str) -> FileIOTaskConfig: + return FileIOTaskConfig( + upload=[UploadItem(src=DEFAULT_OUTPUT_MODEL_PATH, dest=FileSetRef(workspace=None, name=output_fileset_name))], + ) + + +def _build_model_entity_config( + workspace: str, job_spec: RlJobOutput, *, trust_remote_code: bool +) -> ModelEntityTaskConfig: + # DPO is full-weight: no PEFT/adapter config. + return ModelEntityTaskConfig( + name=job_spec.output.name, + workspace=workspace, + description=f"DPO-trained model from nmp-rl job ({job_spec.model})", + fileset=FileSetRef(workspace=None, name=job_spec.output.fileset), + model_entity=job_spec.model, + base_model=job_spec.model, + peft=None, + trust_remote_code=trust_remote_code, + deployment_config=None, + ) + + +def _build_integrations_config(integrations: IntegrationsSpec | None) -> TrainingStepConfig.IntegrationsConfig: + """Map the public ``IntegrationsSpec`` onto the training step's ``IntegrationsConfig``. + + Without this the step config carries the empty default and the driver's + ``build_wandb_config`` / ``build_mlflow_config`` (and the backend's MLFLOW_URI + setup) all see ``None`` — silently disabling W&B/MLflow even when the job + requested them. Field names line up except MLflow's run name (public ``name`` + maps to the step's ``run_name``). Secrets (``api_key_secret``) are NOT copied + here; ``collect_integration_secret_envs`` injects them as env vars in + :func:`_build_training_step`. + """ + if integrations is None: + return TrainingStepConfig.IntegrationsConfig() + + wandb_cfg = None + if integrations.wandb is not None: + w = integrations.wandb + wandb_cfg = WandBConfig( + project=w.project, + name=w.name, + entity=w.entity, + tags=w.tags, + notes=w.notes, + base_url=w.base_url, + ) + + mlflow_cfg = None + if integrations.mlflow is not None: + m = integrations.mlflow + mlflow_cfg = MLflowConfig( + experiment_name=m.experiment_name, + run_name=m.name, + tags=m.tags, + description=m.description, + tracking_uri=m.tracking_uri, + ) + + return TrainingStepConfig.IntegrationsConfig(wandb=wandb_cfg, mlflow=mlflow_cfg) + + +def _build_training_step_config(job_spec: RlJobOutput, *, trust_remote_code: bool) -> TrainingStepConfig: + """Map the canonical DPO spec onto the backend-agnostic step config.""" + t: DPOTraining = job_spec.training + p = t.parallelism + return TrainingStepConfig( + backend=TrainingBackend.NEMO_RL, + model=ModelConfig( + path=DEFAULT_MODEL_PATH, + name=job_spec.model, + max_seq_length=t.max_seq_length, + trust_remote_code=trust_remote_code, + ), + dataset=TrainingStepConfig.DatasetConfig(path=DEFAULT_DATASET_PATH), + training=TrainingStepConfig.TrainingConfig( + training_type=TrainingType.DPO, + finetuning_type=FinetuningType.ALL_WEIGHTS, + dpo=DPOConfig( + ref_policy_kl_penalty=t.ref_policy_kl_penalty, + preference_average_log_probs=t.preference_average_log_probs, + sft_average_log_probs=t.sft_average_log_probs, + preference_loss_weight=t.preference_loss_weight, + sft_loss_weight=t.sft_loss_weight, + max_grad_norm=t.max_grad_norm, + ), + ), + schedule=TrainingStepConfig.ScheduleConfig( + epochs=t.epochs, + max_steps=t.max_steps, + val_check_interval=t.val_check_interval, + val_at_end=t.val_at_end, + keep_top_k=t.keep_top_k, + ), + batch=TrainingStepConfig.BatchConfig(global_batch_size=t.batch_size, micro_batch_size=t.micro_batch_size), + optimizer=TrainingStepConfig.OptimizerConfig( + optimizer_type=t.optimizer_type, + learning_rate=t.learning_rate, + min_learning_rate=t.min_learning_rate, + weight_decay=t.weight_decay, + beta1=t.adam_beta1, + beta2=t.adam_beta2, + eps=t.adam_eps, + warmup_steps=t.warmup_steps, + ), + parallelism=TrainingStepConfig.ParallelismConfig( + num_nodes=p.num_nodes, + num_gpus_per_node=p.num_gpus_per_node, + tensor_parallel_size=p.tensor_parallel_size, + pipeline_parallel_size=p.pipeline_parallel_size, + context_parallel_size=p.context_parallel_size, + sequence_parallel=p.sequence_parallel, + activation_checkpointing=t.activation_checkpointing, + ), + # Carry W&B / MLflow config into the step so the driver actually enables + # them; secrets are injected separately as env vars in _build_training_step. + integrations=_build_integrations_config(job_spec.integrations), + output_model=job_spec.output.name, + seed=t.seed if t.seed is not None else 42, + ) + + +def _build_training_step( + job_spec: RlJobOutput, + base_env: list[EnvironmentVariable], + *, + trust_remote_code: bool, + profile: str | None, +) -> PlatformJobStep: + """Build the Ray DPO training step, selecting the executor by ``num_nodes``. + + Multi-node (``num_nodes > 1``) requires shared storage for Ray's cross-node + ENDED/barrier coordination — fail fast when it is not configured. + """ + p = job_spec.training.parallelism + num_nodes = p.num_nodes + num_gpus_per_node = p.num_gpus_per_node + + step_config = _build_training_step_config(job_spec, trust_remote_code=trust_remote_code) + + container = ContainerSpec( + image=get_training_image(), + entrypoint=RL_PYTHON_ENTRYPOINT, + command=["-m", "nmp.rl.tasks.training"], + ) + + warn_incomplete_integrations(job_spec.integrations) + environment = [*base_env, *collect_integration_secret_envs(job_spec.integrations)] + + executor: GPUExecutionProviderSpec | DistributedGPUExecutionProviderSpec + if num_nodes > 1: + shared_dir = config.multinode_shared_storage_path + if not shared_dir: + raise PlatformJobCompilationError( + f"Multi-node NeMo-RL training (num_nodes={num_nodes}) requires a shared filesystem for Ray's " + "cross-node coordination. Set NMP_RL_MULTINODE_SHARED_STORAGE_PATH to a path mounted on every " + "node (e.g. an NFS mount) before submitting a multi-node job.", + ) + # Ray's bootstrap writes the ENDED marker + barriers under BASE_LOG_DIR. + environment = [*environment, EnvironmentVariable(name=BASE_LOG_DIR_ENVVAR, value=shared_dir)] + executor = { + "provider": "gpu_distributed", + "container": container, + "resources": ResourcesSpec(num_nodes=num_nodes, num_gpus=num_gpus_per_node), + } + resolved_profile = profile or config.default_distributed_execution_profile + else: + executor = { + "provider": "gpu", + "container": container, + "resources": ResourcesSpec(num_gpus=num_gpus_per_node), + } + resolved_profile = profile or config.default_training_execution_profile + + if resolved_profile is not None: + executor["profile"] = resolved_profile + + return PlatformJobStep( + name="dpo-training", + executor=executor, + environment=environment, + config=step_config.model_dump(mode="json"), + ) + + +async def platform_job_config_compiler( + workspace: str, + job_spec: RlJobOutput, + sdk: AsyncNeMoPlatform, + *, + job_name: str | None = None, + profile: str | None = None, +) -> PlatformJobSpec: + """Compile a canonical NeMo-RL job spec into a 4-step ``PlatformJobSpec``.""" + del job_name # reserved for future scheduling decisions + + # Log only non-sensitive, high-level context. The full spec embeds + # `integrations` (W&B / MLflow tokens and tracking URIs), so it must not be + # dumped at INFO. + p = job_spec.training.parallelism + logger.info( + "Compiling NeMo-RL DPO job to PlatformJobSpec: model=%s, dataset=%s, output=%s, " + "num_nodes=%d, num_gpus_per_node=%d", + job_spec.model, + job_spec.dataset, + job_spec.output.name, + p.num_nodes, + p.num_gpus_per_node, + ) + + me = await fetch_model_entity(job_spec.model, workspace, sdk) + trust_remote_code = me.trust_remote_code or False + + cpu_resources = _get_cpu_resources() + base_env = _base_environment() + + def _cpu_task_step( + name: str, command: str, task_config: FileIOTaskConfig | ModelEntityTaskConfig + ) -> PlatformJobStep: + return PlatformJobStep( + name=name, + executor=CPUExecutionProviderSpec( + provider="cpu", + container=ContainerSpec( + image=get_tasks_image(), + entrypoint=RL_PYTHON_ENTRYPOINT, + command=["-m", command], + ), + resources=cpu_resources, + ), + environment=base_env, + config=task_config.model_dump(mode="json"), + ) + + steps: list[PlatformJobStep] = [ + _cpu_task_step( + "model-and-dataset-download", + "nmp.rl.tasks.file_io", + _build_download_config(job_spec, me, workspace=workspace), + ), + _build_training_step(job_spec, base_env, trust_remote_code=trust_remote_code, profile=profile), + _cpu_task_step("model-upload", "nmp.rl.tasks.file_io", _build_upload_config(job_spec.output.fileset)), + _cpu_task_step( + "model-entity-creation", + "nmp.rl.tasks.model_entity", + _build_model_entity_config(workspace, job_spec, trust_remote_code=trust_remote_code), + ), + ] + + return PlatformJobSpec(steps=steps) diff --git a/services/rl/src/nmp/rl/app/jobs/training/schemas.py b/services/rl/src/nmp/rl/app/jobs/training/schemas.py new file mode 100644 index 0000000000..9234dd7733 --- /dev/null +++ b/services/rl/src/nmp/rl/app/jobs/training/schemas.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Backend-agnostic training-step config consumed by the container runner. + +The compiler serializes a :class:`TrainingStepConfig` into the training +``PlatformJobStep``; the runner deserializes it and ``dpo_config.compile_dpo_config`` +turns it into the NeMo-RL YAML. +""" + +from __future__ import annotations + +from enum import Enum + +from nmp.rl.app.constants import DEFAULT_OUTPUT_MODEL_PATH, DEFAULT_SEED, DEFAULT_TRAINING_OUTPUT_PATH +from nmp.rl.entities.values import CheckpointFormat, FinetuningType, Precision, TrainingType +from pydantic import BaseModel, Field + + +class TrainingBackend(str, Enum): + """Training backend identifier.""" + + NEMO_RL = "nemo_rl" + + +class OptimizerType(str, Enum): + """Optimizer and scheduler combination types.""" + + ADAMW_WITH_COSINE_ANNEALING = "adamw_with_cosine_annealing" + ADAM_WITH_COSINE_ANNEALING = "adam_with_cosine_annealing" + ADAMW_WITH_FLAT_LR = "adamw_with_flat_lr" + ADAM_WITH_FLAT_LR = "adam_with_flat_lr" + + +class ModelConfig(BaseModel): + """Internal model configuration with a resolved local path.""" + + path: str = Field(description="Local path to the downloaded model directory.") + name: str | None = Field(default=None, description="Model entity identifier.") + max_seq_length: int = Field(default=2048) + precision: Precision | None = Field(default=None, description="Weight dtype; auto-detected when None.") + chat_template: str | None = Field(default=None, description="Jinja2 chat template override.") + trust_remote_code: bool = Field(default=False) + + +class DPOConfig(BaseModel): + """DPO hyperparameters controlling the loss and optimization behavior.""" + + ref_policy_kl_penalty: float = Field(default=0.05, ge=0.0, description="KL penalty (beta in the DPO paper).") + preference_average_log_probs: bool = Field(default=False) + sft_average_log_probs: bool = Field(default=False) + preference_loss_weight: float = Field(default=1.0, ge=0.0) + sft_loss_weight: float = Field(default=0.0, ge=0.0) + max_grad_norm: float = Field(default=1.0, ge=0.0) + + +class WandBConfig(BaseModel): + project: str | None = None + name: str | None = None + entity: str | None = None + tags: list[str] | None = None + notes: str | None = None + base_url: str | None = None + + +class MLflowConfig(BaseModel): + experiment_name: str | None = None + run_name: str | None = None + tags: dict[str, str] | None = None + description: str | None = None + tracking_uri: str | None = None + + +class TrainingStepConfig(BaseModel): + """Normalized, backend-agnostic training configuration. + + The training container deserializes this and the NeMo-RL backend transforms + it into library-specific YAML at runtime. + """ + + class DatasetConfig(BaseModel): + path: str + prompt_template: str | None = None + add_bos: bool | None = None + add_eos: bool | None = None + + class TrainingConfig(BaseModel): + training_type: TrainingType + finetuning_type: FinetuningType | None = None + dpo: DPOConfig | None = None + + class ScheduleConfig(BaseModel): + epochs: int = 1 + max_steps: int | None = None + val_check_interval: float | None = None + val_at_end: bool = True + keep_top_k: int = 1 + + class BatchConfig(BaseModel): + global_batch_size: int = Field(default=32, gt=0) + micro_batch_size: int = Field(default=1, gt=0) + sequence_packing: bool = False + sequence_packing_max_samples: int = 1000 + + class OptimizerConfig(BaseModel): + optimizer_type: OptimizerType | None = Field(default=None) + learning_rate: float = 1e-4 + min_learning_rate: float | None = None + eps: float = 1e-5 + weight_decay: float = 0.01 + beta1: float = 0.9 + beta2: float = 0.999 + warmup_steps: int = 0 + + class ParallelismConfig(BaseModel): + num_nodes: int = 1 + num_gpus_per_node: int = 1 + tensor_parallel_size: int = 1 + pipeline_parallel_size: int = 1 + context_parallel_size: int = 1 + sequence_parallel: bool = False + activation_checkpointing: bool = False + + class IntegrationsConfig(BaseModel): + wandb: WandBConfig | None = None + mlflow: MLflowConfig | None = None + + # === Main config fields === + backend: TrainingBackend = TrainingBackend.NEMO_RL + model: ModelConfig + dataset: DatasetConfig + training: TrainingConfig + schedule: ScheduleConfig + batch: BatchConfig + optimizer: OptimizerConfig + parallelism: ParallelismConfig + integrations: IntegrationsConfig = Field(default_factory=IntegrationsConfig) + + # === Output paths === + output_model: str + workspace_path: str = Field(default=DEFAULT_TRAINING_OUTPUT_PATH) + output_path: str = Field(default=DEFAULT_OUTPUT_MODEL_PATH) + + # === Misc === + seed: int = Field(default=DEFAULT_SEED) + training_timeout: int | None = None + + +class GPUInfo(BaseModel): + architecture: str + device_name: str + memory_gb: float + cuda_version: str + + +class CheckpointInfo(BaseModel): + path: str + format: CheckpointFormat + precision: Precision | None = None + + +class TrainingMetrics(BaseModel): + final_loss: float | None = None + final_val_loss: float | None = None + best_val_loss: float | None = None + total_steps: int = 0 + total_epochs: int = 0 + + +class TrainingResult(BaseModel): + """Result written by the training task to ``{workspace_path}/training_result.json``.""" + + success: bool + error_message: str | None = None + checkpoint: CheckpointInfo | None = None + gpu_info: GPUInfo | None = None + metrics: TrainingMetrics = Field(default_factory=TrainingMetrics) + training_duration_seconds: float | None = None diff --git a/services/rl/src/nmp/rl/compile.py b/services/rl/src/nmp/rl/compile.py new file mode 100644 index 0000000000..558e3cf1e0 --- /dev/null +++ b/services/rl/src/nmp/rl/compile.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public compile entry for nmp-rl jobs. + +Mirror of :mod:`nmp.unsloth.compile`. Invoked by the plugin's ``RlJob.compile`` +to turn a validated :class:`~nmp.rl.schemas.RlJobOutput` into a 4-step +:class:`PlatformJobSpec` (download → DPO train → upload → model-entity). +""" + +from __future__ import annotations + +from nemo_platform import AsyncNeMoPlatform +from nemo_platform_plugin.jobs.api_factory import PlatformJobSpec +from nmp.rl.app.jobs.compiler import platform_job_config_compiler as _compile_canonical +from nmp.rl.schemas import RlJobOutput + + +async def platform_job_config_compiler( + *, + workspace: str, + spec: RlJobOutput, + sdk: AsyncNeMoPlatform, + job_name: str | None = None, + profile: str | None = None, +) -> PlatformJobSpec: + """Compile a canonical NeMo-RL job spec to a ``PlatformJobSpec``. Container submit only.""" + return await _compile_canonical(workspace, spec, sdk, job_name=job_name, profile=profile) + + +__all__ = ["platform_job_config_compiler"] diff --git a/services/rl/src/nmp/rl/config.py b/services/rl/src/nmp/rl/config.py new file mode 100644 index 0000000000..89837a1694 --- /dev/null +++ b/services/rl/src/nmp/rl/config.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration for the nmp-rl compiler and container tasks. + +Modeled after :mod:`nmp.unsloth.config`. Environment variables use the +``NMP_RL_`` prefix and drive image resolution and the multi-node shared-storage +gate for the ``PlatformJobSpec`` the plugin's ``RlJob.compile`` builds. +""" + +from nmp.common.config import create_service_config_class, get_platform_config, get_service_config +from pydantic import Field + + +class RlConfig(create_service_config_class("rl")): # type: ignore[misc] + """Environment variables use the ``NMP_RL_`` prefix.""" + + image_registry: str | None = Field( + default=None, + description=( + "Registry host/path prefix for nmp-rl-tasks and nmp-rl-training. " + "Override via NMP_RL_IMAGE_REGISTRY; defaults to the platform's image registry." + ), + ) + training_image: str | None = Field( + default=None, + description="Override entire GPU training image (registry/name:tag).", + ) + tasks_image: str | None = Field( + default=None, + description="Override entire CPU tasks image (registry/name:tag).", + ) + + default_job_resource_cpu_request: str = Field(default="1") + default_job_resource_memory_request: str = Field(default="8Gi") + default_job_resource_cpu_limit: str = Field(default="4") + default_job_resource_memory_limit: str = Field(default="16Gi") + + default_training_execution_profile: str = Field( + default="gpu", + description="Default single-node GPU profile when training.execution_profile is omitted.", + ) + default_distributed_execution_profile: str = Field( + default="gpu_distributed", + description="Default multi-node (num_nodes>1) distributed-GPU execution profile.", + ) + + multinode_shared_storage_path: str | None = Field( + default=None, + description=( + "Shared filesystem path (e.g. an NFS mount) used as BASE_LOG_DIR for Ray's " + "cross-node ENDED/barrier coordination. REQUIRED for multi-node jobs " + "(num_nodes>1); compile() fails fast when unset. Single-node jobs ignore it." + ), + ) + + +config = get_service_config(RlConfig) +platform_config = get_platform_config() diff --git a/services/rl/src/nmp/rl/entities/values.py b/services/rl/src/nmp/rl/entities/values.py new file mode 100644 index 0000000000..03a997e033 --- /dev/null +++ b/services/rl/src/nmp/rl/entities/values.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Value enums for the nmp-rl backend. + +``FinetuningType`` is reused from the common package. ``TrainingType`` is +per-backend (RL supports DPO today; GRPO reserved for headroom). +``Precision`` / ``CheckpointFormat`` are RL-local. +""" + +from enum import Enum + +from nmp.customization_common.schemas.values import FinetuningType # noqa: F401 (re-export) + + +class TrainingType(str, Enum): + """RL training algorithm. DPO is wired; GRPO is headroom. + + PPO is intentionally absent: the backend only compiles DPO/GRPO, so + accepting a PPO value here would defer the failure to a runtime crash inside + the training container instead of failing fast at request validation. + """ + + DPO = "dpo" + GRPO = "grpo" + + +class Precision(str, Enum): + """Model weight / compute precision.""" + + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class CheckpointFormat(str, Enum): + """Output checkpoint artifact format.""" + + SAFETENSORS = "safetensors" + HF = "hf" diff --git a/services/rl/src/nmp/rl/images.py b/services/rl/src/nmp/rl/images.py new file mode 100644 index 0000000000..7aea0caeb7 --- /dev/null +++ b/services/rl/src/nmp/rl/images.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Docker image resolution for nmp-rl job steps. + +Unlike unsloth (single image), nmp-rl follows the automodel split: a heavy +``nmp-rl-training`` image (NGC + NeMo-RL + Ray) for the GPU training step and a +lighter ``nmp-rl-tasks`` image for the CPU file_io / model_entity steps. Both +build on ``nmp-rl-base``. Override via ``NMP_RL_TRAINING_IMAGE`` / +``NMP_RL_TASKS_IMAGE``. +""" + +from __future__ import annotations + +from nmp.customization_common.service.images import resolve_qualified_image +from nmp.rl.config import config + +BASE_IMAGE_NAME = "nmp-rl-base" +TASKS_IMAGE_NAME = "nmp-rl-tasks" +TRAINING_IMAGE_NAME = "nmp-rl-training" + +# Must match ENTRYPOINT in Dockerfile.nmp-rl-{tasks,training}. Job specs set this +# explicitly: Docker API create() replaces the image entrypoint when passed []. +RL_PYTHON_ENTRYPOINT = ["/opt/venv/bin/python"] + + +def get_rl_qualified_image(name: str, override: str | None = None) -> str: + """Resolve a job step image reference (see ``resolve_qualified_image``).""" + return resolve_qualified_image(name, override, config.image_registry) + + +def get_tasks_image() -> str: + """CPU task steps (file_io, model_entity) — lighter image, no NeMo-RL/vLLM.""" + return get_rl_qualified_image(TASKS_IMAGE_NAME, config.tasks_image) + + +def get_training_image() -> str: + """GPU training step — NGC + NeMo-RL + Ray.""" + return get_rl_qualified_image(TRAINING_IMAGE_NAME, config.training_image) diff --git a/services/rl/src/nmp/rl/schemas.py b/services/rl/src/nmp/rl/schemas.py new file mode 100644 index 0000000000..348bde37d2 --- /dev/null +++ b/services/rl/src/nmp/rl/schemas.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical NeMo-RL schemas — consumed by the compiler and the DPO driver. + +Why these live in the service, not the plugin (same rationale as +:mod:`nmp.unsloth.schemas`): both compile-time +(:func:`nmp.rl.compile.platform_job_config_compiler`) and runtime (the DPO +driver) consume the canonical shape; the plugin's ``transform.py`` only +produces it from the thin ``RlJobInput``. + +Only DPO is wired today; the discriminated training union leaves a seam for +GRPO/PPO (see ``TrainingMethod``). +""" + +from __future__ import annotations + +from typing import Literal, Self + +from nemo_platform_plugin.integrations import IntegrationsSpec +from nmp.customization_common.schemas.values import OutputNameType +from nmp.rl.app.jobs.training.schemas import OptimizerType +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class ParallelismParams(BaseModel): + """Distributed training parallelism configuration. + + Single-node multi-GPU uses ``num_nodes=1`` with ``num_gpus_per_node>1``; + multi-node sets ``num_nodes>1`` and the compiler emits a distributed-GPU + executor (see :mod:`nmp.rl.app.jobs.compiler`). + """ + + model_config = ConfigDict(extra="forbid") + + num_gpus_per_node: int = Field(default=1, gt=0, description="Number of GPUs per node.") + num_nodes: int = Field(default=1, gt=0, description="Number of nodes (>1 → multi-node Ray cluster).") + tensor_parallel_size: int = Field(default=1, gt=0, description="Tensor parallel size.") + pipeline_parallel_size: int = Field(default=1, gt=0, description="Pipeline parallel size.") + context_parallel_size: int = Field(default=1, gt=0, description="Context parallel size.") + sequence_parallel: bool = Field(default=False, description="Enable sequence parallelism.") + + +class _TrainingBase(BaseModel): + """Common training configuration shared by all RL methods. + + Flat hyperparameters match the ML-practitioner mental model (HuggingFace + ``TrainingArguments`` / TRL configs). Only parallelism is grouped. + """ + + model_config = ConfigDict(protected_namespaces=(), extra="forbid") + + # --- Optimizer --- + optimizer_type: OptimizerType | None = Field( + default=None, + description="Optimizer + LR-scheduler combination (AdamW/Adam × cosine-annealing/flat-LR). " + "Defaults to AdamW with cosine annealing.", + ) + learning_rate: float = Field(default=1e-4, description="Peak learning rate.") + min_learning_rate: float | None = Field(default=None, description="Minimum LR for cosine decay.") + weight_decay: float = Field(default=0.01, description="Weight decay coefficient.") + adam_beta1: float = Field(default=0.9, description="Adam beta1.") + adam_beta2: float = Field(default=0.999, description="Adam beta2.") + adam_eps: float = Field(default=1e-5, gt=0.0, description="Adam epsilon (numerical stability term).") + warmup_steps: int = Field(default=0, ge=0, description="Linear warmup steps.") + + # --- Schedule --- + epochs: int = Field(default=1, gt=0, description="Number of passes through the dataset.") + max_steps: int | None = Field(default=None, gt=0, description="Max training steps (overrides epochs if set).") + val_check_interval: float | None = Field( + default=None, + description="Validation interval. Float <= 1.0 is fraction of epoch; > 1.0 is step count.", + ) + val_at_end: bool = Field( + default=True, + description="Run a final validation pass after the last training step. Keep enabled so the " + "final checkpoint carries validation metrics and best-checkpoint selection works; " + "set False only to skip the extra eval.", + ) + + # --- Checkpointing --- + keep_top_k: int = Field( + default=1, gt=0, description="Number of best checkpoints to retain (ranked by validation loss)." + ) + + # --- Batch --- + batch_size: int = Field(default=32, gt=0, description="Global batch size across all GPUs.") + micro_batch_size: int = Field(default=1, gt=0, description="Per-GPU micro batch size.") + activation_checkpointing: bool = Field( + default=False, + description="Recompute activations during the backward pass to reduce memory at the cost of compute. " + "Enable to fit larger models or longer sequences.", + ) + + # --- Model --- + max_seq_length: int = Field(default=2048, gt=0, description="Maximum token sequence length for training.") + seed: int | None = Field(default=None, description="Random seed for reproducibility.") + + # --- Infrastructure --- + parallelism: ParallelismParams = Field(default_factory=ParallelismParams) + execution_profile: str | None = Field( + default=None, + min_length=1, + description="Execution profile for the GPU training step (operator-configured). " + "Falls back to the service default when omitted.", + ) + + +class DPOTraining(_TrainingBase): + """Direct Preference Optimization (full-weight only — PEFT unsupported).""" + + type: Literal["dpo"] = "dpo" + ref_policy_kl_penalty: float = Field( + default=0.05, ge=0.0, description="KL penalty coefficient (beta in the DPO paper)." + ) + preference_average_log_probs: bool = Field( + default=False, description="Average log probabilities for preference loss calculation." + ) + sft_average_log_probs: bool = Field( + default=False, description="Average log probabilities for SFT regularization loss." + ) + preference_loss_weight: float = Field(default=1.0, ge=0.0, description="Weight for the preference (DPO) loss term.") + sft_loss_weight: float = Field( + default=0.0, ge=0.0, description="Weight for SFT regularization loss (0 = disabled)." + ) + max_grad_norm: float = Field(default=1.0, ge=0.0, description="Maximum gradient norm for clipping.") + + +# GRPO/PPO headroom. DPO carries a ``type: Literal["dpo"]`` discriminator field +# already, so when a second method lands this becomes: +# TrainingMethod = Annotated[Union[DPOTraining, GRPOTraining], Discriminator("type")] +# A single-member Union collapses to the member, and Discriminator requires a +# real Union — so today TrainingMethod is just DPOTraining. +TrainingMethod = DPOTraining + + +class _OutputBase(BaseModel): + name: str = Field( + max_length=255, + description="Name of the output artifact. Used to identify it during deployment and inference.", + examples=["my-dpo-llama"], + ) + + +class OutputRequest(_OutputBase): + """Output artifact configuration provided by the user.""" + + +class OutputResponse(_OutputBase): + """Resolved output artifact details.""" + + type: OutputNameType = Field( + default=OutputNameType.MODEL, + description="Output artifact type. DPO is full-weight, so always `model`.", + ) + fileset: str = Field( + max_length=255, + description="FileSet name where output artifacts are stored.", + ) + + +class RlJobOutput(BaseModel): + """Canonical NeMo-RL job spec (output of the plugin transform). + + The ``dataset`` fileset must contain ``training.jsonl`` and ``validation.jsonl`` + (any of the four supported preference formats); the dataset-preparation step + splits/normalizes them at runtime. + """ + + model_config = ConfigDict(protected_namespaces=()) + + name: str | None = Field(default=None, description="Optional job name; auto-generated when omitted.") + model: str = Field(description="Model entity reference ('name' or 'workspace/name').") + dataset: str = Field(description="Preference dataset fileset reference ('name' or 'workspace/name').") + training: TrainingMethod = Field(description="Training method and hyperparameters (DPO).") + integrations: IntegrationsSpec | None = Field(default=None, description="W&B / MLflow integrations.") + output: OutputResponse = Field(description="Output artifact created by this job.") + + def validate_for_training(self) -> None: + """Validate parallelism/batch consistency before compiling.""" + training = self.training + p = training.parallelism + total_gpus = p.num_gpus_per_node * p.num_nodes + model_parallel_size = p.tensor_parallel_size * p.pipeline_parallel_size * p.context_parallel_size + if total_gpus % model_parallel_size != 0: + raise ValueError( + f"Total GPUs ({total_gpus}) must be divisible by tensor_parallel_size " + f"({p.tensor_parallel_size}) * pipeline_parallel_size ({p.pipeline_parallel_size}) * " + f"context_parallel_size ({p.context_parallel_size}) = {model_parallel_size}" + ) + derived_dp = total_gpus // model_parallel_size + gb, mb = training.batch_size, training.micro_batch_size + divisor = mb * derived_dp + if gb % divisor != 0: + raise ValueError( + f"batch_size ({gb}) must be divisible by micro_batch_size ({mb}) * " + f"data_parallel_size ({derived_dp}) = {divisor}." + ) + + @model_validator(mode="after") + def _dpo_is_full_weight(self) -> Self: + # DPO is full-weight only; surface it early. + if self.output.type != OutputNameType.MODEL: + raise ValueError("DPO produces a full-weight model; output.type must be 'model'.") + return self diff --git a/services/rl/src/nmp/rl/tasks/file_io/__init__.py b/services/rl/src/nmp/rl/tasks/file_io/__init__.py new file mode 100644 index 0000000000..2d2c9ad611 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/file_io/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""File I/O task for rl customization jobs.""" + +from nmp.rl.tasks.file_io.run import run + +__all__ = ["run"] diff --git a/services/rl/src/nmp/rl/tasks/file_io/__main__.py b/services/rl/src/nmp/rl/tasks/file_io/__main__.py new file mode 100644 index 0000000000..c34e051201 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/file_io/__main__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sys + +from nmp.rl.tasks.file_io.run import run + +if __name__ == "__main__": + sys.exit(run()) diff --git a/services/rl/src/nmp/rl/tasks/file_io/callbacks.py b/services/rl/src/nmp/rl/tasks/file_io/callbacks.py new file mode 100644 index 0000000000..e72777198b --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/file_io/callbacks.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Custom fsspec callbacks for progress reporting during file I/O operations.""" + +import logging +import threading +from abc import abstractmethod +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from fsspec.callbacks import Callback, TqdmCallback +from nmp.common.jobs.schemas import PlatformJobStatus +from nmp.customization_common.schemas.file_io import DownloadStats, TaskPhase, UploadStats +from nmp.customization_common.tasks.file_io_progress_reporter import ProgressReporter +from nmp.customization_common.tasks.file_io_utils import list_local_files as _list_local_files + +logger = logging.getLogger(__name__) + + +def get_percentage(current: int, total: int) -> int: + """Get integer percentage 0-100, clamped to the valid range. + + Progress accounting must never abort the underlying transfer. Inputs + can fall outside ``[0, total]`` for benign reasons — most commonly when + the pre-transfer file listing under-counts a source that contains nested + directories, so the live ``current`` count exceeds ``total`` by the + number of nested files. Clamp rather than raise so a cosmetic progress + number can't fail a multi-GB download. + """ + if total <= 0: + return 0 + if current > total or current < 0: + # Benign (see docstring) but worth a breadcrumb now that we no longer + # raise — the old hard error is what previously surfaced count drift. + logger.debug("get_percentage clamping out-of-range progress: current=%s total=%s", current, total) + current = max(0, min(current, total)) + return int((current / total) * 100) + + +@dataclass +class FileInfo: + """A dataclass for file information.""" + + path: str + size: int + + +class TqdmPerFileUploadCallback(Callback): + """A callback that creates a separate tqdm progress bar for each file upload.""" + + def __init__(self, src_path: Path, **kwargs: Any): + self.src_path = src_path + super().__init__(**kwargs) + + def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> TqdmCallback: + if self.src_path.is_file(): + relative_path = self.src_path.name + else: + relative_path = Path(full_src_path).relative_to(self.src_path) + return TqdmCallback( + tqdm_kwargs={ + "desc": f"Uploading {relative_path!s}", + "unit": "B", + "unit_scale": True, + "unit_divisor": 1024, + "miniters": 1, + }, + ) + + +class TqdmPerFileDownloadCallback(Callback): + """A callback that creates a separate tqdm progress bar for each file download. + + Accepts a ``file_sizes`` dict (relative path -> byte size) so each + progress bar can show percent-complete even when the SDK streams the + file without a Content-Length header. + """ + + def __init__(self, dest_path: Path, fileset_path: str, file_sizes: dict[str, int] | None = None, **kwargs: Any): + self.dest_path = dest_path + self.fileset_path = fileset_path.rstrip("/") + self.file_sizes = file_sizes or {} + super().__init__(**kwargs) + + def branched(self, full_src_path: str, full_dest_path: str, **kwargs: Any) -> TqdmCallback: + dest_full_path = Path(full_dest_path) + if self.dest_path.is_file(): + relative_path = dest_full_path.name + else: + try: + relative_path = dest_full_path.relative_to(self.dest_path) + except ValueError: + relative_path = dest_full_path.name + + # full_src_path looks like "workspace/fileset/relative/path/file.txt". + # Strip the prefix to look up the size by relative path. + relative_file_path = full_src_path + if full_src_path.startswith(self.fileset_path): + relative_file_path = full_src_path[len(self.fileset_path) :].lstrip("/") + + file_size = self.file_sizes.get(relative_file_path) + + callback = TqdmCallback( + tqdm_kwargs={ + "desc": f"Downloading {relative_path!s}", + "unit": "B", + "unit_scale": True, + "unit_divisor": 1024, + "miniters": 1, + }, + ) + + # set_size() rather than tqdm_kwargs["total"] so the SDK can also + # call set_size() from a Content-Length header without conflict. + if file_size is not None: + callback.set_size(file_size) + + return callback + + +class BaseProgressCallback(Callback): + """Base class for file upload/download progress callbacks. + + Tracks file transfer progress and reports to the Jobs service. + Subclasses implement upload-vs-download behavior. + + Thread Safety: + Uses ``threading.Lock`` to protect stats updates because + FilesetFileSystem transfers files concurrently. + """ + + progress_reporter: ProgressReporter + fileset_name: str + total_files: int + total_size: int + stats: UploadStats | DownloadStats + _lock: threading.Lock + + def __init__( + self, + progress_reporter: ProgressReporter, + fileset_name: str, + total_files: int, + total_size: int, + stats: UploadStats | DownloadStats, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.progress_reporter = progress_reporter + self.fileset_name = str(fileset_name) + self.total_files = total_files + self.total_size = total_size + self.stats = stats + self._lock = threading.Lock() + + @staticmethod + def list_local_files(src_path: Path) -> list[FileInfo]: + """List all files under *src_path* (see shared ``list_local_files``).""" + return [FileInfo(path=f.path, size=f.size) for f in _list_local_files(src_path)] + + @abstractmethod + def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "BaseSingleFileCallback": + """Create a child callback for a single file transfer.""" + ... + + +class BaseSingleFileCallback(Callback): + """Base class for per-file callbacks within a batch operation. + + Uses the template-method pattern: ``close()`` runs the shared + state-update + progress-report sequence, while subclasses customize + via ``_get_phase``, ``_get_file_display_path``, ``_update_stats``, + ``_get_files_count``, and ``_build_status_details``. + """ + + parent: BaseProgressCallback + source_path: str + dest_path: str + _completed: bool + + def __init__( + self, + parent: BaseProgressCallback, + source_path: str, + dest_path: str, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.parent = parent + self.source_path = source_path + self.dest_path = dest_path + self._completed = False + + @abstractmethod + def _get_phase(self) -> str: + """Return the TaskPhase for this operation.""" + ... + + @abstractmethod + def _get_file_display_path(self) -> str: + """Return the path to use for display/logging.""" + ... + + @abstractmethod + def _update_stats(self) -> None: + """Update the parent's stats for this operation (called within lock).""" + ... + + @abstractmethod + def _get_files_count(self) -> int: + """Return the current files count from stats (called within lock).""" + ... + + @abstractmethod + def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: + """Build the status_details dict for progress reporting.""" + ... + + def close(self) -> None: + """Called when the file transfer completes.""" + if self._completed: + return + + self._completed = True + parent = self.parent + current_file = self._get_file_display_path() + + with parent._lock: + self._update_stats() + files_count = self._get_files_count() + total_bytes = parent.stats.total_bytes + + logger.debug(f"File transferred: {current_file} ({files_count}/{parent.total_files})") + + # Report outside the lock — don't block other threads on the network call. + parent.progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details=self._build_status_details(files_count, total_bytes, current_file), + ) + + def __enter__(self) -> "BaseSingleFileCallback": + return self + + def __exit__(self, *exc_args: object) -> None: + self.close() + + +class FileUploadProgressCallback(BaseProgressCallback): + """Callback for tracking file upload progress and reporting to the Jobs service.""" + + stats: UploadStats + + def __init__( + self, + progress_reporter: ProgressReporter, + src_path: Path, + fileset_name: str, + stats: UploadStats, + **kwargs: Any, + ): + files = self.list_local_files(src_path) + if not files: + logger.warning(f"Source path {src_path} contains no files") + total_files = len(files) + total_size = sum(f.size for f in files) + + super().__init__( + progress_reporter=progress_reporter, + fileset_name=fileset_name, + total_files=total_files, + total_size=total_size, + stats=stats, + **kwargs, + ) + + logger.info(f"Uploading {total_files} files ({total_size} bytes) to {self.fileset_name}") + + progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details={ + "phase": TaskPhase.UPLOADING, + "fileset": self.fileset_name, + "total_files": total_files, + "total_size": total_size, + "uploaded_files": 0, + "uploaded_bytes": 0, + }, + ) + + def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "SingleFileUploadCallback": + return SingleFileUploadCallback( + parent=self, + source_path=source_path, + dest_path=dest_path, + **kwargs, + ) + + +class SingleFileUploadCallback(BaseSingleFileCallback): + """Per-file upload callback. Notifies parent on completion.""" + + parent: FileUploadProgressCallback + + def _get_phase(self) -> str: + return TaskPhase.UPLOADING + + def _get_file_display_path(self) -> str: + return self.dest_path.split("/")[-1] if "/" in self.dest_path else self.dest_path + + def _update_stats(self) -> None: + self.parent.stats.files_uploaded += 1 + if self.size is not None: + self.parent.stats.total_bytes += self.size + + def _get_files_count(self) -> int: + return self.parent.stats.files_uploaded + + def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: + return { + "phase": TaskPhase.UPLOADING, + "fileset": self.parent.fileset_name, + "total_files": self.parent.total_files, + "total_size": self.parent.total_size, + "uploaded_files": files_count, + "uploaded_bytes": total_bytes, + "current_file": current_file, + "progress_pct": get_percentage(files_count, self.parent.total_files), + } + + +class FileDownloadProgressCallback(BaseProgressCallback): + """Callback for tracking file download progress and reporting to the Jobs service.""" + + stats: DownloadStats + + def __init__( + self, + progress_reporter: ProgressReporter, + fileset_name: str, + total_files: int, + total_size: int, + stats: DownloadStats, + **kwargs: Any, + ): + super().__init__( + progress_reporter=progress_reporter, + fileset_name=fileset_name, + total_files=total_files, + total_size=total_size, + stats=stats, + **kwargs, + ) + + logger.info(f"Downloading {total_files} files ({total_size} bytes) from {self.fileset_name}") + + progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details={ + "phase": TaskPhase.DOWNLOADING, + "fileset": self.fileset_name, + "total_files": total_files, + "total_size": total_size, + "downloaded_files": 0, + "downloaded_bytes": 0, + }, + ) + + def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "SingleFileDownloadCallback": + return SingleFileDownloadCallback( + parent=self, + source_path=source_path, + dest_path=dest_path, + **kwargs, + ) + + +class SingleFileDownloadCallback(BaseSingleFileCallback): + """Per-file download callback. Notifies parent on completion.""" + + parent: FileDownloadProgressCallback + + def _get_phase(self) -> str: + return TaskPhase.DOWNLOADING + + def _get_file_display_path(self) -> str: + return self.source_path.split("/")[-1] if "/" in self.source_path else self.source_path + + def _update_stats(self) -> None: + self.parent.stats.files_downloaded += 1 + if self.size is not None: + self.parent.stats.total_bytes += self.size + + def _get_files_count(self) -> int: + return self.parent.stats.files_downloaded + + def _build_status_details(self, files_count: int, total_bytes: int, current_file: str) -> dict[str, Any]: + return { + "phase": TaskPhase.DOWNLOADING, + "fileset": self.parent.fileset_name, + "total_files": self.parent.total_files, + "total_size": self.parent.total_size, + "downloaded_files": files_count, + "downloaded_bytes": total_bytes, + "current_file": current_file, + "progress_pct": get_percentage(files_count, self.parent.total_files), + } + + +class CompositeCallback(Callback): + """A callback that delegates to multiple child callbacks. + + Lets us combine console-side ``TqdmCallback`` and Jobs-service + ``File{Upload,Download}ProgressCallback`` into one callback object + passed to fsspec operations. + """ + + def __init__(self, *callbacks: Callback, **kwargs: Any): + super().__init__(**kwargs) + self.callbacks = list(callbacks) + + def set_size(self, size: int) -> None: + self.size = size + for cb in self.callbacks: + cb.set_size(size) + + def absolute_update(self, value: int) -> None: + self.value = value + for cb in self.callbacks: + cb.absolute_update(value) + + def relative_update(self, inc: int = 1) -> None: + self.value += inc + for cb in self.callbacks: + cb.relative_update(inc) + + def branched(self, source_path: str, dest_path: str, **kwargs: Any) -> "CompositeCallback": + child_callbacks = [cb.branched(source_path, dest_path, **kwargs) for cb in self.callbacks] + return CompositeCallback(*child_callbacks) + + def call(self, hook_name: str | None = None, **kwargs: Any) -> None: + for cb in self.callbacks: + cb.call(hook_name, **kwargs) + + def close(self) -> None: + for cb in self.callbacks: + cb.close() + + def __enter__(self) -> "CompositeCallback": + for cb in self.callbacks: + cb.__enter__() + return self + + def __exit__(self, *exc_args: object) -> None: + for cb in self.callbacks: + cb.__exit__(*exc_args) diff --git a/services/rl/src/nmp/rl/tasks/file_io/run.py b/services/rl/src/nmp/rl/tasks/file_io/run.py new file mode 100644 index 0000000000..581567d2c2 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/file_io/run.py @@ -0,0 +1,537 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""File I/O task entry point. + +Handles file operations between NeMo Platform Files Service and the job's shared PVC. + +The task reads configuration and performs: +- Downloads: If config.download is non-empty, download files from FileSets to local paths +- Uploads: If config.upload is non-empty, upload files from local paths to FileSets + +Usage: + export NEMO_JOB_STEP_CONFIG_FILE_PATH= + python -m nmp.rl.tasks.file_io +""" + +import logging +from pathlib import Path + +import httpx + +# https://docs.nvidia.com/nemo/microservices/latest/pysdk/index.html#handling-errors +from nemo_platform import ( + APIConnectionError, + APITimeoutError, + ConflictError, + InternalServerError, + NeMoPlatform, + NotFoundError, +) +from nemo_platform.types.files.fileset_file import FilesetFile +from nmp.common.jobs.schemas import PlatformJobStatus +from nmp.common.sdk_factory import get_task_sdk +from nmp.customization_common.schemas.file_io import ( + DownloadItem, + DownloadStats, + FileDownloadError, + FileSetRef, + FileUploadError, + PathTraversalError, + TaskPhase, + UploadItem, + UploadStats, +) +from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.tasks.file_io_progress_reporter import JobsServiceProgressReporter, ProgressReporter +from nmp.customization_common.tasks.file_io_utils import ( + filesystem_sdk_error_handler, + get_config, + sdk_error_handler, + validate_safe_path, + validate_storage_path, +) +from nmp.rl.app.constants import SERVICE_NAME +from nmp.rl.tasks.file_io.callbacks import ( + CompositeCallback, + FileDownloadProgressCallback, + FileUploadProgressCallback, + TqdmPerFileDownloadCallback, + TqdmPerFileUploadCallback, +) +from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential + +logger = logging.getLogger(__name__) + +# Service-source tag stamped onto every upload-created fileset. Lets operators +# filter filesets by training backend. +SERVICE_SOURCE = "rl" + +# Timeout configurations for SDK operations (httpx.Timeout for API calls). +CREATE_FILESET_TIMEOUT = httpx.Timeout(10.0, connect=10.0) +LIST_FILES_TIMEOUT = httpx.Timeout(10.0, connect=10.0) + +# Timeout configurations for FilesetFileSystem operations. Passed via +# sdk.with_options(timeout=...). httpx.Timeout(read=...) is per-chunk +# (the SDK chunks at 16MB), NOT total transfer time — it's a socket-level +# timeout. SDK defaults are httpx.Timeout(timeout=60, connect=5.0). +DOWNLOAD_TIMEOUT = httpx.Timeout(30.0, read=5 * 60) +UPLOAD_TIMEOUT = httpx.Timeout(30.0, write=10 * 60, read=5 * 60) + +# Retry configuration. +MAX_RETRIES = 3 +INITIAL_BACKOFF_SECONDS = 1.0 +MAX_BACKOFF_SECONDS = 30.0 + +# Transient exceptions that should trigger retries for filesystem operations. +# FilesetFileSystem uses httpx under the hood, so we retry httpx transients +# in addition to SDK-level transients. +TRANSIENT_FILESYSTEM_EXCEPTIONS = ( + httpx.TimeoutException, + httpx.ConnectError, + httpx.ReadTimeout, + # Connection dropped mid-transfer (CDN/proxy closed the socket before the + # full body arrived). Common on large multi-GB model shards; safe to retry. + httpx.RemoteProtocolError, + httpx.ReadError, +) + + +class FileIORunner: + """Runner for file I/O operations against the Files service.""" + + def __init__( + self, + sdk: NeMoPlatform, + progress_reporter: ProgressReporter, + job_ctx: NMPJobContext, + ): + self.sdk = sdk + self.progress_reporter = progress_reporter + self.job_ctx = job_ctx + + def list_fileset_files(self, fileset: FileSetRef) -> list[FilesetFile]: + """List files in a FileSet. Returns a list of ``FilesetFile`` objects.""" + try: + with sdk_error_handler(FileDownloadError, f"list files in fileset {fileset}", passthrough=(NotFoundError,)): + response = self.sdk.with_options(timeout=LIST_FILES_TIMEOUT).files.list( + fileset=fileset.name, + workspace=fileset.workspace, + ) + logger.info(f"Found {len(response.data)} files in FileSet {fileset!s}") + return response.data + except NotFoundError as e: + raise FileDownloadError( + f"FileSet {fileset!s} not found. Please ensure the FileSet exists and contains the expected files.", + ) from e + + def download_fileset(self, fileset: FileSetRef, dest_dir: Path) -> DownloadStats: + """Download all files from a FileSet to a destination directory. + + Uses ``FilesetFileSystem.get()`` with ``recursive=True`` for efficient batch + downloads. Progress is tracked via two callbacks combined in a + ``CompositeCallback``: + + - ``TqdmPerFileDownloadCallback`` — separate console progress bar per file + - ``FileDownloadProgressCallback`` — reports to Jobs service after each file + + Raises: + FileDownloadError: If the download fails. + """ + fileset_name = str(fileset) + + files = self.list_fileset_files(fileset) + + if not files: + logger.warning(f"FileSet {fileset_name} contains no files") + return DownloadStats() + + total_files = len(files) + total_size = sum(f.size for f in files) + + dest_dir.mkdir(parents=True, exist_ok=True) + + # Maps relative file paths to byte sizes for tqdm percent display. + file_sizes = {f.path.lstrip("/"): f.size for f in files} + + with filesystem_sdk_error_handler( + FileDownloadError, + f"download from '{fileset_name}' to '{dest_dir}'", + ): + # Progress state (stats + callbacks) is built inside the retried call so + # each attempt starts from zero. A retry restarts the transfer from + # scratch, so reusing a single stats/callback object across attempts + # would double-count earlier progress and inflate the reported totals. + stats = self._download_with_retry( + fileset_name=fileset.name, + fileset_workspace=fileset.workspace, + dest_dir=str(dest_dir), + fileset_display_name=fileset_name, + dest_path=dest_dir, + file_sizes=file_sizes, + total_files=total_files, + total_size=total_size, + ) + + logger.info(f"Download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") + return stats + + @retry( + stop=stop_after_attempt(MAX_RETRIES), + wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), + retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), + reraise=True, + before_sleep=before_sleep_log(logger, logging.WARNING), + ) + def _download_with_retry( + self, + fileset_name: str, + fileset_workspace: str | None, + dest_dir: str, + fileset_display_name: str, + dest_path: Path, + file_sizes: dict[str, int], + total_files: int, + total_size: int, + ) -> DownloadStats: + """Internal method with retry logic for downloading from FilesetFileSystem. + + Builds fresh ``DownloadStats`` and callbacks on every attempt so retried + transfers report progress from zero rather than accumulating across attempts. + """ + stats = DownloadStats() + tqdm_callback = TqdmPerFileDownloadCallback( + dest_path=dest_path, + fileset_path=fileset_display_name, + file_sizes=file_sizes, + ) + jobs_callback = FileDownloadProgressCallback( + progress_reporter=self.progress_reporter, + fileset_name=fileset_display_name, + total_files=total_files, + total_size=total_size, + stats=stats, + ) + composite_callback = CompositeCallback(tqdm_callback, jobs_callback) + + self.sdk.with_options(timeout=DOWNLOAD_TIMEOUT).files.download( + fileset=fileset_name, + workspace=fileset_workspace, + local_path=dest_dir, + callback=composite_callback, + ) + return stats + + def upload_fileset(self, fileset: FileSetRef, src_path: Path) -> UploadStats: + """Upload all files from a source path (file or directory) to a FileSet. + + Uses ``FilesetFileSystem.put()`` with ``recursive=True`` for efficient batch + uploads. Progress is tracked via the same composite-callback pattern as + downloads. + + Raises: + FileUploadError: If the upload fails. + """ + fileset_name = str(fileset) + + # Build local and remote paths for upload. ``remote_path`` is relative within + # the fileset ("" for root, "filename" for single file). Trailing slash on + # ``local_path`` follows rsync/scp convention: "dir/" copies contents, + # "dir" copies the directory itself. + if src_path.is_dir(): + local_path = f"{src_path}/" + remote_path = "" + else: + local_path = str(src_path) + remote_path = src_path.name + + with filesystem_sdk_error_handler( + FileUploadError, + f"upload from '{src_path}' to '{fileset_name}'", + ): + # Fresh progress state per attempt — see _download_with_retry for why + # reusing stats/callbacks across retries double-counts progress. + stats = self._upload_with_retry( + local_path=local_path, + remote_path=remote_path, + fileset_name=fileset.name, + fileset_workspace=fileset.workspace, + fileset_display_name=fileset_name, + src_path=src_path, + ) + + logger.info(f"Upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") + return stats + + @retry( + stop=stop_after_attempt(MAX_RETRIES), + wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), + retry=retry_if_exception_type(TRANSIENT_FILESYSTEM_EXCEPTIONS), + reraise=True, + before_sleep=before_sleep_log(logger, logging.WARNING, exc_info=True), + ) + def _upload_with_retry( + self, + local_path: str, + remote_path: str, + fileset_name: str, + fileset_workspace: str | None, + fileset_display_name: str, + src_path: Path, + ) -> UploadStats: + """Internal method with retry logic for uploading to FilesetFileSystem. + + Builds fresh ``UploadStats`` and callbacks on every attempt so retried + transfers report progress from zero rather than accumulating across attempts. + """ + stats = UploadStats() + tqdm_callback = TqdmPerFileUploadCallback(src_path=src_path) + jobs_callback = FileUploadProgressCallback( + progress_reporter=self.progress_reporter, + src_path=src_path, + fileset_name=fileset_display_name, + stats=stats, + ) + composite_callback = CompositeCallback(tqdm_callback, jobs_callback) + + self.sdk.with_options(timeout=UPLOAD_TIMEOUT).files.upload( + local_path=local_path, + remote_path=remote_path, + fileset=fileset_name, + workspace=fileset_workspace, + callback=composite_callback, + ) + return stats + + def create_fileset(self, fileset: FileSetRef, metadata: dict | None = None) -> None: + """Create a FileSet. Skip if it already exists. + + Wraps the retry with ``sdk_error_handler`` to convert exceptions after + all retries exhaust. + """ + with sdk_error_handler(FileUploadError, f"create fileset {fileset}", passthrough=(ConflictError,)): + self._create_fileset_with_retry(fileset, metadata) + + # We don't use the SDK's built-in retry: it would retry on ConflictError, + # which is expected here and would just waste calls. + @retry( + stop=stop_after_attempt(MAX_RETRIES), + wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), + retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), + reraise=True, + ) + def _create_fileset_with_retry(self, fileset: FileSetRef, metadata: dict | None = None) -> None: + """Internal method with retry logic for creating a FileSet.""" + try: + create_kwargs: dict = { + "workspace": fileset.workspace, + "name": fileset.name, + "timeout": CREATE_FILESET_TIMEOUT, + "custom_fields": {"service_source": SERVICE_SOURCE}, + } + if metadata is not None: + create_kwargs["metadata"] = metadata + result = self.sdk.with_options(max_retries=0).files.filesets.create(**create_kwargs) + logger.info(f"Created FileSet: {result.workspace}/{result.name}") + except ConflictError: + # Fileset already exists — patch metadata so tool_calling etc. aren't lost. + workspace = fileset.workspace or self.job_ctx.workspace + if metadata is not None: + update_kwargs: dict = { + "name": fileset.name, + "workspace": workspace, + "metadata": metadata, + "timeout": CREATE_FILESET_TIMEOUT, + } + try: + self.sdk.with_options(max_retries=0).files.filesets.update(**update_kwargs) + logger.info(f"Patched existing FileSet metadata: {workspace}/{fileset.name}") + except Exception as e: + logger.warning( + f"Could not patch metadata on existing fileset {workspace}/{fileset.name}: {e}. " + "Upload will continue; downstream consumers may lack the latest metadata.", + ) + + def run_download(self, downloads: list[DownloadItem]) -> None: + """Execute download operations.""" + if not downloads: + logger.info("No downloads configured, skipping download operation") + return + + storage_path = validate_storage_path(self.job_ctx.storage_path) + + logger.info(f"Starting download operation: {len(downloads)} fileset(s) to download") + + self.progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details={ + "phase": TaskPhase.DOWNLOADING, + "total_filesets": len(downloads), + "completed_filesets": 0, + }, + ) + + total_stats = DownloadStats() + + for idx, item in enumerate(downloads): + fileset = item.src + dest_dir = validate_safe_path(storage_path, item.dest) + + logger.info(f"[{idx + 1}/{len(downloads)}] Downloading from {fileset!s} to {dest_dir}") + + self.progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details={ + "phase": TaskPhase.DOWNLOADING, + "total_filesets": len(downloads), + "completed_filesets": idx, + "current_fileset": f"{fileset!s}", + }, + ) + + stats = self.download_fileset(fileset, dest_dir) + total_stats.files_downloaded += stats.files_downloaded + total_stats.total_bytes += stats.total_bytes + + logger.info(f"FileSet download complete: {stats.files_downloaded} files, {stats.total_bytes} bytes") + + logger.info( + f"All downloads complete: {total_stats.files_downloaded} files, {total_stats.total_bytes} bytes total", + ) + + def run_upload(self, uploads: list[UploadItem]) -> None: + """Execute upload operations.""" + if not uploads: + logger.info("No uploads configured, skipping upload operation") + return + + storage_path = validate_storage_path(self.job_ctx.storage_path) + + logger.info(f"Starting upload operation: {len(uploads)} fileset(s) to upload") + + self.progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details={ + "phase": TaskPhase.UPLOADING, + "total_filesets": len(uploads), + "completed_filesets": 0, + }, + ) + + total_stats = UploadStats() + + for idx, item in enumerate(uploads): + if item.dest.workspace is None: + item.dest.workspace = self.job_ctx.workspace + fileset = item.dest + src_path = validate_safe_path(storage_path, item.src) + if not src_path.exists(): + raise FileUploadError(f"Source path does not exist: {src_path}. Ensure the source path exists.") + if not src_path.is_dir() and not src_path.is_file(): + raise FileUploadError( + f"Source path is not a file or directory: {src_path}. " + "Ensure the source path is a file or directory.", + ) + + logger.info(f"[{idx + 1}/{len(uploads)}] Uploading from {src_path} to {fileset!s}") + + self.progress_reporter.update_progress( + status=PlatformJobStatus.ACTIVE, + status_details={ + "phase": TaskPhase.UPLOADING, + "total_filesets": len(uploads), + "completed_filesets": idx, + "current_fileset": str(fileset), + }, + ) + + self.create_fileset(fileset, metadata=item.metadata) + + stats = self.upload_fileset(fileset, src_path) + total_stats.files_uploaded += stats.files_uploaded + total_stats.total_bytes += stats.total_bytes + + logger.info(f"FileSet upload complete: {stats.files_uploaded} files, {stats.total_bytes} bytes") + + logger.info(f"All uploads complete: {total_stats.files_uploaded} files, {total_stats.total_bytes} bytes total") + + +def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: + """Execute the file I/O task. + + Args: + sdk: Optional SDK instance for dependency injection (for testing). + If None, creates one via get_task_sdk(). + job_ctx: Optional job context for dependency injection (for testing). + If None, creates one via NMPJobContext.from_env(). + + Returns: + Exit code (0 for success, non-zero for failure). + """ + job_ctx = job_ctx or NMPJobContext.from_env() + validate_storage_path(job_ctx.storage_path) + + sdk_owned = sdk is None + progress_reporter: ProgressReporter | None = None + try: + sdk = sdk or get_task_sdk(SERVICE_NAME) + progress_reporter = JobsServiceProgressReporter.create_progress_reporter(sdk, job_ctx) + runner = FileIORunner(sdk=sdk, progress_reporter=progress_reporter, job_ctx=job_ctx) + + config = get_config(job_ctx.config_path) + + logger.info(f"Starting file I/O task with job context: {job_ctx}") + logger.info(f"Config: {config.model_dump_json(indent=2)}") + logger.info(f"NeMo Platform service URL: {sdk.base_url}") + + runner.run_upload(config.upload) + runner.run_download(config.download) + + progress_reporter.update_progress( + status=PlatformJobStatus.COMPLETED, + status_details={"phase": TaskPhase.COMPLETED, "message": "File I/O task completed successfully"}, + ) + + return 0 + except PathTraversalError as e: + logger.error(f"Security error - path traversal detected: {e}") + if progress_reporter: + progress_reporter.update_progress( + status=PlatformJobStatus.ERROR, + error_details={"message": str(e), "type": type(e).__name__}, + ) + return 1 + except (FileDownloadError, FileUploadError) as e: + logger.exception(f"File operation failed: {e}") + if progress_reporter: + progress_reporter.update_progress( + status=PlatformJobStatus.ERROR, + error_details={"message": str(e), "type": type(e).__name__}, + ) + return 1 + except Exception as e: + logger.exception(f"File I/O task failed: {e}") + if progress_reporter: + progress_reporter.update_progress( + status=PlatformJobStatus.ERROR, + error_details={"message": str(e), "type": type(e).__name__}, + ) + return 1 + finally: + if sdk_owned and sdk is not None: + sdk.close() + + +def build_output_metadata(spec) -> dict: + """Build the metadata dict stamped onto the output fileset. + + Captures the bits a downstream consumer (model-entity creation, + deployment) needs about this artefact without re-deriving them + from the training spec. + """ + return { + "model": spec.model.name, + "finetuning_type": spec.training.finetuning_type, + "save_method": spec.output.save_method, + "output_type": spec.output.type, + } diff --git a/services/rl/src/nmp/rl/tasks/model_entity/__init__.py b/services/rl/src/nmp/rl/tasks/model_entity/__init__.py new file mode 100644 index 0000000000..ebe5a7cf8f --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/model_entity/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model entity task for creating model entities after rl customization.""" + +from nmp.rl.tasks.model_entity.run import run + +__all__ = ["run"] diff --git a/services/rl/src/nmp/rl/tasks/model_entity/__main__.py b/services/rl/src/nmp/rl/tasks/model_entity/__main__.py new file mode 100644 index 0000000000..3a2739e1b9 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/model_entity/__main__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Entry point for model_entity task. + +Usage: + python -m nmp.rl.tasks.model_entity +""" + +import sys + +from .run import run + +if __name__ == "__main__": + sys.exit(run()) diff --git a/services/rl/src/nmp/rl/tasks/model_entity/run.py b/services/rl/src/nmp/rl/tasks/model_entity/run.py new file mode 100644 index 0000000000..acb3888e0a --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/model_entity/run.py @@ -0,0 +1,477 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model entity task entry point. + +Handles creating model entities in the Models service after customization completes. + +The task reads configuration and creates a Model Entity that references the +uploaded model artifacts in the Files service. When ``deployment_config`` is set +on the task config, the task also launches an inference deployment. + +Usage: + export NEMO_JOB_STEP_CONFIG_FILE_PATH= + python -m nmp.rl.tasks.model_entity +""" + +import json +import logging +import re +import time +from pathlib import Path + +from nemo_platform import ( + APIConnectionError, + APITimeoutError, + ConflictError, + InternalServerError, + NeMoPlatform, + NotFoundError, +) +from nemo_platform.types.inference import ( + ContainerExecutorConfigParam, + ModelDeploymentConfig, + ModelDeploymentConfigFilterParam, + ModelDeploymentConfigModelSpecParam, + ModelDeploymentFilterParam, +) +from nemo_platform.types.models import LoraParam, ModelEntity +from nemo_platform.types.shared_params.tool_call_config import ToolCallConfig as ToolCallConfigParam +from nmp.common.sdk_factory import get_task_sdk +from nmp.customization_common.schemas.model_entity import ( + DeploymentParameters, + ModelEntityCreationError, + ModelEntityTaskConfig, +) +from nmp.customization_common.service.context import NMPJobContext +from nmp.rl.app.constants import SERVICE_NAME +from nmp.rl.entities.values import FinetuningType +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential + +logger = logging.getLogger(__name__) + +# Retry configuration. +MAX_RETRIES = 3 +INITIAL_BACKOFF_SECONDS = 1.0 +MAX_BACKOFF_SECONDS = 30.0 + +ACTIVE_DEPLOYMENT_STATUSES = frozenset({"CREATED", "PENDING", "READY"}) + +SPEC_POLL_INTERVAL_SECONDS = 10 +SPEC_POLL_TIMEOUT_SECONDS = 600 + + +def get_config(config_path: Path) -> ModelEntityTaskConfig: + """Load and validate the model_entity step config from disk.""" + with open(config_path) as f: + return ModelEntityTaskConfig.model_validate(json.load(f)) + + +def sanitize_name(prefix: str, name: str) -> str: + """Build a deployment-safe name from a free-form model name. + + Must match the API's ``{'pattern': '^[a-z](?!.*--)[a-z0-9\\-@.+_]{1,62}(? ModelEntity: + """Poll until the model_spec task has populated the model's spec. + + The spec must be populated before creating a deployment because the + inference service relies on ``spec.family`` and ``spec.base_num_parameters`` + to select the correct NIM profile. + + Raises: + ModelEntityCreationError: If the spec is not populated within the timeout. + """ + logger.info(f"Waiting for model_spec to populate spec on {workspace}/{name}") + start = time.monotonic() + + while time.monotonic() - start < SPEC_POLL_TIMEOUT_SECONDS: + try: + target = self.sdk.models.retrieve(name=name, workspace=workspace) + spec = target.spec + # Deployment/NIM profile selection needs both spec.family and + # spec.base_num_parameters; a partially-populated spec isn't enough. + if spec and getattr(spec, "family", None) and getattr(spec, "base_num_parameters", None) is not None: + logger.info(f"Spec populated on {workspace}/{name}") + return target + except (APIConnectionError, APITimeoutError, InternalServerError) as e: + logger.warning(f"Transient error polling spec for {workspace}/{name}: {e}") + time.sleep(SPEC_POLL_INTERVAL_SECONDS) + + raise ModelEntityCreationError( + f"Timed out waiting for model spec on {workspace}/{name} " + f"after {SPEC_POLL_TIMEOUT_SECONDS}s. The platform could not auto-detect the " + f"model's specifications. Verify the model checkpoint is valid and in a supported format." + ) + + def get_model_entity(self, model_entity: str, fileset_workspace: str) -> ModelEntity: + """Resolve ``"workspace/name"`` (or bare ``"name"``) to a ``ModelEntity``.""" + parts = model_entity.split("/") + if len(parts) == 1 and parts[0]: + me_workspace, me_name = fileset_workspace, parts[0] + elif len(parts) == 2 and all(parts): + me_workspace, me_name = parts[0], parts[1] + else: + # Reject anything that isn't exactly 'name' or 'workspace/name' (e.g. + # 'a/b/c', '/b', 'a/') instead of silently dropping extra segments. + raise ModelEntityCreationError( + f"Invalid model entity reference '{model_entity}': expected 'name' or 'workspace/name'." + ) + + try: + me: ModelEntity = self.sdk.models.retrieve(name=me_name, workspace=me_workspace) + except NotFoundError as e: + raise ModelEntityCreationError(f"Model entity {me_workspace}/{me_name} not found") from e + + return me + + @retry( + stop=stop_after_attempt(MAX_RETRIES), + wait=wait_exponential(multiplier=2, min=INITIAL_BACKOFF_SECONDS, max=MAX_BACKOFF_SECONDS), + retry=retry_if_exception_type((InternalServerError, APITimeoutError, APIConnectionError)), + reraise=True, + ) + def create_model_entity(self, config: ModelEntityTaskConfig) -> tuple[dict, ModelEntity]: + """Create a model entity in the Models service. + + Returns: + Tuple of (result dict, deploy target). For LoRA the deploy target is the + *base* model entity; for SFT it is the newly created output model entity. + + Raises: + ModelEntityCreationError: If creation fails. + """ + # The output entity is created in the workspace declared on the config + # (the "workspace of the model entity to create" contract), not the + # ambient job workspace — the two can differ for cross-workspace jobs. + output_workspace = config.workspace + logger.info(f"Creating model entity: {output_workspace}/{config.name}") + + fileset_workspace = config.fileset.workspace or self.job_ctx.workspace + fileset_ref = f"{fileset_workspace}/{config.fileset.name}" + + logger.info(f"Validating fileset exists: {fileset_workspace}/{config.fileset.name}") + try: + self.sdk.files.filesets.retrieve(workspace=fileset_workspace, name=config.fileset.name) + logger.info(f"Fileset validation successful: {fileset_workspace}/{config.fileset.name}") + except (InternalServerError, APITimeoutError, APIConnectionError): + # Transient API failures must propagate so the @retry wrapping + # create_model_entity can retry them, instead of being masked as a + # permanent (non-retryable) ModelEntityCreationError. + raise + except Exception as e: + logger.error(f"Fileset validation failed: {fileset_workspace}/{config.fileset.name}") + raise ModelEntityCreationError( + f"Cannot create model entity: fileset '{fileset_workspace}/{config.fileset.name}' " + "does not exist or is not accessible" + ) from e + + base_me: ModelEntity = self.get_model_entity(config.model_entity, fileset_workspace) + + if config.peft is not None and config.peft.type == FinetuningType.LORA: + return self._create_or_update_adapter(config, base_me, fileset_ref) + return self._create_or_update_full_entity(config, fileset_ref, output_workspace) + + def _create_or_update_adapter( + self, + config: ModelEntityTaskConfig, + base_me: ModelEntity, + fileset_ref: str, + ) -> tuple[dict, ModelEntity]: + """Create or update a LoRA adapter on ``base_me``. Returns (result, base_me).""" + assert config.peft is not None # type narrowing — caller already checked + try: + output_me = self.sdk.models.adapters.create( + model_name=base_me.name, + workspace=base_me.workspace, + name=config.name, + description=config.description, + fileset=fileset_ref, + finetuning_type=config.peft.type.value, + lora_config=LoraParam( + alpha=config.peft.alpha, + rank=config.peft.rank, + ), + enabled=True, + ) + return output_me.model_dump(), base_me + except ConflictError: + logger.warning( + f"Adapter {base_me.workspace}/{config.name} already exists for model " + f"{base_me.workspace}/{base_me.name}, updating with new fileset" + ) + try: + output_me = self.sdk.models.adapters.update( + adapter=config.name, + model_name=base_me.name, + workspace=base_me.workspace, + fileset=fileset_ref, + description=config.description, + enabled=True, + ) + logger.info( + f"Successfully updated adapter: {base_me.workspace}/{config.name} " + f"for base model {base_me.workspace}/{base_me.name}" + ) + return output_me.model_dump(), base_me + except (InternalServerError, APITimeoutError, APIConnectionError): + raise + except Exception as update_error: + logger.exception( + f"Failed to update existing adapter, {base_me.workspace}/{config.name}: {update_error}" + ) + raise ModelEntityCreationError( + f"Adapter '{config.name}' already exists but update failed: {update_error}" + ) from update_error + except Exception as e: + logger.exception(f"Failed to create model adapter: {e}") + raise ModelEntityCreationError(f"Failed to create model adapter: {e}") from e + + def _create_or_update_full_entity( + self, + config: ModelEntityTaskConfig, + fileset_ref: str, + workspace: str, + ) -> tuple[dict, ModelEntity]: + """Create or update a full / merged model entity. Returns (result, output_me).""" + ft_type = config.peft.type.value if config.peft else FinetuningType.ALL_WEIGHTS.value + + request_body: dict = { + "name": config.name, + "description": config.description, + "fileset": fileset_ref, + "finetuning_type": ft_type, + # Honor the task config's flag (resolved by the compiler from the base + # model entity) rather than re-reading it off a freshly fetched entity. + "trust_remote_code": config.trust_remote_code, + } + if config.base_model: + request_body["base_model"] = config.base_model + + try: + output_me = self.sdk.models.create(workspace=workspace, **request_body) + logger.info(f"Successfully created model entity: {output_me.workspace}/{output_me.name}") + return output_me.model_dump(), output_me + except ConflictError: + logger.warning(f"Model entity already exists: {workspace}/{config.name}, updating existing model") + try: + update_body = {k: v for k, v in request_body.items() if k != "name"} + output_me = self.sdk.models.update( + name=config.name, + workspace=workspace, + **update_body, + ) + logger.info(f"Successfully updated model entity: {output_me.workspace}/{output_me.name}") + return output_me.model_dump(), output_me + except (InternalServerError, APITimeoutError, APIConnectionError): + raise + except Exception as update_error: + logger.exception(f"Failed to update existing model entity: {update_error}") + raise ModelEntityCreationError( + f"Model entity '{config.name}' already exists and update failed: {update_error}" + ) from update_error + except Exception as e: + logger.exception(f"Failed to create model entity: {e}") + raise ModelEntityCreationError(f"Failed to create model entity: {e}") from e + + def launch_model(self, config: ModelEntityTaskConfig, me: ModelEntity) -> None: + """Deploy a model entity after creation. + + For LoRA jobs, ``me`` should be the base model entity. + For SFT jobs, ``me`` should be the output model entity. + """ + dc = config.deployment_config + if dc is None: + return + + # LORA_MERGED produces a full-weight model, so it's deployed like SFT and + # is intentionally excluded from the LoRA-only checks below. + is_lora = config.peft is not None and config.peft.type == FinetuningType.LORA + if is_lora and self._has_active_deployment(me): + return + + if is_lora and isinstance(dc, DeploymentParameters) and not dc.lora_enabled: + logger.warning(f"Deployment requested but lora_enabled is false for a LoRA job: {dc}") + return + + if isinstance(dc, str): + logger.info(f"Resolving deployment config reference: {dc}") + deployment_config = self._resolve_config_ref(dc, me.workspace) + logger.info(f"Using deployment config: {deployment_config.workspace}/{deployment_config.name}") + else: + deployment_config = self._create_deployment_config(dc, me) + + self._create_deployment(deployment_config, me) + + def _has_active_deployment(self, me: ModelEntity) -> bool: + """Check if the model entity already has an active deployment.""" + deployment_configs = self.sdk.inference.deployment_configs.list( + workspace=me.workspace, + filter=ModelDeploymentConfigFilterParam(model_entity_id=f"{me.workspace}/{me.name}"), + ).data + + for c in deployment_configs: + deployments = self.sdk.inference.deployments.list( + filter=ModelDeploymentFilterParam(config=c.name, workspace=me.workspace) + ).data + for d in deployments: + if d.status in ACTIVE_DEPLOYMENT_STATUSES: + logger.info(f"Active deployment (status={d.status}) exists for config {c.name}, skipping") + return True + + return False + + def _resolve_config_ref(self, config_ref: str, me_workspace: str) -> ModelDeploymentConfig: + """Resolve a ``name`` or ``workspace/name`` reference to a ``ModelDeploymentConfig``.""" + parts = config_ref.split("/") + if len(parts) == 2: + workspace, name = parts[0], parts[1] + elif len(parts) == 1: + workspace, name = me_workspace, parts[0] + else: + raise ModelEntityCreationError( + f"Invalid deployment config reference '{config_ref}': expected 'name' or 'workspace/name'" + ) + + try: + return self.sdk.inference.deployment_configs.retrieve(workspace=workspace, name=name) + except Exception as e: + raise ModelEntityCreationError( + f"Failed to resolve deployment config '{config_ref}' in workspace '{workspace}': {e}" + ) from e + + def _create_deployment_config(self, deploy_params: DeploymentParameters, me: ModelEntity) -> ModelDeploymentConfig: + """Create (or update) a ``ModelDeploymentConfig`` from inline parameters.""" + model_spec = ModelDeploymentConfigModelSpecParam( + model_name=me.name, + model_namespace=me.workspace, + lora_enabled=deploy_params.lora_enabled, + ) + executor_config = ContainerExecutorConfigParam( + image_name=deploy_params.image_name, + image_tag=deploy_params.image_tag, + gpu=deploy_params.gpu, + additional_envs=deploy_params.additional_envs, + ) + + if deploy_params.tool_call_config: + model_spec["tool_call_config"] = ToolCallConfigParam( + **deploy_params.tool_call_config.model_dump(exclude_none=True) + ) + + deployment_cfg_name = sanitize_name("sft-cfg", me.name) + try: + return self.sdk.inference.deployment_configs.create( + workspace=me.workspace, + name=deployment_cfg_name, + engine="nim", + model_spec=model_spec, + executor_config=executor_config, + ) + except ConflictError: + logger.info(f"Deployment config {me.workspace}/{deployment_cfg_name} already exists, updating") + return self.sdk.inference.deployment_configs.update( + workspace=me.workspace, + name=deployment_cfg_name, + engine="nim", + model_spec=model_spec, + executor_config=executor_config, + ) + + def _create_deployment(self, deployment_config: ModelDeploymentConfig, me: ModelEntity) -> None: + """Create a deployment from the given ``ModelDeploymentConfig``.""" + # Log identifiers only: the full ModelDeploymentConfig embeds + # executor_config.additional_envs (deployment secrets), which would + # otherwise become durable in the job logs. + logger.info(f"Using deployment config: {deployment_config.workspace}/{deployment_config.name}") + + if not me.spec: + _ = self._wait_for_spec(me.workspace, me.name) + + deployment_name = sanitize_name("sft-deploy", me.name) + try: + deployment = self.sdk.inference.deployments.create( + workspace=deployment_config.workspace, + name=deployment_name, + config=deployment_config.name, + ) + logger.info(f"Deployment created: {deployment.workspace}/{deployment.name}") + except ConflictError: + logger.info(f"Deployment {deployment_config.workspace}/{deployment_name} already exists") + deployment = self.sdk.inference.deployments.retrieve( + workspace=deployment_config.workspace, + name=deployment_name, + ) + + deployment_status = self.sdk.inference.deployments.retrieve( + workspace=deployment.workspace, + name=deployment.name, + ) + logger.info( + f"Deployment {deployment_status.workspace}/{deployment_status.name} status: {deployment_status.status}" + ) + + +def run(sdk: NeMoPlatform | None = None, job_ctx: NMPJobContext | None = None) -> int: + """Execute the model entity creation task. + + Args: + sdk: Optional SDK instance for dependency injection (for testing). + If None, creates one via get_task_sdk(). + job_ctx: Optional job context for dependency injection (for testing). + If None, creates one via NMPJobContext.from_env(). + + Returns: + Exit code (0 for success, non-zero for failure). + """ + job_ctx = job_ctx or NMPJobContext.from_env() + + sdk_owned = sdk is None + try: + sdk = sdk or get_task_sdk(SERVICE_NAME).with_options(workspace=job_ctx.workspace) + runner = ModelEntityRunner(sdk=sdk, job_ctx=job_ctx) + + config = get_config(job_ctx.config_path) + + # Log only a non-sensitive summary. The full job context carries service + # URLs/identifiers and the config's deployment_config.additional_envs can + # carry deployment secrets, so neither is dumped wholesale. + logger.info( + "Starting model entity task: job_id=%s, name=%s, workspace=%s, fileset=%s/%s, deployment_configured=%s", + job_ctx.job_id, + config.name, + config.workspace, + config.fileset.workspace or job_ctx.workspace, + config.fileset.name, + config.deployment_config is not None, + ) + logger.info(f"NeMo Platform service URL: {sdk.base_url}") + + result, deploy_target = runner.create_model_entity(config) + logger.info(f"Model entity creation complete: {result}") + + runner.launch_model(config, deploy_target) + return 0 + + except ModelEntityCreationError as e: + logger.exception(f"Model entity creation failed: {e}") + return 1 + except Exception as e: + logger.exception(f"Model entity task failed: {e}") + return 1 + finally: + if sdk_owned and sdk is not None: + sdk.close() diff --git a/services/rl/src/nmp/rl/tasks/training/__main__.py b/services/rl/src/nmp/rl/tasks/training/__main__.py new file mode 100644 index 0000000000..315e60b9c6 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/__main__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Training task entry point. + +Usage: + python -m nmp.rl.tasks.training + +The runner reads the platform Jobs step config (``NEMO_JOB_STEP_CONFIG_FILE_PATH``), +builds a :class:`~nmp.rl.app.jobs.training.schemas.TrainingStepConfig`, and runs +the :class:`~nmp.rl.tasks.training.runner.TrainingRunner`. + +In distributed (multi-node) training, all pods run this entry point. +The DistributedContext handles role detection and coordination: +- Rank 0 (coordinator): Runs all phases, reports progress +- Rank > 0 (workers): Participate in training, wait at barriers +""" + +import logging +import sys + +from .runner import TrainingRunner + +logger = logging.getLogger(__name__) + + +def main() -> int: + """Execute the training task.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + try: + with TrainingRunner() as runner: + result = runner.run() + return 0 if result.success else 1 + except Exception as e: + logger.exception(f"Training task failed: {e}") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py new file mode 100644 index 0000000000..fbb75a5b07 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/backend.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""TrainingBackend protocol implementation for NeMo RL (DPO).""" + +import logging +import os +import signal +from pathlib import Path +from typing import Any, Optional, cast + +from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager +from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.training.progress import JobsServiceProgressReporter +from nmp.rl.app.jobs.training.schemas import ( + CheckpointFormat, + CheckpointInfo, + TrainingMetrics, + TrainingStepConfig, + TrainingType, +) +from nmp.rl.app.jobs.training.schemas import ( + TrainingBackend as TrainingBackendEnum, +) +from nmp.rl.tasks.training.backends.nemo_rl.checkpoints import convert_dcp_to_huggingface +from nmp.rl.tasks.training.chat_templates import apply_chat_template_to_checkpoint +from nmp.rl.tasks.training.errors.parser import parse_error_from_output +from nmp.rl.tasks.training.protocol import LibraryConfig, TrainingBackend + +from .dpo_config import compile_dpo_config +from .ray_bootstrap import create_bootstrap_from_env + +logger = logging.getLogger(__name__) + +# Path to driver scripts (relative to this module) +_DRIVER_DIR = Path(__file__).parent + + +class NemoRLBackend(TrainingBackend): + """TrainingBackend implementation for NeMo RL (DPO). + + This backend handles DPO (Direct Preference Optimization) training using NeMo RL. + + Key responsibilities: + - Run pre-training conversions (model to HF format) via injected converter + - Compile TrainingStepConfig to NeMo RL YAML format + - Bootstrap Ray cluster on Volcano-provisioned pods + - Execute appropriate training driver (DPO) + - Process checkpoints to standard output format + + Args: + job_ctx: Job context with job metadata + """ + + def __init__( + self, + job_ctx: NMPJobContext, + ) -> None: + """Initialize the backend. + + Args: + job_ctx: Job context with job metadata + """ + self._job_ctx = job_ctx + + @property + def backend_type(self) -> TrainingBackendEnum: + return TrainingBackendEnum.NEMO_RL + + def compile_config( + self, + customizer_config: TrainingStepConfig, + workspace_dir: Path, + ) -> dict[str, Any]: + """Compile TrainingStepConfig to NeMo RL YAML format. + + Args: + customizer_config: The training step configuration + workspace_dir: Directory for storing generated config files + + Returns: + Configuration dict for NeMo RL (will be serialized to YAML) + """ + training_type = customizer_config.training.training_type + + if training_type == TrainingType.DPO: + return compile_dpo_config(customizer_config, self._job_ctx) + + # GRPO is reserved headroom in the schema but not yet implemented. Reject + # it here — the earliest training-type-specific wiring point — so the job + # fails fast with a clear message instead of routing to an unfinished stub + # and crashing deep inside the training container. + raise NotImplementedError( + f"NemoRLBackend does not yet support training type {training_type.value!r}. " + f"Only {TrainingType.DPO.value!r} is currently available." + ) + + def execute_training( + self, + customizer_config: TrainingStepConfig, + library_config: LibraryConfig, + progress: JobsServiceProgressReporter, + ) -> TrainingMetrics: + """Execute NeMo RL training via Ray bootstrap. + + Args: + customizer_config: The training step configuration + library_config: The compiled library configuration + progress: Progress reporter for status updates + + Returns: + TrainingMetrics with results from the training run + """ + progress.report_running("training", backend=self.backend_type.value) + + # Get the workspace directory from config path + workspace_dir = library_config.config_path.parent + + # Environment overrides the driver subprocess inherits. We snapshot and + # restore them (below) so a reused worker process doesn't leak this run's + # values — e.g. a stale MLFLOW_URI — into a subsequent run. + env_overrides = { + "BASE_LOG_DIR": str(workspace_dir), + "GPUS_PER_NODE": str(customizer_config.parallelism.num_gpus_per_node), + } + # MLflow integration (if configured) + if customizer_config.integrations and customizer_config.integrations.mlflow: + mlflow_config = customizer_config.integrations.mlflow + if mlflow_config.tracking_uri: + env_overrides["MLFLOW_URI"] = mlflow_config.tracking_uri + + # Build driver arguments + driver_path = self._get_driver_path(customizer_config) + driver_args = [ + "--config", + str(library_config.config_path), + "--id", + self._job_ctx.job_id, + "--output-model", + customizer_config.model.name or "output_model", + ] + + # Bootstrap Ray cluster and run driver + logger.info(f"Starting Ray cluster and running driver: {driver_path}") + logger.info(f"Driver args: {driver_args}") + + bootstrap = create_bootstrap_from_env() + + # Set up signal handler for cleanup — terminate the driver subprocess + # explicitly so it doesn't become orphaned, then let SystemExit propagate + # to trigger Ray cluster cleanup in the bootstrap's finally block. + def cleanup(signum, frame): + logger.warning(f"Signal {signum} received, terminating driver and cleaning up") + bootstrap.terminate_driver(signum) + raise SystemExit(signum) + + # Snapshot env vars and signal handlers so we can restore them after the + # run; otherwise stale state persists on workers reused across runs (a + # leftover handler closure could even terminate a later, unrelated driver). + saved_env = {key: os.environ.get(key) for key in env_overrides} + previous_sigint = signal.getsignal(signal.SIGINT) + previous_sigterm = signal.getsignal(signal.SIGTERM) + try: + os.environ.update(env_overrides) + signal.signal(signal.SIGINT, cleanup) + signal.signal(signal.SIGTERM, cleanup) + + exit_code = bootstrap.run_with_driver(str(driver_path), driver_args) + finally: + signal.signal(signal.SIGINT, previous_sigint) + signal.signal(signal.SIGTERM, previous_sigterm) + for key, original in saved_env.items(): + if original is None: + os.environ.pop(key, None) + else: + os.environ[key] = original + + if exit_code != 0: + parsed = parse_error_from_output(bootstrap.driver_output, exit_code) + raise parsed.to_exception() + + logger.info("Training completed successfully") + + # Return empty metrics (actual metrics are logged during training) + return TrainingMetrics(total_steps=0, total_epochs=0) + + def _get_driver_path(self, config: TrainingStepConfig) -> Path: + """Get the appropriate driver script path based on training type. + + Args: + config: Training configuration with type information + + Returns: + Path to the driver script (dpo_driver.py or grpo_driver.py) + """ + training_type = config.training.training_type + + if training_type == TrainingType.DPO: + return _DRIVER_DIR / "dpo_driver.py" + + raise NotImplementedError( + f"No training driver available for training type {training_type.value!r}; " + f"only {TrainingType.DPO.value!r} is currently supported." + ) + + def find_best_checkpoint( + self, + workspace_dir: Path, + customizer_config: TrainingStepConfig, + library_config: Optional[LibraryConfig] = None, + ) -> Path: + """Find the best checkpoint after training. + + NeMo RL driver converts the best checkpoint to HuggingFace format + and saves it to {workspace_dir}/output. This method returns that path. + + Args: + workspace_dir: Directory containing training artifacts + customizer_config: Training configuration + + Returns: + Path to the converted HF checkpoint + """ + if library_config is None: + raise ValueError("Library config is required to find the best checkpoint") + + checkpointing_config = library_config.config_dict["checkpointing"] + if checkpointing_config is None or not isinstance(checkpointing_config, dict): + raise ValueError("Checkpointing config is required to find the best checkpoint") + + checkpointing_config = cast(CheckpointingConfig, checkpointing_config) + checkpointer = CheckpointManager(checkpointing_config) + + # get_best_checkpoint_path() handles the missing-metric case internally: it + # filters out checkpoints lacking the metric (with a warning) and, if none + # have it, returns the latest checkpoint. It only returns None when there + # are no checkpoints at all + best_checkpoint = checkpointer.get_best_checkpoint_path() + + if best_checkpoint is None: + raise ValueError("No best checkpoint found") + + best_checkpoint_path = Path(best_checkpoint) + if not best_checkpoint_path.exists(): + raise ValueError(f"Best checkpoint not found at {best_checkpoint_path}") + + return best_checkpoint_path + + def process_checkpoint( + self, + checkpoint_path: Path, + output_path: Path, + customizer_config: TrainingStepConfig, + library_config: LibraryConfig | None = None, + ) -> CheckpointInfo: + """Process NeMo RL checkpoint to standard output format. + + The NeMo RL driver already converts checkpoints to HuggingFace format. + This method copies the output and applies the chat template. + + Args: + checkpoint_path: Path to the checkpoint directory in the DCP format + output_path: Where to write the processed checkpoint in the HF format + customizer_config: Training configuration + library_config: Library-specific config (contains chat template) + + Returns: + CheckpointInfo with output path, format, and precision + """ + logger.info("Processing created checkpoint") + hf_checkpoint_path = convert_dcp_to_huggingface(checkpoint_path, output_path) + + # Apply chat template if available + chat_template = None + if library_config and library_config.config_dict: + chat_template = library_config.config_dict.get("policy", {}).get("tokenizer", {}).get("chat_template") + + if chat_template: + apply_chat_template_to_checkpoint(hf_checkpoint_path, chat_template) + logger.debug("Applied chat template to checkpoint") + + return CheckpointInfo( + path=str(hf_checkpoint_path), + format=CheckpointFormat.HF, + precision=customizer_config.model.precision, + ) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py new file mode 100644 index 0000000000..8ac0760844 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import logging +from typing import Any + +from nmp.customization_common.training.progress import JobsServiceProgressReporter + +logger = logging.getLogger(__name__) + + +class TrainingProgressCallback: + """ + Callback for reporting NeMo RL training progress to the Jobs service. + + This class composes JobsServiceProgressReporter and provides training-specific + methods for reporting detailed metrics during training. + """ + + def __init__(self, reporter: JobsServiceProgressReporter): + self._reporter = reporter + + def report_training_start(self, max_steps: int, num_epochs: int) -> None: + """Report that training has started with schedule information.""" + self._reporter.configure_progress_tracking(max_steps, num_epochs) + self._reporter.report_running(phase="training", step=0, max_steps=max_steps, num_epochs=num_epochs) + + def report_train_step( + self, + step: int, + epoch: int, + loss: float, + lr: float | None = None, + grad_norm: float | None = None, + **additional_metrics: Any, + ) -> None: + """Report training step with metrics. + + Args: + step: Training step number + epoch: Current epoch number + loss: Training loss value + lr: Learning rate (optional) + grad_norm: Gradient norm (optional) + **additional_metrics: Additional training metrics to report (e.g., num_valid_samples, + preference_loss, rewards_rejected_mean, global_valid_seqs, global_valid_toks) + """ + self._reporter.report_running( + phase="training", + step=step, + epoch=epoch, + train_loss=loss, + lr=lr, + grad_norm=grad_norm, + **additional_metrics, + ) + + def report_validation( + self, + step: int, + epoch: int, + val_loss: float, + **additional_metrics: Any, + ) -> None: + """Report validation results. + + Args: + step: Training step number + epoch: Current epoch number + val_loss: Validation loss value + **additional_metrics: Additional validation metrics to report (e.g., accuracy, + num_valid_samples, or any other validation-specific metrics) + """ + self._reporter.report_running( + phase="validation", + step=step, + epoch=epoch, + val_loss=val_loss, + **additional_metrics, + ) + + def report_checkpoint_saved(self, step: int, epoch: int, checkpoint_path: str | None = None) -> None: + """Report that a checkpoint was saved.""" + self._reporter.report_running(phase="checkpoint_saved", step=step, epoch=epoch, checkpoint_path=checkpoint_path) + + def close(self) -> None: + """Clean up resources.""" + self._reporter.close() diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py new file mode 100644 index 0000000000..6352a0f231 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/checkpoints.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""DCP → HuggingFace checkpoint conversion utilities. + +This module handles conversion of Distributed Checkpoint (DCP) format +used by PyTorch/NeMo to HuggingFace format for model serving and distribution. +""" + +import glob +import logging +import os +from pathlib import Path + +import yaml +from nemo_rl.utils.native_checkpoint import convert_dcp_to_hf +from transformers import AutoModelForCausalLM + +logger = logging.getLogger(__name__) + + +def convert_dcp_to_huggingface( + dcp_checkpoint_path: Path, + output_path: Path, +) -> Path: + """Convert a DCP checkpoint to HuggingFace format. + + Args: + dcp_checkpoint_path: Path to the DCP checkpoint directory + output_path: Path for the output HuggingFace checkpoint + model_config: Optional model configuration overrides + + Returns: + Path to the converted HuggingFace checkpoint + """ + with open(dcp_checkpoint_path / "config.yaml", "r") as f: + config = yaml.safe_load(f) + + model_name_or_path = config["policy"]["model_name"] + tokenizer_name_or_path = f"{dcp_checkpoint_path}/policy/tokenizer" + + # It saves the weights as a single pytorch_model.bin file (pickle-based PyTorch format). + hf_ckpt = convert_dcp_to_hf( + dcp_ckpt_path=f"{dcp_checkpoint_path}/policy/weights", + hf_ckpt_path=str(output_path), + model_name_or_path=model_name_or_path, + tokenizer_name_or_path=tokenizer_name_or_path, + overwrite=True, + ) + + saved_hf_checkpoint_path = Path(hf_ckpt) + if not saved_hf_checkpoint_path.exists(): + raise FileNotFoundError( + f"HF checkpoint not found at {saved_hf_checkpoint_path} after conversion from DCP to HF" + ) + # Compare resolved paths: convert_dcp_to_hf() may return an absolute path while + # output_path is relative, and string inequality would then falsely trip even + # when both point at the same directory. + if output_path.resolve() != saved_hf_checkpoint_path.resolve(): + raise ValueError( + f"Output path {output_path} does not match the saved HF checkpoint path {saved_hf_checkpoint_path}" + ) + + # Convert pickle-based .bin format to safetensors format + # Shards the model into multiple files if larger than 4GB + model = AutoModelForCausalLM.from_pretrained(saved_hf_checkpoint_path) + model.save_pretrained( + saved_hf_checkpoint_path, + safe_serialization=True, + max_shard_size="4GB", + ) + + # Remove unnecessary files from DCP checkpoint + # *.bin files come from the DCP format, which is not needed in the HF safetensors format + for f in glob.glob(os.path.join(saved_hf_checkpoint_path, "*.bin")) + glob.glob( + os.path.join(saved_hf_checkpoint_path, "*.bin.index.json") + ): + os.remove(f) + + logger.info("Saved HF checkpoint successfully") + + return saved_hf_checkpoint_path diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py new file mode 100644 index 0000000000..08b2cfe836 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py @@ -0,0 +1,588 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""TrainingStepConfig -> NeMo-RL DPO YAML generation. + +Converts the internal :class:`TrainingStepConfig` into the complete YAML config +NeMo-RL's DPO training expects. + +The full config is generated here in one place — there is no external base file +to merge against. Fields driven by the job spec (model, batch sizes, parallelism, +schedule, DPO hyperparameters, optimizer/scheduler, integrations) are computed +from ``TrainingStepConfig``; every other key NeMo-RL's schema requires +(``policy.megatron_cfg``, ``dtensor_cfg.lora_cfg``, ``fp8_cfg``, +``dpo.val_at_end``, ``checkpointing.save_optimizer``, the logger subsections, …) +is set explicitly to a known-good default. The Megatron backend block is inert +(``enabled: False``) since training runs on DTensor, but must still be fully +populated to satisfy the schema. +""" + +import logging +from pathlib import Path +from typing import Any + +from nmp.customization_common.service.context import NMPJobContext +from nmp.rl.app.jobs.training.schemas import ( + DPOConfig, + OptimizerType, + TrainingStepConfig, +) +from nmp.rl.tasks.training.chat_templates import resolve_chat_template +from nmp.rl.tasks.training.datasets.preparation import ( + PreparedDataset, + compute_val_check_interval, + prepare_dataset, +) +from nmp.rl.tasks.training.datasets.validation import DatasetValidator, detect_dpo_schema_name +from nmp.rl.tasks.training.integrations import ( + build_mlflow_config, + build_wandb_config, +) + +logger = logging.getLogger(__name__) + + +def compile_dpo_config( + customizer_config: TrainingStepConfig, + job_ctx: NMPJobContext, +) -> dict[str, Any]: + """ + Compile TrainingStepConfig to NeMo RL DPO configuration dict. + + This transforms the standardized TrainingStepConfig into the format + expected by NeMo RL's DPO training. The output dict will be serialized + to YAML by the training runner. + + Args: + customizer_config: The training step configuration + job_ctx: Job context + + Returns: + Configuration dict for NeMo RL DPO training + + Reference: https://github.com/NVIDIA-NeMo/RL/blob/main/examples/configs/dpo.yaml + """ + cfg: dict[str, Any] = {} + workspace_dir = Path(customizer_config.workspace_path) + + # === Dataset Preparation === + prepared = prepare_dataset( + dataset_path=Path(customizer_config.dataset.path), + output_dir=workspace_dir / "dataset", + ) + logger.info( + f"Prepared dataset: train={prepared.train_samples} samples, validation={prepared.validation_samples} samples" + ) + validator = DatasetValidator(training_type=customizer_config.training.training_type) + validator.validate_dataset(str(prepared.train_file)) + validator.validate_dataset(str(prepared.validation_file)) + logger.info("Validated datasets successfully") + + # === Training Schedule Calculations === + batch_size = customizer_config.batch.global_batch_size + micro_batch_size = customizer_config.batch.micro_batch_size + epochs = customizer_config.schedule.epochs + + # Compute steps per epoch (round up to ensure all samples are used) + steps_per_epoch = max((prepared.train_samples + batch_size - 1) // batch_size, 1) + total_steps = steps_per_epoch * epochs + + # Determine effective max_steps + user_max_steps = customizer_config.schedule.max_steps + if user_max_steps and user_max_steps > 0: + max_steps = min(user_max_steps, total_steps) + else: + max_steps = total_steps + + # Compute validation interval + val_check_interval = compute_val_check_interval( + steps_per_epoch=steps_per_epoch, + max_steps=max_steps, + val_check_interval=customizer_config.schedule.val_check_interval, + ) + + logger.info( + f"Training schedule: {prepared.train_samples} samples, batch_size={batch_size}, " + f"steps_per_epoch={steps_per_epoch}, epochs={epochs}, max_steps={max_steps}, " + f"val_period={val_check_interval}" + ) + + # === Get DPO Hyperparameters === + dpo_hp = customizer_config.training.dpo or DPOConfig() + + # Checkpoint selection ranks by the validation metric (`metric_name` below), + # so the saved checkpoint must carry validation metrics. NeMo-RL always saves a + # checkpoint on the LAST step (is_last_step), and the last step is frequently + # not a validation step — e.g. with the small-run default (max_steps caps the + # run at 7 steps with 200 rows / batch 32) the last step never aligns with + # val_period, so that checkpoint has no `val:...` metric. NeMo-RL then warns and + # falls back to "latest" instead of best. + # + # Two measures keep validation aligned with checkpointing: + # 1. `val_at_end` (the dpo section sets it from schedule.val_at_end, which + # defaults to True) forces a validation pass on the FINAL step, so the + # is_last_step checkpoint carries the metric regardless of whether + # max_steps is a multiple of val_period. This is the primary guarantee. + # A user CAN opt out (val_at_end=False, to skip the final eval); in that + # case the last-step checkpoint may lack the metric and best-checkpoint + # selection degrades to "latest" — the accepted trade-off of opting out. + # 2. `save_period == val_period` so intermediate saves also land on validation + # steps. (val_period=steps_per_epoch when val_check_interval is unset.) + # NeMo-RL's get_best_checkpoint_path() itself is the safety net: it filters out + # checkpoints missing the metric and, if none have it, returns the latest — so + # the val_at_end=False path degrades gracefully to "latest" instead of crashing. + val_period = val_check_interval + + # === DPO Section === + cfg["dpo"] = { + "max_num_epochs": epochs, + "max_num_steps": max_steps, + "steps_per_epoch": steps_per_epoch, + "val_period": val_period, + "val_batches": 0, # Run the entire validation dataset + "val_global_batch_size": batch_size, + "val_micro_batch_size": micro_batch_size, + "val_at_start": True, + "val_at_end": customizer_config.schedule.val_at_end, + "seed": customizer_config.seed, + # DPO-specific hyperparameters + "reference_policy_kl_penalty": dpo_hp.ref_policy_kl_penalty, + "preference_average_log_probs": dpo_hp.preference_average_log_probs, + "sft_average_log_probs": dpo_hp.sft_average_log_probs, + "preference_loss_weight": dpo_hp.preference_loss_weight, + "sft_loss_weight": dpo_hp.sft_loss_weight, + } + + # === Checkpointing Section === + # save_period == val_period so intermediate saves land on validation steps; the + # always-saved last-step checkpoint gets its validation metric from `val_at_end` + # (see the val_period comment above). `metric_name` + `keep_top_k` then select + # the best checkpoint by validation loss. + cfg["checkpointing"] = { + "enabled": True, + "checkpoint_dir": str(workspace_dir / "checkpoints"), + "metric_name": "val:validation-default_loss", + "higher_is_better": False, + "keep_top_k": customizer_config.schedule.keep_top_k, + "save_period": val_period, + "checkpoint_must_save_by": None, + "save_optimizer": True, + } + + # === Policy Section === + model_path = customizer_config.model.path + precision = _adapt_precision(customizer_config.model.precision) + parallelism = customizer_config.parallelism + + # Resolve chat template with priority: + # 1. Fileset metadata chat_template (from model entity spec) + # 2. Custom template from DEFAULT_CHAT_TEMPLATES (if model.name matches) + # 3. Model's built-in tokenizer template (fallback) + chat_template = resolve_chat_template( + model_path=model_path, + model_name=customizer_config.model.name, + user_template=customizer_config.model.chat_template, + trust_remote_code=customizer_config.model.trust_remote_code, + ) + + cfg["policy"] = { + "model_name": model_path, + "tokenizer": { + "name": model_path, + "chat_template": chat_template, + "chat_template_kwargs": None, + }, + "train_global_batch_size": batch_size, + "train_micro_batch_size": micro_batch_size, + "max_total_sequence_length": customizer_config.model.max_seq_length, + "precision": precision, + "offload_optimizer_for_logprob": False, + # Training runs on the DTensor backend. We propagate tensor / sequence / + # context parallelism from the parallelism config; the remaining keys are + # NeMo-RL defaults (LoRA disabled — DPO is full-weight here). + "dtensor_cfg": { + "env_vars": {"PYTORCH_CUDA_ALLOC_CONF": ""}, + "enabled": True, + "cpu_offload": False, + "sequence_parallel": parallelism.sequence_parallel, + "activation_checkpointing": parallelism.activation_checkpointing, + "tensor_parallel_size": parallelism.tensor_parallel_size, + "context_parallel_size": parallelism.context_parallel_size, + "custom_parallel_plan": None, + "clear_cache_every_n_steps": None, + "automodel_kwargs": {}, + "lora_cfg": { + "enabled": False, + "target_modules": [], + "exclude_modules": [], + "match_all_linear": True, + "dim": 8, + "alpha": 32, + "dropout": 0.0, + "dropout_position": "post", + "lora_A_init": "xavier", + "use_triton": True, + }, + }, + "dynamic_batching": {"enabled": False}, + "sequence_packing": _build_sequence_packing_config(customizer_config), + "make_sequence_length_divisible_by": parallelism.tensor_parallel_size, + "max_grad_norm": dpo_hp.max_grad_norm, + # Optimizer and scheduler + "optimizer": _build_optimizer_config(customizer_config), + # Schedule LR over the steps that will actually execute (max_steps after + # user capping), so warmup + cosine decay complete within the run instead of + # being stretched across the uncapped epoch length and never reaching min_lr. + "scheduler": _build_scheduler_config(customizer_config, max_steps), + # Megatron backend is disabled (we train on DTensor). NeMo-RL's config + # schema still requires this block to be fully populated, so it is + # reproduced inert here; none of these values take effect while + # ``enabled`` is False. + "megatron_cfg": _megatron_cfg_disabled(precision, dpo_hp.max_grad_norm), + } + + # === Data Section === + cfg["data"] = _build_data_config(customizer_config, prepared) + + # === Logger Section === + cfg["logger"] = _build_logger_config(customizer_config, job_ctx, workspace_dir) + + # === Cluster Section === + cfg["cluster"] = { + "gpus_per_node": parallelism.num_gpus_per_node, + "num_nodes": parallelism.num_nodes, + } + + return cfg + + +def _megatron_cfg_disabled(precision: str, max_grad_norm: float) -> dict[str, Any]: + """Return the (inert) Megatron backend config block. + + Training uses the DTensor backend, so ``enabled`` is False and none of these + values take effect. NeMo-RL's config schema still requires the block to be + present and fully populated, so it is reproduced here. ``pipeline_dtype`` and + ``optimizer.clip_grad`` track ``policy.precision`` / ``policy.max_grad_norm``. + """ + return { + "enabled": False, + "use_linear_ce_fusion_loss": False, + "linear_ce_fusion_chunk_size": 256, + "force_reconvert_from_hf": False, + "empty_unused_memory_level": 1, + "activation_checkpointing": False, + "tensor_model_parallel_size": 2, + "expert_tensor_parallel_size": 1, + "expert_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "pipeline_dtype": precision, + "num_layers_in_first_pipeline_stage": None, + "num_layers_in_last_pipeline_stage": None, + "sequence_parallel": True, + "freeze_moe_router": False, + "moe_router_dtype": "fp64", + "moe_router_load_balancing_type": "aux_loss", + "moe_router_bias_update_rate": 1e-3, + "moe_permute_fusion": False, + "apply_rope_fusion": True, + "bias_activation_fusion": True, + "defer_fp32_logits": False, + "moe_per_layer_logging": False, + "moe_enable_deepep": False, + "moe_token_dispatcher_type": "alltoall", + "moe_shared_expert_overlap": False, + "gradient_accumulation_fusion": False, + "peft": { + "enabled": False, + "target_modules": [], + "exclude_modules": [], + "dim": 8, + "alpha": 32, + "dropout": 0.0, + "dropout_position": "post", + "lora_A_init_method": "xavier", + "lora_B_init_method": "zero", + "a2a_experimental": False, + "lora_dtype": None, + }, + "optimizer": { + "optimizer": "adam", + "lr": 5.0e-6, + "min_lr": 5.0e-6, + "weight_decay": 0.1, + "bf16": True, + "fp16": False, + "params_dtype": "float32", + "adam_beta1": 0.9, + "adam_beta2": 0.98, + "adam_eps": 1e-8, + "sgd_momentum": 0.9, + "use_distributed_optimizer": True, + "use_precision_aware_optimizer": True, + "clip_grad": max_grad_norm, + "optimizer_cpu_offload": False, + "optimizer_offload_fraction": 0.0, + }, + "scheduler": { + "start_weight_decay": 0.1, + "end_weight_decay": 0.1, + "weight_decay_incr_style": "constant", + "lr_decay_style": "constant", + "lr_warmup_iters": 1, + "lr_warmup_init": 0.00000001, + }, + "distributed_data_parallel_config": { + "grad_reduce_in_fp32": False, + "overlap_grad_reduce": True, + "overlap_param_gather": True, + "data_parallel_sharding_strategy": "optim_grads_params", + "use_custom_fsdp": False, + }, + "fp8_cfg": { + "enabled": False, + "fp8": "e4m3", + "fp8_recipe": "blockwise", + "fp8_param": False, + }, + } + + +def _build_data_config(customizer_config: TrainingStepConfig, prepared: PreparedDataset) -> dict[str, Any]: + """Build the NeMo-RL ``data`` config. + + NeMo-RL's ``setup_preference_data`` reads nested ``train`` / ``validation`` + dataset specs (``dataset_name`` + ``data_path``) and builds each split by + instantiating the class registered under ``dataset_name`` in NeMo-RL's + ``DATASET_REGISTRY`` as ``cls(**spec)`` — no custom preprocessor. + ``detect_dpo_schema_name`` returns that registry key, one of + ``BinaryPreferenceDataset`` / ``PreferenceDataset`` / ``HelpSteer3`` / + ``Tulu3Preference``. + + All four load from the local ``data_path`` here. ``BinaryPreferenceDataset`` and + ``PreferenceDataset`` accept a local path natively; ``HelpSteer3`` and + ``Tulu3Preference`` only do so because the DPO driver re-points those registry + entries to our local-file-capable subclasses via ``register_preference_datasets()`` + (NeMo-RL's built-ins for those two always download from HuggingFace and ignore + ``data_path``). If that registration is ever removed, a user-uploaded + HelpSteer3/Tulu3 dataset would silently train on the public HF dataset instead. + + The schema is detected per split: train and validation may legitimately be in + different supported formats, so inferring once from the train file and reusing + it for validation would point validation at the wrong loader. + """ + + def _dataset_spec(path: Path) -> dict[str, Any]: + dataset_name = detect_dpo_schema_name(path) + spec: dict[str, Any] = {"dataset_name": dataset_name, "data_path": str(path)} + # BinaryPreferenceDataset reads explicit prompt/chosen/rejected keys; our + # datasets use those exact field names. + if dataset_name == "BinaryPreferenceDataset": + spec.update({"prompt_key": "prompt", "chosen_key": "chosen", "rejected_key": "rejected"}) + return spec + + return { + "max_input_seq_length": customizer_config.model.max_seq_length, + # Disable dataloader shuffling for deterministic, reproducible training + # order (and consistent ordering across distributed ranks). The + # train/validation split is already randomized at preparation time. + "shuffle": False, + "num_workers": 1, + "train": _dataset_spec(prepared.train_file), + "validation": _dataset_spec(prepared.validation_file), + } + + +def _adapt_precision(precision: str | None) -> str: + """ + + Returns in the format that is expected by NeMo FW: + ('transformer-engine', 'transformer-engine-float16', '16-true', '16-mixed', + 'bf16-true', 'bf16-mixed', '32-true', '64-true', 64, 32, 16, '64', '32', '16', 'bf16') + """ + precision_map = { + "bf16": "bfloat16", + "bf16-mixed": "bfloat16", + "fp16": "float16", + "fp32": "float32", + None: "bfloat16", # Default + } + result = precision_map.get(precision) + if result is None: + logger.warning(f"Unknown precision '{precision}', defaulting to bfloat16") + return "bfloat16" + return result + + +def _build_sequence_packing_config(customizer_config: TrainingStepConfig) -> dict[str, Any]: + """Build sequence packing configuration.""" + logger.warning("Sequence packing is currently not supported with DPO.") + return {"enabled": False} + + ## TODO: uncomment below code when sequence packing is supported by nemo-rl + ## Sequence packing is currently not supported with DPO. See https://github.com/NVIDIA-NeMo/RL/issues/719 + # if not customizer_config.batch.sequence_packing: + # return {"enabled": False} + + # return { + # "enabled": True, + # "train_mb_tokens": 2048, + # "logprob_mb_tokens": 2048, + # "algorithm": "modified_first_fit_decreasing", + # "sequence_length_round": 64, # Hardware alignment + # } + + +def _build_optimizer_config(customizer_config: TrainingStepConfig) -> dict[str, Any]: + """Build optimizer configuration for NeMo RL. + + Supports: + - AdamW (with weight decay) + - Adam (without weight decay correction) + + The optimizer type is determined by the optimizer_type field in OptimizerConfig. + """ + opt = customizer_config.optimizer + optimizer_type = opt.optimizer_type or OptimizerType.ADAMW_WITH_COSINE_ANNEALING + + # Determine optimizer name based on type + if optimizer_type in (OptimizerType.ADAM_WITH_COSINE_ANNEALING, OptimizerType.ADAM_WITH_FLAT_LR): + optimizer_name = "torch.optim.Adam" + else: + # Default: AdamW for ADAMW_WITH_COSINE_ANNEALING and ADAMW_WITH_FLAT_LR + optimizer_name = "torch.optim.AdamW" + + return { + "name": optimizer_name, + "kwargs": { + "lr": opt.learning_rate, + "weight_decay": opt.weight_decay, + "betas": [opt.beta1, opt.beta2], + "eps": opt.eps, + "foreach": False, + "fused": False, + }, + } + + +def _build_scheduler_config( + customizer_config: TrainingStepConfig, + num_steps: int, +) -> list[dict[str, Any]] | dict[str, Any]: + """ + Build learning rate scheduler configuration. + + Supports two scheduler types based on optimizer_type: + - Cosine Annealing: LinearLR warmup followed by CosineAnnealingLR decay + - Flat LR: ConstantLR (constant learning rate throughout training) + + ``num_steps`` is the number of steps that will actually run (``max_steps`` after + capping), so the warmup + decay horizon matches the executed run. + """ + opt = customizer_config.optimizer + optimizer_type = opt.optimizer_type or OptimizerType.ADAMW_WITH_COSINE_ANNEALING + warmup_steps = opt.warmup_steps + lr = opt.learning_rate + min_lr = opt.min_learning_rate or 0.0 + + # Check if using flat LR scheduler + if optimizer_type in (OptimizerType.ADAM_WITH_FLAT_LR, OptimizerType.ADAMW_WITH_FLAT_LR): + # Flat LR: Use ConstantLR scheduler + return { + "name": "torch.optim.lr_scheduler.ConstantLR", + "kwargs": { + "factor": 1.0, + "total_iters": num_steps, + }, + } + + if optimizer_type in (OptimizerType.ADAM_WITH_COSINE_ANNEALING, OptimizerType.ADAMW_WITH_COSINE_ANNEALING): + # Default: Cosine Annealing with warmup + # Compute start_factor for warmup (avoid division by zero) + start_factor = max(min_lr / lr, 1e-5) if lr > 0 else 1e-5 + # Clamp warmup_steps to >= 1 for cosine schedulers; LinearLR(total_iters=0) + # and milestones=[0] produce invalid scheduler behavior + effective_warmup_steps = max(warmup_steps or 0, 1) + + return [ + { + "name": "torch.optim.lr_scheduler.LinearLR", + "kwargs": { + "start_factor": start_factor, + "end_factor": 1.0, + "total_iters": effective_warmup_steps, + }, + }, + { + "name": "torch.optim.lr_scheduler.CosineAnnealingLR", + "kwargs": { + "T_max": max(num_steps - effective_warmup_steps, 1), + "eta_min": min_lr, + }, + }, + { + "milestones": [effective_warmup_steps], + }, + ] + + return {} + + +def _build_logger_config( + customizer_config: TrainingStepConfig, + job_ctx: NMPJobContext, + workspace_dir: Path, +) -> dict[str, Any]: + """Build logger configuration for NeMo RL. + + WandB logging is handled by nemo-rl's Logger class when wandb_enabled is True. + The wandb config is passed directly to wandb.init(). + """ + wandb_config = build_wandb_config( + customizer_config=customizer_config, + job_ctx=job_ctx, + framework="nemo_rl", + ) + wandb_enabled = wandb_config is not None + # NeMo-RL's WandbLogger always passes `dir=` when initializing wandb. + # Avoid duplicate keyword errors by removing it from shared config here. + if wandb_config is not None: + wandb_config.pop("dir", None) + mlflow_config = build_mlflow_config( + customizer_config=customizer_config, + job_ctx=job_ctx, + framework="nemo_rl", + ) + mlflow_enabled = mlflow_config is not None + + # All four backend subsections (wandb / swanlab / tensorboard / mlflow) are + # always present — NeMo-RL's config schema expects them even when the + # corresponding ``*_enabled`` flag is False. We overlay the user's wandb / + # mlflow config when those integrations are enabled; the rest carry inert + # defaults. + return { + "log_dir": str(workspace_dir / "logs"), + "num_val_samples_to_print": 0, + "monitor_gpus": False, + "wandb_enabled": wandb_enabled, + "tensorboard_enabled": False, + "mlflow_enabled": mlflow_enabled, + "swanlab_enabled": False, + "wandb": wandb_config if (wandb_enabled and wandb_config) else {"project": "dpo", "name": "dpo"}, + "swanlab": {"project": "dpo", "name": "dpo"}, + "tensorboard": {"log_dir": str(workspace_dir / "tb_logs")}, + "mlflow": mlflow_config + if (mlflow_enabled and mlflow_config) + else {"experiment_name": "dpo", "run_name": "dpo", "tracking_uri": "http://localhost:5000"}, + "gpu_monitoring": { + "collection_interval": 10, + "flush_interval": 10, + }, + } diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py new file mode 100644 index 0000000000..148ea5c26d --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DPO training driver (ray run entry point). + +Entry point for DPO (Direct Preference Optimization) training, invoked via +ray run in a distributed environment. + +Preference data handling lives in NeMo-RL itself (``setup_preference_data`` plus +the config-driven ``BinaryPreferenceDataset`` / ``PreferenceDataset`` loaders): +the ``data`` config emitted by ``dpo_config.compile_dpo_config`` drives the +built-in loaders, so no custom preprocessor is needed. On top of NeMo-RL's DPO +loop we add the ``NemoRLLogger`` that streams progress back to the NeMo Platform +Jobs service. +""" + +import argparse +import logging +from typing import cast + +from nemo_rl.algorithms.dpo import MasterConfig, dpo_train, setup +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data.utils import setup_preference_data +from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.utils.config import load_config, parse_hydra_overrides +from nemo_rl.utils.logger import get_next_experiment_dir +from nmp.customization_common.service.context import NMPJobContext +from nmp.rl.tasks.training.backends.nemo_rl.nemo_rl_logger import NemoRLLogger +from nmp.rl.tasks.training.backends.nemo_rl.preference_datasets import register_preference_datasets +from omegaconf import OmegaConf + +logger = logging.getLogger(__name__) + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser(description="Run DPO training with configuration") + parser.add_argument("--config", type=str, required=True, help="Path to YAML config file") + parser.add_argument("--id", type=str, help="Customization ID") + parser.add_argument("--output-model", type=str, help="Output Model") + + # Parse known args for the script + args, overrides = parser.parse_known_args() + + return args, overrides + + +def main(): + """Main entry point.""" + args, overrides = parse_args() + + cfg = load_config(args.config) + print(f"Loaded configuration from: {args.config}") + + if overrides: + print(f"Overrides: {overrides}") + cfg = parse_hydra_overrides(cfg, overrides) + + config = cast(MasterConfig, OmegaConf.to_container(cfg, resolve=True)) + print("Applied CLI overrides") + + # Log only the top-level config section names. The resolved config carries + # integration secrets (W&B / MLflow tokens, tracking URIs), so never dump the + # full structure to stdout. + print(f"Config sections loaded: {sorted(config.keys())}") + + config["logger"]["log_dir"] = get_next_experiment_dir(config["logger"]["log_dir"]) + print(f"📊 Using log directory: {config['logger']['log_dir']}") + if config["checkpointing"]["enabled"]: + print(f"📊 Using checkpoint directory: {config['checkpointing']['checkpoint_dir']}") + + init_ray() + + # setup tokenizer + tokenizer = get_tokenizer(config["policy"]["tokenizer"]) + + # Register our local-file-capable HelpSteer3 / Tulu3 datasets into NeMo-RL's + # DATASET_REGISTRY before building data. Without this, setup_preference_data + # resolves those two formats to NeMo-RL's built-in classes, which always + # download from HuggingFace and ignore the uploaded local files. + register_preference_datasets() + + # setup data — NeMo-RL builds the datasets from the `data` config (per-split + # dataset specs). The compiler emits one of BinaryPreferenceDataset / + # PreferenceDataset / HelpSteer3 / Tulu3Preference per detected schema, each + # pointing at the prepared local training.jsonl / validation.jsonl. + dataset, val_dataset = setup_preference_data(tokenizer, config["data"]) + ( + policy, + cluster, + train_dataloader, + val_dataloader, + loss_fn, + logger, + checkpointer, + dpo_save_state, + master_config, + ) = setup(config, tokenizer, dataset, val_dataset) + + # Add NemoRLLogger for progress reporting if Jobs service is configured + job_ctx = NMPJobContext.from_env() + # Log only the non-sensitive job id; the full context carries service URLs + # and identifiers that should not be dumped to stdout. + print(f"Job context loaded (job_id={job_ctx.job_id})") + if job_ctx.jobs_url: + # Extract training parameters for progress reporting + max_steps = config["dpo"].get("max_num_steps", 0) + num_epochs = config["dpo"].get("max_num_epochs", 1) + steps_per_epoch = config["dpo"]["steps_per_epoch"] # type: ignore - we need to pass this additional parameter to the logger + log_interval = (config["dpo"]["val_period"] // 10) + 1 + + customizer_logger = NemoRLLogger( + steps_per_epoch=steps_per_epoch, + job_ctx=job_ctx, + log_interval=log_interval, + max_steps=max_steps, + num_epochs=num_epochs, + ) + # The setup() logger is a composite with a `.loggers` list; guard in case + # that internal shape changes. + if hasattr(logger, "loggers"): + logger.loggers.append(customizer_logger) + else: + print("WARNING: logger has no `.loggers`; NeMo Platform progress reporting disabled.") + + logger.log_hyperparams(config) + + dpo_train( + policy, + train_dataloader, + val_dataloader, + tokenizer, + loss_fn, + master_config, + logger, + checkpointer, + dpo_save_state, + ) + + +if __name__ == "__main__": + main() diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py new file mode 100644 index 0000000000..114b1c48d1 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""TrainingStepConfig → NeMo RL YAML configuration generation. + +This module handles configuration generation for GRPO training type, +converting the internal TrainingStepConfig format to NeMo RL's YAML format. + +Example of similar config but for DPO training type +- services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_config.py +""" + +import logging +from typing import Any + +from nmp.customization_common.service.context import NMPJobContext +from nmp.rl.app.jobs.training.schemas import ( + TrainingStepConfig, +) + +logger = logging.getLogger(__name__) + + +def compile_grpo_config( + training_config: TrainingStepConfig, + job_ctx: NMPJobContext, +) -> dict[str, Any]: + """Compile TrainingStepConfig to GRPO configuration. + + Args: + training_config: The training step configuration + job_ctx: Job context + + Returns: + Configuration dict for NeMo RL GRPO training + """ + # GRPO is not yet implemented; NemoRLBackend gates it out before this is + # reached. Kept as a placeholder for the future GRPO wiring. Do not log the + # config/job context here — they may carry integration secrets. + raise NotImplementedError("GRPO config compilation not yet implemented") diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py new file mode 100644 index 0000000000..0306a44fe5 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""GRPO training driver (torchrun entry point). + +This module serves as the entry point for GRPO (Group Relative Policy Optimization) +training, designed to be invoked via torchrun in a distributed environment. + +Migration source: customizer_training/rl/run_grpo_penguin.py +""" + +import argparse +from pathlib import Path + + +def parse_args() -> argparse.Namespace: + """Parse command line arguments. + + Returns: + Parsed arguments namespace + """ + parser = argparse.ArgumentParser(description="GRPO Training Driver for NeMo RL") + parser.add_argument( + "--config", + type=Path, + required=True, + help="Path to NeMo RL configuration YAML file", + ) + parser.add_argument( + "--environment", + type=str, + choices=["math", "code", "reward_model"], + help="Override environment type from config", + ) + parser.add_argument( + "--output-dir", + type=Path, + help="Override output directory from config", + ) + parser.add_argument( + "--resume-from-checkpoint", + type=Path, + help="Path to checkpoint to resume training from", + ) + return parser.parse_args() + + +def load_config(config_path: Path) -> dict: + """Load NeMo RL configuration from YAML file. + + Args: + config_path: Path to the configuration file + + Returns: + Configuration dictionary + """ + # TODO: Implement YAML config loading + raise NotImplementedError + + +def get_environment(env_type: str): + """Get the GRPO environment based on type. + + Args: + env_type: Environment type (math, code, reward_model) + + Returns: + Configured environment instance + """ + # TODO: Import and instantiate appropriate environment + # from .environments import math, code, reward_model + raise NotImplementedError + + +def run_grpo_training(config: dict) -> dict: + """Execute GRPO training with the given configuration. + + Args: + config: NeMo RL configuration dictionary + + Returns: + Training metrics dictionary + """ + # TODO: Implement GRPO training execution + # - Initialize model and tokenizer + # - Load training dataset + # - Configure GRPO environment + # - Configure GRPO trainer + # - Run training loop with group sampling + # - Save checkpoints + # - Return metrics + raise NotImplementedError + + +def main() -> None: + """Main entry point for GRPO training.""" + raise NotImplementedError + + +if __name__ == "__main__": + main() diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py new file mode 100644 index 0000000000..f7534bc2c7 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import logging +import math +from typing import Any, Mapping, Optional + +from nemo_rl.utils.logger import LoggerInterface +from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.training.progress import JobsServiceProgressReporter +from nmp.rl.app.constants import SERVICE_NAME +from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback + +_logger = logging.getLogger(__name__) + + +def has_metric_value(metric: Any) -> bool: + """Check if a metric has a valid value.""" + if metric is not None and not math.isnan(metric): + return True + return False + + +class NemoRLLogger(LoggerInterface): + """ + NemoRLLogger is a logger implementation that reports training updates to Jobs Service. + + It implements the LoggerInterface from nemo_rl.utils.logger to provide a consistent + logging interface while maintaining compatibility with the Jobs Service. + + This implementation uses TrainingProgressCallback with JobsServiceProgressReporter + to report progress via the NeMo Platform SDK. + """ + + 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, + ): + """Initialize the NemoRL logger. + + Args: + steps_per_epoch: Number of steps per epoch (required for accurate epoch calculation). + job_ctx: NeMo Platform job context for progress reporting (defaults to environment variables). + log_interval: Number of steps between progress updates. + max_steps: Total number of training steps (optional, used for progress reporting). + num_epochs: Total number of epochs (optional, used for progress reporting). + + Raises: + ValueError: If ``steps_per_epoch`` or ``log_interval`` is < 1. Both are + used as divisors/moduli in ``log_metrics`` (epoch derivation and + log-interval throttling), so non-positive values are rejected up + front to fail fast instead of raising ZeroDivisionError mid-training. + """ + if steps_per_epoch < 1: + raise ValueError(f"steps_per_epoch must be >= 1, got {steps_per_epoch}") + if log_interval < 1: + raise ValueError(f"log_interval must be >= 1, got {log_interval}") + + self._job_ctx = job_ctx or NMPJobContext.from_env() + self._log_interval = log_interval + self._max_steps = max_steps + self._num_epochs = num_epochs + self._steps_per_epoch = steps_per_epoch + + # Create the callback for progress reporting + self._reporter = JobsServiceProgressReporter(self._job_ctx, SERVICE_NAME) + self._callback = TrainingProgressCallback(self._reporter) + + # Track best metrics for monitoring + self._best_metric_value = float("inf") + self._best_epoch: int | None = None + self._closed = False + + _logger.info( + f"Initialized NemoRLLogger with jobs_url={self._job_ctx.jobs_url}, " + f"log_interval={log_interval}, max_steps={max_steps}, num_epochs={num_epochs}, " + f"steps_per_epoch={steps_per_epoch}" + ) + + def log_metrics( + self, + metrics: dict[str, Any], + step: int, + prefix: Optional[str] = "", + step_metric: Optional[str] = None, + step_finished: bool = False, + ) -> None: + """Log metrics to NeMo Customizer. + + Args: + metrics: Dict of metrics to log + step: Global step value + prefix: Optional prefix for metric names (e.g. "train", "validation", "timing/train") + step_metric: Optional step metric name (ignored in this implementation) + step_finished: Whether the step is finished (part of NeMo-RL's LoggerInterface; ignored here) + """ + step = step + 1 # Increment step since we start counting from 1 + + # Calculate epoch from step (epochs start from 1) + epoch = ((step - 1) // self._steps_per_epoch) + 1 + + # Handle training loss + if prefix == "train" and has_metric_value(metrics.get("loss")): + # Only report at log_interval to reduce output + if step % self._log_interval == 0: + # Extract core metrics + loss = metrics["loss"] + lr = metrics.get("lr") + grad_norm = metrics.get("grad_norm") + + # Extract additional training metrics (whitelisted only) + additional_metrics = {} + for key in [ + "num_valid_samples", + "preference_loss", + "rewards_rejected_mean", + "global_valid_seqs", + "global_valid_toks", + ]: + if has_metric_value(metrics.get(key)): + additional_metrics[key] = metrics[key] + + self._callback.report_train_step( + step=step, + epoch=epoch, + loss=loss, + lr=lr, + grad_norm=grad_norm, + **additional_metrics, + ) + + # Handle validation metrics + elif prefix and prefix.startswith("validation"): + if has_metric_value(metrics.get("loss")): + val_loss = metrics["loss"] + + # Extract additional validation metrics (whitelisted only) + additional_metrics = {} + for key in [ + "num_valid_samples", + "preference_loss", + "rewards_rejected_mean", + "global_valid_seqs", + "global_valid_toks", + ]: + if has_metric_value(metrics.get(key)): + additional_metrics[key] = metrics[key] + + self._callback.report_validation( + step=step, + epoch=epoch, + val_loss=val_loss, + **additional_metrics, + ) + # Track best validation loss + if val_loss < self._best_metric_value: + self._best_metric_value = val_loss + self._best_epoch = epoch + + _logger.debug(f"log_metrics: step={step}, prefix={prefix}, metrics={metrics}") + + def log_hyperparams(self, params: Mapping[str, Any]) -> None: + """Log hyperparameters and report training start. + + Args: + params: Dictionary of hyperparameters to log + """ + # Extract max_steps and num_epochs from params if not already set + max_steps = self._max_steps or params.get("max_steps", 0) + num_epochs = self._num_epochs or params.get("num_epochs", 1) + + # Update internal tracking if extracted from params + if not self._max_steps and max_steps: + self._max_steps = max_steps + if not self._num_epochs and num_epochs: + self._num_epochs = num_epochs + + self._callback.report_training_start(max_steps=max_steps, num_epochs=num_epochs) + _logger.debug(f"log_hyperparams: max_steps={max_steps}, num_epochs={num_epochs}") + + def log_histogram(self, histogram: list[Any], step: int, name: str) -> None: + """No-op: required by NeMo-RL's LoggerInterface. + + Jobs Service progress reporting has no histogram concept, so there is + nothing to forward. Implemented only to satisfy the abstract base class. + """ + return None + + def log_plot(self, figure: Any, step: int, name: str) -> None: + """No-op: required by NeMo-RL's LoggerInterface. + + ``figure`` is a ``matplotlib.figure.Figure``; typed ``Any`` so we don't + import matplotlib. Jobs Service has no figure/plot concept, so this is a + no-op implemented only to satisfy the abstract base class. + """ + return None + + def close(self) -> None: + """Clean up resources.""" + if self._closed: + return + self._closed = True + _logger.info("NemoRLLogger closing") + self._callback.close() + + def __del__(self): + """Cleanup when the logger is destroyed.""" + try: + if hasattr(self, "_closed") and not self._closed: + self.close() + except Exception: + # Silently ignore errors during interpreter shutdown + pass diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt new file mode 100644 index 0000000000..c6adce77aa --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/no_override_requirements.txt @@ -0,0 +1,22 @@ +# Packages whose base-image versions MUST be preserved when adding the platform +# glue into NeMo-RL's venv. NeMo-RL pins these to versions tuned for its +# Ray/torch ML stack; upgrading any of them can break training. +# +# How preservation is actually enforced: the RL training image +# (docker/Dockerfile.nmp-rl-training) installs the glue with `uv pip install +# --no-deps`, which skips dependency resolution entirely and therefore never +# upgrades these packages. This file is the documented "protected set" (parity +# with docker/unsloth and docker/automodel). +# +# NOTE: the `sys_platform == 'never'` marker makes each line a deliberate no-op +# under `uv pip install --overrides` — the marker is always false, so the line +# never contributes a requirement and, on its own, pins nothing. Version +# preservation comes from `--no-deps` above, not from these markers. If this +# file is ever switched to a hard-pinning `--overrides`/constraints flow, +# replace these lines with explicit `==` pins captured from the base +# image at build time. +prometheus-client; sys_platform == 'never' +starlette; sys_platform == 'never' +cryptography; sys_platform == 'never' +mlflow; sys_platform == 'never' +ray; sys_platform == 'never' diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.py new file mode 100644 index 0000000000..435d69a3c2 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/__init__.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Preference datasets for DPO training. + +NeMo-RL's built-in ``HelpSteer3Dataset`` / ``Tulu3PreferenceDataset`` only download +their datasets from HuggingFace and ignore a local ``data_path``. This package +provides local-file-capable subclasses and registers them into NeMo-RL's +``DATASET_REGISTRY`` so the library's ``setup_preference_data`` / +``load_preference_dataset`` resolve a user-uploaded HelpSteer3/Tulu3 dataset to the +local file instead of silently training on the public HuggingFace dataset. + +``BinaryPreferenceDataset`` / ``PreferenceDataset`` already load from a local path, +so they are re-exported unchanged. +""" + +from nemo_rl.data.datasets.preference_datasets import ( + DATASET_REGISTRY, + BinaryPreferenceDataset, + PreferenceDataset, +) +from nmp.rl.tasks.training.backends.nemo_rl.preference_datasets.helpsteer3 import HelpSteer3Dataset +from nmp.rl.tasks.training.backends.nemo_rl.preference_datasets.tulu3 import Tulu3PreferenceDataset + + +def register_preference_datasets() -> None: + """Override NeMo-RL's HF-only HelpSteer3 / Tulu3 with local-file-capable subclasses. + + NeMo-RL's ``load_preference_dataset`` resolves ``dataset_name`` via the + module-level ``DATASET_REGISTRY`` dict and constructs ``cls(**data_config)``. + Re-pointing the "HelpSteer3" / "Tulu3Preference" entries at our subclasses makes + ``setup_preference_data`` honor the compiled ``data_path`` for those formats. + + Must be called before ``setup_preference_data`` (e.g. at driver start-up). + Idempotent — safe to call more than once. + """ + DATASET_REGISTRY["HelpSteer3"] = HelpSteer3Dataset + DATASET_REGISTRY["Tulu3Preference"] = Tulu3PreferenceDataset + + +__all__ = [ + "BinaryPreferenceDataset", + "HelpSteer3Dataset", + "PreferenceDataset", + "Tulu3PreferenceDataset", + "register_preference_datasets", +] diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py new file mode 100644 index 0000000000..f918d29fa4 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/helpsteer3.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""HelpSteer3 preference dataset with local-file support.""" + +from typing import Optional + +from nemo_rl.data.datasets.preference_datasets.helpsteer3 import ( + HelpSteer3Dataset as BaseHelpSteer3Dataset, +) +from nemo_rl.data.datasets.utils import load_dataset_from_path + + +class HelpSteer3Dataset(BaseHelpSteer3Dataset): + """HelpSteer3 preference dataset for DPO training, extended for local files. + + NeMo-RL's base ``HelpSteer3Dataset`` only downloads ``nvidia/HelpSteer3`` from + HuggingFace and ignores any local path. This subclass adds local-file support: + when ``data_path`` is provided it loads that JSONL (HelpSteer3 schema — + ``context`` / ``response1`` / ``response2`` / ``overall_preference``) and reuses + the base class's :meth:`format_data` to produce the canonical + ``{context, completions, task_name}`` shape. With no ``data_path`` it falls back + to the base HuggingFace download. + + NeMo-RL's ``load_preference_dataset`` instantiates the registered class via + ``cls(**data_config)``, so ``__init__`` accepts the per-split spec keys + (``data_path``, plus ``dataset_name`` and any others swallowed by ``**kwargs``). + The ``task_spec`` is bound afterwards by the library via ``set_task_spec``. + """ + + def __init__( + self, + data_path: Optional[str] = None, + subset: Optional[str] = None, + split: str = "train", + **kwargs, + ) -> None: + if data_path is None: + # No local file → keep the base HuggingFace download behavior. + super().__init__(split=split, **kwargs) + return + + self.task_name = "HelpSteer3" + # Load from the local file (or HuggingFace) and apply the base formatting. + self.dataset = load_dataset_from_path(data_path, subset, split) + self.dataset = self.dataset.map( + self.format_data, + remove_columns=self.dataset.column_names, + ) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py new file mode 100644 index 0000000000..1796ccabc3 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/preference_datasets/tulu3.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tulu3 preference dataset with local-file support.""" + +from typing import Optional + +from nemo_rl.data.datasets.preference_datasets.tulu3 import ( + Tulu3PreferenceDataset as BaseTulu3PreferenceDataset, +) +from nemo_rl.data.datasets.utils import load_dataset_from_path + + +class Tulu3PreferenceDataset(BaseTulu3PreferenceDataset): + """Tulu3 preference dataset for DPO training, extended for local files. + + NeMo-RL's base ``Tulu3PreferenceDataset`` only downloads + ``allenai/llama-3.1-tulu-3-8b-preference-mixture`` from HuggingFace and ignores + any local path. This subclass adds local-file support: when ``data_path`` is + provided it loads that JSONL (Tulu3 schema — ``chosen`` / ``rejected`` message + lists) and reuses the base class's :meth:`format_data` to produce the canonical + ``{context, completions, task_name}`` shape. With no ``data_path`` it falls back + to the base HuggingFace download. + + NeMo-RL's ``load_preference_dataset`` instantiates the registered class via + ``cls(**data_config)``, so ``__init__`` accepts the per-split spec keys + (``data_path``, plus ``dataset_name`` and any others swallowed by ``**kwargs``). + The ``task_spec`` is bound afterwards by the library via ``set_task_spec``. + """ + + def __init__( + self, + data_path: Optional[str] = None, + subset: Optional[str] = None, + split: str = "train", + **kwargs, + ) -> None: + if data_path is None: + # No local file → keep the base HuggingFace download behavior. Forward + # split for parity with the local branch and HelpSteer3Dataset (the base + # currently hard-codes its split, but absorbs the kwarg harmlessly). + super().__init__(split=split, **kwargs) + return + + self.task_name = "Tulu3Preference" + # Load from the local file (or HuggingFace) and apply the base formatting. + self.dataset = load_dataset_from_path(data_path, subset, split) + self.dataset = self.dataset.map( + self.format_data, + remove_columns=self.dataset.column_names, + ) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py new file mode 100644 index 0000000000..a909d8c476 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/ray_bootstrap.py @@ -0,0 +1,876 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. +"""Python equivalent of run-ray.sh for Ray cluster bootstrap. + +This module provides a Python implementation of Ray cluster bootstrapping +on Volcano-provisioned pods, replacing the shell script approach for +better integration and error handling. +""" + +from __future__ import annotations + +import logging +import os +import re +import signal +import subprocess +import sys +import threading +import time +from collections import deque +from dataclasses import dataclass, field +from pathlib import Path +from types import FrameType + +from nmp.rl.tasks.training.errors.exceptions import format_exception_string +from nmp.rl.tasks.training.errors.parser import ( + MAX_OUTPUT_LINES, + read_subprocess_output, +) + +logger = logging.getLogger(__name__) + +# Timeouts (seconds) for Ray CLI subprocesses so a hung `ray` invocation can't +# block the bootstrap — or its non-daemon cleanup thread — indefinitely. +# `ray start` may pull images / initialize for a while; `ray status` / `ray +# memory` are quick queries. +RAY_START_TIMEOUT_SECONDS = 300 +RAY_STATUS_TIMEOUT_SECONDS = 60 + + +def _pause(seconds: float) -> None: + time.sleep(seconds) + + +@dataclass +class RayPortConfig: + """Port configuration for Ray cluster services. + + All ports are configurable via environment variables with sensible defaults. + Head nodes use port+1 offset for manager ports to avoid conflicts with workers. + """ + + node_manager_port: int = field(default_factory=lambda: int(os.getenv("NODE_MANAGER_PORT", "53001"))) + object_manager_port: int = field(default_factory=lambda: int(os.getenv("OBJECT_MANAGER_PORT", "53003"))) + runtime_env_agent_port: int = field(default_factory=lambda: int(os.getenv("RUNTIME_ENV_AGENT_PORT", "53005"))) + dashboard_agent_grpc_port: int = field(default_factory=lambda: int(os.getenv("DASHBOARD_AGENT_GRPC_PORT", "53007"))) + metrics_export_port: int = field(default_factory=lambda: int(os.getenv("METRICS_EXPORT_PORT", "53009"))) + gcs_port: int = field(default_factory=lambda: int(os.getenv("GCS_PORT", "6379"))) + ray_client_server_port: int = field(default_factory=lambda: int(os.getenv("RAY_CLIENT_SERVER_PORT", "10001"))) + dashboard_port: int = field(default_factory=lambda: int(os.getenv("DASHBOARD_PORT", "8265"))) + dashboard_agent_listen_port: int = field( + default_factory=lambda: int(os.getenv("DASHBOARD_AGENT_LISTEN_PORT", "52365")) + ) + min_worker_port: int = field(default_factory=lambda: int(os.getenv("MIN_WORKER_PORT", "54001"))) + max_worker_port: int = field(default_factory=lambda: int(os.getenv("MAX_WORKER_PORT", "54257"))) + + +@dataclass +class RayClusterBootstrap: + """Bootstrap Ray cluster on Volcano-provisioned pods. + + This class handles starting Ray head nodes and worker nodes in a + distributed training environment managed by Volcano job scheduler. + + The bootstrap process: + - Head (rank 0): Start Ray head -> wait for workers -> run driver -> cleanup + - Worker (rank > 0): Start Ray worker -> monitor for ENDED file -> exit + + Attributes: + rank: The rank of this node (0 = head, >0 = worker) + world_size: Total number of nodes in the cluster + master_addr: IP address of the head node + gpus_per_node: Number of GPUs per node + log_dir: Directory for logs and coordination files. For multi-node clusters, + this MUST be a shared filesystem (e.g., NFS) accessible by all nodes. + The head node writes an ENDED marker file here that workers poll for + graceful shutdown coordination. Set via BASE_LOG_DIR environment variable. + ports: Port configuration for Ray services + num_retries: Number of retries for Ray start commands + retry_sleep: Seconds to sleep between retries + driver_python: Python executable for running driver scripts (allows using + a different virtual environment). Defaults to DRIVER_PYTHON env var + or current Python interpreter. + driver_extra_pythonpath: Additional paths to append to PYTHONPATH when running + driver scripts. Useful for accessing packages from other environments. + Defaults to DRIVER_EXTRA_PYTHONPATH env var. + + Example: + Basic usage with environment variables (recommended for Volcano jobs):: + + # Environment variables set by Volcano: RANK, WORLD_SIZE, MASTER_ADDR + bootstrap = create_bootstrap_from_env() + exit_code = bootstrap.run_with_driver( + driver_script="/path/to/dpo_driver.py", + driver_args=["--config", "/path/to/config.yaml", "--id", "job-123"], + ) + sys.exit(exit_code) + + Manual configuration for testing:: + + bootstrap = RayClusterBootstrap( + rank=0, # Head node + world_size=2, # 2-node cluster + master_addr="10.0.0.1", + gpus_per_node=8, + ) + + # Option 1: Start cluster and run driver script + exit_code = bootstrap.run_with_driver( + driver_script="train_dpo.py", + driver_args=["--config", "config.yaml"], + ) + + # Option 2: Just start the cluster (for workers or manual control) + bootstrap.start() + + Using a different Python virtual environment for the driver:: + + # Via environment variables + os.environ["DRIVER_PYTHON"] = "/opt/nemo-venv/bin/python" + os.environ["DRIVER_EXTRA_PYTHONPATH"] = "/opt/venv/lib/python3.12/site-packages" + bootstrap = create_bootstrap_from_env() + + # Or via direct configuration + bootstrap = RayClusterBootstrap( + rank=0, + world_size=1, + master_addr="127.0.0.1", + driver_python="/opt/nemo-venv/bin/python", # Custom venv + driver_extra_pythonpath="/opt/venv/lib/python3.12/site-packages", # Extra packages + ) + + Command-line invocation:: + + # Start cluster and run driver + python -m nmp.rl.tasks.training.backends.nemo_rl.ray_bootstrap \\ + /path/to/driver.py --config config.yaml --id job-123 + + # Just start cluster node (head or worker based on RANK env var) + python -m nmp.rl.tasks.training.backends.nemo_rl.ray_bootstrap + """ + + rank: int + world_size: int + master_addr: str + gpus_per_node: int = field(default_factory=lambda: int(os.getenv("GPUS_PER_NODE", "1"))) + log_dir: Path = field(default_factory=lambda: Path(os.getenv("BASE_LOG_DIR", "/tmp")) / "logs") + ports: RayPortConfig = field(default_factory=RayPortConfig) + num_retries: int = 3 + retry_sleep: int = 20 + attempt_id: str = field(default_factory=lambda: os.getenv("NEMO_JOB_ATTEMPT_ID", "attempt-0")) + """Job attempt id, used to scope the ENDED coordination marker so a stale + marker left by a previous attempt cannot short-circuit a retry's startup.""" + driver_python: str = field(default_factory=lambda: os.getenv("DRIVER_PYTHON", sys.executable)) + """Python executable path for running driver scripts. + + This allows running driver scripts in a different virtual environment. + Can be set via DRIVER_PYTHON environment variable or passed directly. + Defaults to sys.executable (current Python interpreter). + """ + + ray_executable: str = field(default="") + """Path to the ray executable. If empty, derived from driver_python's directory.""" + + driver_extra_pythonpath: str = field(default_factory=lambda: os.getenv("DRIVER_EXTRA_PYTHONPATH", "")) + """Additional paths to append to PYTHONPATH when running driver scripts. + + Multiple paths can be separated by colons (Unix) or semicolons (Windows). + Can be set via DRIVER_EXTRA_PYTHONPATH environment variable or passed directly. + Example: "/opt/venv/lib/python3.12/site-packages:/other/path" + """ + + # Internal state + _stop_event: threading.Event = field(default_factory=threading.Event, repr=False) + _driver_output: deque[str] = field(default_factory=lambda: deque(maxlen=MAX_OUTPUT_LINES), repr=False) + _driver_process: subprocess.Popen | None = field(default=None, repr=False) + + def __post_init__(self) -> None: + """Initialize the bootstrap environment.""" + # Warn if multi-node setup might have coordination issues + if self.world_size > 1 and not os.getenv("BASE_LOG_DIR"): + logger.warning( + "Multi-node Ray cluster detected (world_size=%d) but BASE_LOG_DIR not set. " + "The ENDED coordination file requires a shared filesystem across all nodes. " + "Workers may not detect graceful termination if log_dir (%s) is not shared.", + self.world_size, + self.log_dir, + ) + + # Ensure log directory exists + self.log_dir.mkdir(parents=True, exist_ok=True) + + # Derive ray executable from driver_python if not specified + if not self.ray_executable: + # Get the bin directory from driver_python path + driver_bin_dir = Path(self.driver_python).parent + self.ray_executable = str(driver_bin_dir / "ray") + + # Disable proxy environment variables for local pod communication + self._unset_proxy_env() + + def _unset_proxy_env(self) -> None: + """Unset proxy environment variables for local pod communication.""" + proxy_vars = ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"] + for var in proxy_vars: + os.environ.pop(var, None) + + @property + def driver_output(self) -> deque[str]: + """Rolling buffer of recent driver output lines for error extraction.""" + return self._driver_output + + def terminate_driver(self, signum: int = signal.SIGTERM, timeout: int = 30) -> None: + """Terminate the driver subprocess if it is running. + + Sends the specified signal to the driver process and waits for it to exit. + If it doesn't exit within the timeout, it is forcefully killed. + + Args: + signum: Signal to send (default SIGTERM). + timeout: Seconds to wait for graceful exit before killing. + """ + process = self._driver_process + if process is None or process.poll() is not None: + return + logger.warning(f"Terminating driver process (pid={process.pid}) with signal {signum}") + try: + process.send_signal(signum) + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + logger.warning(f"Driver process did not exit within {timeout}s, killing") + process.kill() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.warning("Killed driver process did not terminate within 5s") + + @property + def ended_file(self) -> Path: + """Path to the ENDED coordination file. + + Scoped by the job attempt id: the head writes this marker during cleanup, + and on shared (multi-node) storage it would otherwise persist into the + next attempt and immediately short-circuit the retry's head/worker + startup. A new attempt id yields a fresh path with no stale marker. + """ + return self.log_dir / f"ENDED.{self.attempt_id}" + + def _clear_ended_marker(self) -> None: + """Remove a lingering ENDED marker before the head starts a fresh run. + + The marker is attempt-scoped so a normal retry already gets a new path; + this additionally covers attempt-id reuse (e.g. local runs / resume) where + a marker from the prior run could otherwise linger on shared storage. + """ + try: + self.ended_file.unlink(missing_ok=True) + except OSError as e: + logger.warning(f"Failed to clear stale ENDED marker {self.ended_file}: {e}") + + @property + def expected_worker_units(self) -> int: + """Total expected worker units (nodes * GPUs per node).""" + return self.world_size * self.gpus_per_node + + @property + def max_wait_seconds(self) -> int: + """Maximum wait time for workers to connect. + + Multi-node clusters get 40 minutes to accommodate slow image downloads. + Single-node clusters only need 4 minutes. + """ + return 2400 if self.world_size > 1 else 240 + + def start(self) -> None: + """Start Ray head (rank 0) or worker (rank > 0).""" + if self.rank == 0: + self._run_as_head() + else: + self._run_as_worker() + + def run_with_driver(self, driver_script: str, driver_args: list[str]) -> int: + """Start Ray cluster and run driver script on head node. + + This is the main entry point for executing training with Ray. + + Args: + driver_script: Path to the Python driver script + driver_args: Arguments to pass to the driver script + + Returns: + Exit code from the driver script (0 for success) + """ + if self.rank == 0: + return self._run_head_with_driver(driver_script, driver_args) + else: + return self._run_as_worker() + + def _run_as_head(self) -> None: + """Run as head node: start head, wait for workers, then exit.""" + self._clear_ended_marker() + if not self._start_head_background(): + raise RuntimeError("Failed to start Ray head node") + self._wait_for_workers() + + def _run_head_with_driver(self, driver_script: str, driver_args: list[str]) -> int: + """Run as head node with driver execution. + + Args: + driver_script: Path to the Python driver script + driver_args: Arguments to pass to the driver script + + Returns: + Exit code from driver execution + """ + exit_code = 1 + try: + self._clear_ended_marker() + if not self._start_head_background(): + raise RuntimeError("Failed to start Ray head node") + + # Wait a bit for Ray to fully initialize before checking status + print("[ray_bootstrap] Waiting for Ray to initialize...", flush=True) + _pause(5) + + print("[ray_bootstrap] Waiting for workers to connect...", flush=True) + self._wait_for_workers() + + logger.info("--- All workers connected! ---") + print("[ray_bootstrap] --- All workers connected! ---", flush=True) + self._log_ray_status() + + logger.info("--- Starting driver ---") + print(f"[ray_bootstrap] --- Starting driver: {driver_script} ---", flush=True) + exit_code = self._run_driver(driver_script, driver_args) + logger.info(f"Driver completed with exit code: {exit_code}") + print(f"[ray_bootstrap] Driver completed with exit code: {exit_code}", flush=True) + + except Exception as e: + logger.exception(f"Error in head node execution: {e}") + print(f"[ray_bootstrap] Error in head node execution: {e}", flush=True) + # Capture the exception message into the output buffer so the + # backend's parse_error_from_output can surface it to the user. + self._driver_output.append(format_exception_string(e)) + exit_code = 1 + finally: + self._cleanup_with_timeout() + + return exit_code + + def _run_as_worker(self) -> int: + """Run as worker node: start worker and monitor for termination. + + Returns: + Exit code (0 for graceful termination, non-zero otherwise) + """ + if not self._start_worker_background(): + return 1 + return self._monitor_for_termination() + + def _start_head_background(self) -> bool: + """Start Ray head node with retry logic. + + Since ray start returns immediately (without --block), we run this + synchronously with retries rather than in a background thread. + + Returns: + True if head started successfully, False otherwise + """ + for attempt in range(self.num_retries): + if self._stop_event.is_set() or self.ended_file.exists(): + logger.info("Head node stopping due to termination signal") + return False + + logger.info(f"Launching Head Node (attempt {attempt + 1}/{self.num_retries})") + print( + f"[ray_bootstrap] Launching Head Node (attempt {attempt + 1}/{self.num_retries})", + flush=True, + ) + try: + result = self._start_head_process() + if result is not None: + return True + logger.warning(f"Head start failed, attempt {attempt + 1}/{self.num_retries}") + except Exception as e: + logger.exception(f"Head node error: {e}") + print(f"[ray_bootstrap] Head node error: {e}", flush=True) + + if not self._stop_event.is_set() and not self.ended_file.exists(): + _pause(self.retry_sleep) + + logger.error("Head Node failed to start after all retries") + print("[ray_bootstrap] Head Node failed to start after all retries", flush=True) + return False + + def _start_worker_background(self) -> bool: + """Start Ray worker node with retry logic. + + Since ray start returns immediately (without --block), we run this + synchronously with retries rather than in a background thread. + + Returns: + True if worker started successfully, False otherwise + """ + for attempt in range(self.num_retries): + if self._stop_event.is_set() or self.ended_file.exists(): + logger.info("Worker node stopping due to termination signal") + return False + + logger.info(f"Launching Worker Node (attempt {attempt + 1}/{self.num_retries})") + print( + f"[ray_bootstrap] Launching Worker Node (attempt {attempt + 1}/{self.num_retries})", + flush=True, + ) + try: + result = self._start_worker_process() + if result is not None: + return True + logger.warning(f"Worker start failed, attempt {attempt + 1}/{self.num_retries}") + except Exception as e: + logger.exception(f"Worker node error: {e}") + print(f"[ray_bootstrap] Worker node error: {e}", flush=True) + + if not self._stop_event.is_set() and not self.ended_file.exists(): + _pause(self.retry_sleep) + + logger.error("Worker Node failed to start after all retries") + print("[ray_bootstrap] Worker Node failed to start after all retries", flush=True) + return False + + def _start_head_process(self) -> subprocess.CompletedProcess | None: + """Start the Ray head process. + + Note: Unlike the bash script which uses --block, we don't need it here + because ray start returns immediately and Ray continues running in the + background. The bash script needed --block to keep the background job alive. + + Returns: + The CompletedProcess result from ray start, or None on failure + """ + p = self.ports + cmd = [ + self.ray_executable, + "start", + "--head", + "--disable-usage-stats", + "--include-dashboard=false", + f'--resources={{"worker_units": {self.gpus_per_node}}}', + f"--node-ip-address={self.master_addr}", + f"--port={p.gcs_port}", + f"--ray-client-server-port={p.ray_client_server_port}", + f"--dashboard-port={p.dashboard_port}", + # Head uses port+1 offset to avoid conflicts + f"--node-manager-port={p.node_manager_port + 1}", + f"--object-manager-port={p.object_manager_port + 1}", + f"--runtime-env-agent-port={p.runtime_env_agent_port + 1}", + f"--dashboard-agent-grpc-port={p.dashboard_agent_grpc_port + 1}", + f"--dashboard-agent-listen-port={p.dashboard_agent_listen_port + 1}", + f"--metrics-export-port={p.metrics_export_port + 1}", + ] + logger.info(f"Starting head: {' '.join(cmd)}") + print(f"[ray_bootstrap] Starting head: {' '.join(cmd)}", flush=True) + try: + result = subprocess.run(cmd, check=False, timeout=RAY_START_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + logger.warning(f"Ray head start timed out after {RAY_START_TIMEOUT_SECONDS}s") + print(f"[ray_bootstrap] Ray head start timed out after {RAY_START_TIMEOUT_SECONDS}s", flush=True) + return None + if result.returncode != 0: + print(f"[ray_bootstrap] Head start failed with code {result.returncode}", flush=True) + return None + print("[ray_bootstrap] Head node started successfully", flush=True) + return result + + def _start_worker_process(self) -> subprocess.CompletedProcess | None: + """Start the Ray worker process. + + Note: Unlike the bash script which uses --block, we don't need it here + because ray start returns immediately and Ray continues running in the + background. + + Returns: + The CompletedProcess result from ray start, or None on failure + """ + p = self.ports + cmd = [ + self.ray_executable, + "start", + f"--address={self.master_addr}:{p.gcs_port}", + "--disable-usage-stats", + f'--resources={{"worker_units": {self.gpus_per_node}}}', + f"--min-worker-port={p.min_worker_port}", + f"--max-worker-port={p.max_worker_port}", + f"--node-manager-port={p.node_manager_port}", + f"--object-manager-port={p.object_manager_port}", + f"--runtime-env-agent-port={p.runtime_env_agent_port}", + f"--dashboard-agent-grpc-port={p.dashboard_agent_grpc_port}", + f"--dashboard-agent-listen-port={p.dashboard_agent_listen_port}", + f"--metrics-export-port={p.metrics_export_port}", + ] + logger.info(f"Starting worker: {' '.join(cmd)}") + print(f"[ray_bootstrap] Starting worker: {' '.join(cmd)}", flush=True) + try: + result = subprocess.run(cmd, check=False, timeout=RAY_START_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + logger.warning(f"Ray worker start timed out after {RAY_START_TIMEOUT_SECONDS}s") + print(f"[ray_bootstrap] Ray worker start timed out after {RAY_START_TIMEOUT_SECONDS}s", flush=True) + return None + if result.returncode != 0: + print(f"[ray_bootstrap] Worker start failed with code {result.returncode}", flush=True) + return None + print("[ray_bootstrap] Worker node started successfully", flush=True) + return result + + def _wait_for_workers(self) -> None: + """Poll until all workers have connected to the cluster. + + Raises: + TimeoutError: If workers don't connect within max_wait_seconds + """ + poll_interval = 2 + elapsed = 0 + + while elapsed < self.max_wait_seconds: + if self.ended_file.exists(): + raise RuntimeError("ENDED file detected during worker wait") + + worker_units = self._get_worker_units() + logger.info(f"[INFO] Number of actors online: {worker_units}/{self.expected_worker_units}") + print( + f"[ray_bootstrap] Workers online: {worker_units}/{self.expected_worker_units}", + flush=True, + ) + + if worker_units >= self.expected_worker_units: + return + + _pause(poll_interval) + elapsed += poll_interval + + raise TimeoutError( + f"Timed out waiting for all workers to connect after {self.max_wait_seconds}s. " + f"Expected {self.expected_worker_units} worker_units." + ) + + def _get_worker_units(self) -> int: + """Extract worker_units from ray status output. + + Returns: + Total number of worker_units available in the cluster + """ + try: + result = subprocess.run( + [self.ray_executable, "status"], + capture_output=True, + text=True, + check=False, + timeout=RAY_STATUS_TIMEOUT_SECONDS, + ) + except subprocess.TimeoutExpired: + logger.warning(f"ray status timed out after {RAY_STATUS_TIMEOUT_SECONDS}s") + print(f"[ray_bootstrap] ray status timed out after {RAY_STATUS_TIMEOUT_SECONDS}s", flush=True) + return 0 + + if result.returncode != 0: + logger.warning(f"ray status failed: {result.stderr}") + print(f"[ray_bootstrap] ray status failed: {result.stderr}", flush=True) + return 0 + + return self._parse_worker_units(result.stdout) + + @staticmethod + def _parse_worker_units(status_output: str) -> int: + """Parse worker_units from ray status output. + + The ray status output contains lines like: + 0.0/1.0 worker_units + + Where the format is: usage/total resource_name + We want the TOTAL (second number) as that's how many worker_units are available. + + Args: + status_output: Output from `ray status` command + + Returns: + Total number of worker_units available in the cluster + """ + # Match pattern: " 0.0/1.0 worker_units" - extract the total (second number) + match = re.search(r"(\d+(?:\.\d+)?)/(\d+(?:\.\d+)?)\s+worker_units", status_output) + if match: + total = int(float(match.group(2))) + print( + f"[ray_bootstrap] Parsed worker_units: {match.group(1)}/{match.group(2)} -> total={total}", flush=True + ) + return total + + # Fallback: look for lines containing worker_units + for line in status_output.splitlines(): + if "worker_units" in line: + # Try to extract numbers from format "X/Y worker_units" + parts = line.strip().split() + if len(parts) >= 2 and "/" in parts[0]: + usage_total = parts[0].split("/") + if len(usage_total) == 2: + try: + total = int(float(usage_total[1])) + print( + f"[ray_bootstrap] Fallback parsed worker_units: {parts[0]} -> total={total}", flush=True + ) + return total + except ValueError: + pass + + print("[ray_bootstrap] Could not parse worker_units from ray status", flush=True) + return 0 + + def _run_driver(self, driver_script: str, driver_args: list[str]) -> int: + """Execute the Python driver script with output capture. + + Runs the driver as a subprocess, streaming output to console in real-time + while capturing recent lines in a rolling buffer for error extraction. + The captured output is available via the ``driver_output`` property. + + Args: + driver_script: Path to the driver script + driver_args: Arguments for the driver + + Returns: + Exit code from the driver + """ + cmd = [self.driver_python, driver_script] + driver_args + logger.info(f"Running driver with python={self.driver_python}: {' '.join(cmd)}") + + # Build environment with extended PYTHONPATH if configured + env = os.environ.copy() + if self.driver_extra_pythonpath: + existing_pythonpath = env.get("PYTHONPATH", "") + if existing_pythonpath: + env["PYTHONPATH"] = f"{existing_pythonpath}{os.pathsep}{self.driver_extra_pythonpath}" + else: + env["PYTHONPATH"] = self.driver_extra_pythonpath + logger.info(f"Driver PYTHONPATH: {env['PYTHONPATH']}") + + # Reset the output buffer for this driver run + self._driver_output.clear() + + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + self._driver_process = process + + reader_thread = threading.Thread( + target=read_subprocess_output, + args=(process, self._driver_output), + daemon=True, + ) + reader_thread.start() + + try: + process.wait() + except BaseException: + # If interrupted (e.g. SystemExit from signal handler), terminate the + # driver process so it doesn't become orphaned. + self.terminate_driver() + raise + finally: + self._driver_process = None + + # Wait for reader thread to finish capturing remaining output + if reader_thread.is_alive(): + reader_thread.join(timeout=5) + + return process.returncode + + def _monitor_for_termination(self) -> int: + """Monitor for ENDED file and handle worker termination. + + Returns: + Exit code (0 for graceful termination) + """ + logger.info("Worker monitoring for termination signal") + + while not self._stop_event.is_set(): + if self.ended_file.exists(): + logger.info("Detected ENDED file, terminating worker...") + self._stop_ray() + return 0 + + _pause(1) + + return 0 + + def _signal_termination(self) -> None: + """Signal termination by creating the ENDED file.""" + logger.info(f"Creating termination signal: {self.ended_file}") + self.ended_file.touch() + + def _stop_ray(self, grace_period: int = 60, timeout: int | None = None) -> None: + """Stop Ray with grace period. + + Args: + grace_period: Seconds to wait for graceful shutdown. + timeout: Hard wall-clock bound for the ``ray stop`` subprocess. Defaults + to ``grace_period + 30``; callers under a tight cleanup budget pass a + smaller value so the (non-daemon) cleanup thread can't keep the + process alive past that budget. + """ + effective_timeout = timeout if timeout is not None else grace_period + 30 + logger.info(f"Stopping Ray with {grace_period}s grace period") + # Bound the call so a wedged `ray stop` can't keep the (non-daemon) + # cleanup thread — and therefore the whole process — alive indefinitely. + try: + subprocess.run( + [self.ray_executable, "stop", "--force", f"--grace-period={grace_period}"], + check=False, + capture_output=True, + timeout=effective_timeout, + ) + except subprocess.TimeoutExpired: + logger.warning(f"ray stop did not complete within {effective_timeout}s") + + def _cleanup_with_timeout(self, timeout: int = 30) -> None: + """Cleanup with timeout, force kill if necessary. + + Args: + timeout: Maximum seconds to wait for cleanup + """ + logger.info(f"[INFO] Cleaning up Ray cluster from RANK {self.rank}") + self._stop_event.set() + self._signal_termination() + + def cleanup() -> None: + # Keep stop + sleep within the outer `timeout` budget: bound `ray stop` + # to timeout-10 and sleep the remaining 10, so the non-daemon cleanup + # thread can't keep the process alive well past `timeout`. + self._stop_ray(grace_period=20, timeout=max(1, timeout - 10)) + _pause(10) # Wait for ray to stop + logger.info("[INFO] Cleanup complete.") + + cleanup_thread = threading.Thread(target=cleanup) + cleanup_thread.start() + cleanup_thread.join(timeout=timeout) + + if cleanup_thread.is_alive(): + logger.warning("[WARN] Cleanup timed out. Forcing termination.") + + def _log_ray_status(self) -> None: + """Log current Ray cluster status.""" + try: + status_result = subprocess.run( + [self.ray_executable, "status"], + capture_output=True, + text=True, + check=False, + timeout=RAY_STATUS_TIMEOUT_SECONDS, + ) + logger.info(f"Ray status:\n{status_result.stdout}") + + memory_result = subprocess.run( + [self.ray_executable, "memory"], + capture_output=True, + text=True, + check=False, + timeout=RAY_STATUS_TIMEOUT_SECONDS, + ) + logger.info(f"Ray memory:\n{memory_result.stdout}") + except Exception as e: + logger.warning(f"Failed to log Ray status: {e}") + + def cleanup(self) -> None: + """Public cleanup method.""" + self._cleanup_with_timeout() + + +def create_bootstrap_from_env() -> RayClusterBootstrap: + """Create RayClusterBootstrap from environment variables. + + Expected environment variables: + RANK: Node rank (0 for head, >0 for workers) + WORLD_SIZE: Total number of nodes + MASTER_ADDR: IP address of the head node + GPUS_PER_NODE: Number of GPUs per node (optional, default 1) + BASE_LOG_DIR: Base directory for logs (optional, default /tmp) + + Returns: + Configured RayClusterBootstrap instance + """ + return RayClusterBootstrap( + rank=int(os.getenv("RANK", "0")), + world_size=int(os.getenv("WORLD_SIZE", "1")), + master_addr=os.getenv("MASTER_ADDR", "127.0.0.1"), + driver_python=os.getenv("DRIVER_PYTHON", sys.executable), + driver_extra_pythonpath=os.getenv("DRIVER_EXTRA_PYTHONPATH", ""), + ) + + +def main() -> int: + """Main entry point for Ray cluster bootstrap. + + This can be invoked directly to start a Ray cluster node, or + with driver arguments to start the cluster and run a training script. + + Usage: + # Start cluster node (head or worker based on RANK) + python -m nmp.rl.tasks.training.backends.nemo_rl.ray_bootstrap + + # Start cluster and run driver + python -m nmp.rl.tasks.training.backends.nemo_rl.ray_bootstrap \ + driver_script.py --config config.yaml --id job-123 + + Returns: + Exit code (0 for success) + """ + import argparse + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + + parser = argparse.ArgumentParser(description="Ray cluster bootstrap") + parser.add_argument( + "driver_script", + nargs="?", + help="Optional driver script to run after cluster is ready", + ) + parser.add_argument( + "driver_args", + nargs="*", + help="Arguments to pass to the driver script", + ) + + args = parser.parse_args() + + bootstrap = create_bootstrap_from_env() + + # Setup signal handlers + def signal_handler(signum: int, frame: FrameType | None) -> None: + logger.warning(f"Received signal {signum}, initiating cleanup") + bootstrap.cleanup() + sys.exit(signum) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + if args.driver_script: + return bootstrap.run_with_driver(args.driver_script, args.driver_args) + else: + bootstrap.start() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/services/rl/src/nmp/rl/tasks/training/chat_templates.py b/services/rl/src/nmp/rl/tasks/training/chat_templates.py new file mode 100644 index 0000000000..4b13425dc0 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/chat_templates.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Chat template resolution and application for training backends. + +This module provides: +1. Priority-based chat template selection (resolve_chat_template) +2. Applying chat templates to output checkpoints (apply_chat_template_to_checkpoint) + +Chat template priority order: +1. User-provided template (via API) +2. Custom template from DEFAULT_CHAT_TEMPLATES map (enhanced for tool calling) +3. Model's built-in tokenizer template (fallback) + +The custom templates in the templates/ directory extend base model templates with: +- Tool calling support: , , formatting +- Generation markers: {% generation %}...{% endgeneration %} blocks for loss masking +- Enhanced compatibility across models +""" + +import json +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + +# Directory containing custom chat template jinja files +TEMPLATES_DIR = Path(__file__).parent / "templates" + +# ============================================================================ +# Model Name Constants +# ============================================================================ + +# Meta Llama models +META_LLAMA_31_8B_INSTRUCT = "meta/llama-3.1-8b-instruct" +META_LLAMA_31_70B_INSTRUCT = "meta/llama-3.1-70b-instruct" +META_LLAMA_31_405B_INSTRUCT = "meta/llama-3.1-405b-instruct" +META_LLAMA_32_1B = "meta/llama-3.2-1b" +META_LLAMA_32_1B_INSTRUCT = "meta/llama-3.2-1b-instruct" +META_LLAMA_32_3B_INSTRUCT = "meta/llama-3.2-3b-instruct" +META_LLAMA_33_70B_INSTRUCT = "meta/llama-3.3-70b-instruct" +# NVIDIA Nemotron models +NVIDIA_NEMOTRON_31_8B = "nvidia/nemotron-nano-llama-3.1-8b" +NVIDIA_NEMOTRON_31_70B = "nvidia/nemotron-llama-3.1-70b" +NVIDIA_NEMOTRON_33_49B = "nvidia/nemotron-super-llama-3.3-49b" +NVIDIA_NEMOTRON_33_49B_V1_5 = "nvidia/nemotron-super-llama-3.3-49b-v1.5" +# NIM model names (alternative naming) +NIM_NVIDIA_NEMOTRON_31_8B = "nvidia/llama-3.1-nemotron-nano-8b-v1" +NIM_NVIDIA_NEMOTRON_31_70B = "nvidia/llama-3.1-nemotron-70b-instruct" +NIM_NVIDIA_NEMOTRON_33_49B = "nvidia/llama-3.3-nemotron-super-49b-v1" +NIM_NVIDIA_NEMOTRON_33_49B_V1_5 = "nvidia/llama-3.3-nemotron-super-49b-v1.5" +# Microsoft models +PHI_4 = "microsoft/phi-4" + +# ============================================================================ +# Default Chat Templates Map +# ============================================================================ + +# Maps model names to custom jinja template filenames. +# These templates extend the base model templates with: +# - Tool calling support +# - Generation markers for loss masking +# - Enhanced compatibility +DEFAULT_CHAT_TEMPLATES: dict[str, str] = { + # Llama 3.1 family + META_LLAMA_31_8B_INSTRUCT: "llama-3.1-instruct.jinja", + META_LLAMA_31_70B_INSTRUCT: "llama-3.1-instruct.jinja", + META_LLAMA_31_405B_INSTRUCT: "llama-3.1-instruct.jinja", + # Llama 3.2 family + META_LLAMA_32_1B: "llama-3.2-instruct.jinja", + META_LLAMA_32_1B_INSTRUCT: "llama-3.2-instruct.jinja", + META_LLAMA_32_3B_INSTRUCT: "llama-3.2-instruct.jinja", + # Llama 3.3 family + META_LLAMA_33_70B_INSTRUCT: "llama-3.3-instruct.jinja", + # Nemotron family + NVIDIA_NEMOTRON_31_8B: "nemotron-3.1.jinja", + NVIDIA_NEMOTRON_31_70B: "nemotron-3.1.jinja", + NVIDIA_NEMOTRON_33_49B: "nemotron-super-3.3.jinja", + NVIDIA_NEMOTRON_33_49B_V1_5: "nemotron-super-3.3.jinja", + # NIM Nemotron (alternative naming) + NIM_NVIDIA_NEMOTRON_31_8B: "nemotron-3.1.jinja", + NIM_NVIDIA_NEMOTRON_31_70B: "nemotron-3.1.jinja", + NIM_NVIDIA_NEMOTRON_33_49B: "nemotron-super-3.3.jinja", + NIM_NVIDIA_NEMOTRON_33_49B_V1_5: "nemotron-super-3.3.jinja", + # Microsoft + PHI_4: "phi-4.jinja", +} + + +def _load_template_file(template_filename: str) -> str | None: + """Load a custom template from the templates directory.""" + template_path = TEMPLATES_DIR / template_filename + if template_path.exists(): + with open(template_path, "r", encoding="utf-8") as f: + return f.read() + logger.warning(f"Template file not found: {template_path}") + return None + + +def _get_tokenizer_chat_template(model_path: str, trust_remote_code: bool = False) -> str | None: + """ + Get chat template from model's tokenizer. + + Uses AutoTokenizer which handles all model formats (HF, NeMo, custom). + + Args: + model_path: Path to the model directory. + trust_remote_code: Whether to allow executing the model repo's custom code + when loading its tokenizer. Sourced from the model entity's + ``trust_remote_code`` flag (operator/registry-controlled, threaded down + from the API) — not free-form end-user input — because some models + require it to load their tokenizer at all. + """ + try: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust_remote_code) + template = getattr(tokenizer, "chat_template", None) + if template: + logger.debug(f"Found chat template in tokenizer for {model_path}") + return template + except Exception as e: + logger.warning(f"Could not load tokenizer to get chat template: {e}") + return None + + +def resolve_chat_template( + model_path: str, + model_name: str | None = None, + user_template: str | None = None, + trust_remote_code: bool = False, +) -> str | None: + """ + Resolve chat template using priority-based selection. + + Priority order: + 1. User-provided template (highest priority) + 2. Custom template from DEFAULT_CHAT_TEMPLATES (if model_name matches) + 3. Model's built-in tokenizer template (fallback) + + Args: + model_path: Path to the model directory (for tokenizer fallback). + model_name: Canonical model name (e.g., "meta/llama-3.1-8b-instruct"). + Used to look up custom templates. + user_template: User-provided template string (takes highest priority). + trust_remote_code: Forwarded to the tokenizer fallback. Sourced from the + model entity's ``trust_remote_code`` flag; some models require it to + load their tokenizer. + + Returns: + The resolved chat template string, or None if no template found. + """ + # Priority 1: User-provided template + if user_template: + logger.info("Using user-provided chat template") + return user_template + + # Priority 2: Custom template from DEFAULT_CHAT_TEMPLATES + if model_name and model_name in DEFAULT_CHAT_TEMPLATES: + template_filename = DEFAULT_CHAT_TEMPLATES[model_name] + template = _load_template_file(template_filename) + if template: + logger.info(f"Using custom chat template for {model_name}: {template_filename}") + return template + + # Priority 3: Model's built-in tokenizer template + template = _get_tokenizer_chat_template(model_path, trust_remote_code=trust_remote_code) + if template: + logger.info(f"Using model's built-in chat template from {model_path}") + return template + + logger.warning(f"No chat template found for model_name={model_name}, model_path={model_path}") + return None + + +def apply_chat_template_to_checkpoint( + output_path: Path, + chat_template: str | None, +) -> None: + """ + Apply chat template to the output checkpoint's tokenizer_config.json. + + Also ensures pad_token is set if missing (uses eos_token as fallback), + which is required by many inference frameworks. + + Args: + output_path: Path to the checkpoint directory containing tokenizer_config.json. + chat_template: The chat template string to apply. If None, skips application. + """ + if not chat_template: + logger.warning("No chat template provided, skipping") + return + + tokenizer_config = output_path / "tokenizer_config.json" + if not tokenizer_config.exists(): + logger.warning(f"tokenizer_config.json not found at {output_path}") + return + + with open(tokenizer_config, "r") as f: + config = json.load(f) + + config["chat_template"] = chat_template + + # Backfill pad_token from eos_token when missing. Many inference frameworks + # require a pad_token, and full-weight checkpoints sometimes ship without one. + if not config.get("pad_token") and config.get("eos_token"): + config["pad_token"] = config["eos_token"] + logger.info("Backfilled missing pad_token from eos_token") + + with open(tokenizer_config, "w") as f: + json.dump(config, f, indent=2) + + logger.info("Applied chat template to output checkpoint") diff --git a/services/rl/src/nmp/rl/tasks/training/datasets/preparation.py b/services/rl/src/nmp/rl/tasks/training/datasets/preparation.py new file mode 100644 index 0000000000..ec08deb3f1 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/datasets/preparation.py @@ -0,0 +1,559 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +""" +Shared dataset utilities for training backends. + +This module provides schema detection, sample counting, and training schedule +utilities that are shared across all training backends (automodel, megatron_bridge, nemo_rl). +""" + +import json +import logging +import random +import re +import shutil +import subprocess +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any, Optional, Union + +from nmp.rl.app.constants import DEFAULT_SEED + +logger = logging.getLogger(__name__) + +# Dataset directory constants for merged files (we control this structure) +MERGED_DIR = "merged" +TRAIN_FILE = "train.jsonl" +VAL_FILE = "validation.jsonl" + +# Heuristic patterns for discovering training files. JSONL only: everything +# downstream (schema detection, _merge_files, _create_val_split) is line-delimited, +# so a plain .json array/object would produce invalid merged data or split failures. +TRAIN_PATTERNS = [ + "train*.jsonl", + "training*.jsonl", +] +TRAIN_DIRS = ["train", "training"] + +# Heuristic patterns for discovering validation files (JSONL only — see above). +VAL_PATTERNS = [ + "val*.jsonl", + "validation*.jsonl", + "dev*.jsonl", +] +VAL_DIRS = ["val", "validation", "dev"] + + +class DatasetSchema(str, Enum): + """Detected dataset schema type.""" + + CHAT = "chat" # OpenAI messages format: {"messages": [...]} + SFT = "sft" # Prompt/completion: {"prompt": ..., "completion": ...} + CUSTOM = "custom" # Custom columns via prompt_template + EMBEDDING = "embedding" # Retrieval format: {"query": ..., "pos_doc": ..., "neg_doc": [...]} + + +class DatasetFormatError(Exception): + """Raised when dataset format is invalid or unsupported.""" + + pass + + +def detect_dataset_schema( + file_path: Path, + prompt_template: str | None = None, +) -> tuple[DatasetSchema, tuple[str, ...] | None]: + """ + Detect dataset schema by sampling the first line. + + Supports four formats: + 1. Chat format: {"messages": [{"role": "user", ...}, {"role": "assistant", ...}]} + 2. Embedding format: {"query": "...", "pos_doc": "...", "neg_doc": ["...", "..."]} + 3. SFT format: {"prompt": "...", "completion": "..."} + 4. Custom format: Any two-column format specified via prompt_template like "{input} {output}" + + Args: + file_path: Path to the JSONL dataset file. + prompt_template: Optional template string with two placeholders like "{input} {output}". + + Returns: + Tuple of (schema_type, column_keys) where: + - CHAT: column_keys is None + - EMBEDDING: column_keys is ("query", "pos_doc", "neg_doc") + - SFT/CUSTOM: column_keys is (question_col, answer_col) + + Raises: + DatasetFormatError: If the dataset format cannot be detected or is invalid. + """ + with open(file_path, "r", encoding="utf-8") as f: + line = f.readline() + + try: + obj: dict[str, Any] = json.loads(line) + except json.JSONDecodeError as e: + raise DatasetFormatError(f"Invalid JSON in {file_path}: {e}") + + # Check for chat format (OpenAI messages) + if "messages" in obj and isinstance(obj["messages"], list): + if len(obj["messages"]) > 0 and isinstance(obj["messages"][0], dict): + if "role" in obj["messages"][0]: + logger.info(f"Detected chat dataset format in {file_path}") + return DatasetSchema.CHAT, None + + # Check for embedding/retrieval format + # Format: {"query": "...", "pos_doc": "...", "neg_doc": ["...", "..."]} + if "query" in obj and "pos_doc" in obj and "neg_doc" in obj: + if isinstance(obj["query"], str) and isinstance(obj["pos_doc"], str) and isinstance(obj["neg_doc"], list): + logger.info(f"Detected embedding/retrieval dataset format in {file_path}") + return DatasetSchema.EMBEDDING, ("query", "pos_doc", "neg_doc") + + # Check for custom prompt_template format + if prompt_template: + keys = re.findall(r"\{(.*?)\}", prompt_template) + if len(keys) == 2: + # Validate keys exist in data + if all(k in obj for k in keys): + logger.info(f"Detected custom template format with keys {keys}") + return DatasetSchema.CUSTOM, (keys[0], keys[1]) + else: + raise DatasetFormatError( + f"prompt_template keys {keys} not found in dataset. Available keys: {list(obj.keys())}" + ) + else: + raise DatasetFormatError(f"prompt_template must have exactly 2 placeholders, got: {prompt_template}") + + # Check for standard SFT format (prompt/completion) + if "prompt" in obj and "completion" in obj: + logger.info(f"Detected SFT (prompt/completion) format in {file_path}") + return DatasetSchema.SFT, ("prompt", "completion") + + # Fallback - try to find any two string columns + string_cols = [k for k, v in obj.items() if isinstance(v, str)] + if len(string_cols) >= 2: + logger.warning(f"Could not detect standard format, using first two string columns: {string_cols[:2]}") + return DatasetSchema.SFT, (string_cols[0], string_cols[1]) + + raise DatasetFormatError( + f"Could not detect dataset format. Expected 'messages' (chat) or " + f"'prompt'/'completion' (SFT) columns. Found: {list(obj.keys())}" + ) + + +def _count_jsonl_samples_python(file_path: Path) -> int: + """Pure Python implementation of line counting (fallback).""" + count = 0 + with open(file_path, "r", encoding="utf-8") as f: + for line in f: + if line.strip(): # Non-empty line + count += 1 + return count + + +def count_jsonl_samples(file_path: Path) -> int: + """ + Count the number of non-empty lines in a JSONL file. + + Uses grep for efficiency with large files when available, + falls back to pure Python implementation otherwise. + + Args: + file_path: Path to the JSONL file. + + Returns: + Number of non-empty lines (samples) in the file. + """ + # Check if grep is available + if shutil.which("grep") is None: + return _count_jsonl_samples_python(file_path) + + try: + # Use `grep -c "\S"` to count non-empty lines (excludes trailing empty lines) + result = subprocess.check_output(["grep", "-c", r"\S", str(file_path)], text=True) + return int(result.strip()) + except subprocess.CalledProcessError as e: + # grep exits 1 when there are no matching (non-empty) lines → genuinely empty. + # Any other exit code (e.g. 2 = file unreadable / regex error) is a real + # failure, so fall back to the Python counter rather than reporting 0. + if e.returncode == 1: + return 0 + return _count_jsonl_samples_python(file_path) + except OSError: + # Fallback if subprocess fails for any reason + return _count_jsonl_samples_python(file_path) + + +def compute_val_check_interval( + steps_per_epoch: int, + max_steps: int, + val_check_interval: Optional[Union[int, float]] = None, +) -> int: + """ + Compute how often to run validation (in steps). + + This handles the semantic difference between: + - float <= 1.0: Fraction of epoch (e.g., 0.5 = validate at 50% of each epoch) + - int or float > 1.0: Absolute step count + + Args: + steps_per_epoch: Number of gradient steps per epoch. + max_steps: Maximum training steps. + val_check_interval: User-provided interval (float for fraction, int for steps). + + Returns: + Integer step count for validation interval. + + Raises: + ValueError: If val_check_interval is negative. + """ + effective_steps = min(steps_per_epoch, max_steps) + + if val_check_interval is None or val_check_interval == 0: + # Default: validate once per epoch (or at end if max_steps < steps_per_epoch) + return effective_steps + + if val_check_interval < 0: + raise ValueError("val_check_interval cannot be negative") + + # Float <= 1.0: interpret as fraction of epoch + if isinstance(val_check_interval, float) and val_check_interval <= 1.0: + interval = max(1, int(val_check_interval * steps_per_epoch)) + else: + # Integer or float > 1.0: treat as absolute step count + interval = int(val_check_interval) + + # Cap at effective_steps + interval = min(interval, effective_steps) + + # Ensure validation happens at least once before training ends + if interval >= max_steps: + interval = max(1, max_steps - 1) + + return interval + + +@dataclass +class PreparedDataset: + """Result of dataset preparation.""" + + merged_dir: Path + train_file: Path + validation_file: Path + train_samples: int + validation_samples: int + + +def _discover_files_by_patterns(base_path: Path, patterns: list[str], dirs: list[str]) -> list[Path]: + """ + Discover files matching patterns or in specific directories. + + Searches for: + 1. Files matching glob patterns in base_path + 2. All .jsonl/.json files in specified subdirectories + + Args: + base_path: Root directory to search. + patterns: Glob patterns to match (e.g., ["train*.jsonl"]). + dirs: Subdirectory names to search (e.g., ["train", "training"]). + + Returns: + Sorted list of discovered file paths. + """ + files: set[Path] = set() + + # Pattern matching in base directory + for pattern in patterns: + for match in base_path.glob(pattern): + if match.is_file(): + files.add(match.resolve()) + + # Files in subdirectories + for dir_name in dirs: + subdir = base_path / dir_name + if subdir.is_dir(): + for f in subdir.iterdir(): + if f.is_file() and f.suffix.lower() in (".jsonl", ".json"): + files.add(f.resolve()) + + return sorted(files) # Sorted for deterministic ordering + + +def discover_dataset_files(dataset_path: Path) -> tuple[list[Path], list[Path]]: + """ + Discover training and validation files using heuristics. + + Heuristics applied (in order): + 1. Files matching train*/training* patterns → training + 2. Files in train/ or training/ directories → training + 3. Files matching val*/validation*/dev* patterns → validation + 4. Files in val/, validation/, or dev/ directories → validation + 5. If only one .jsonl file found → treat as training (will auto-split) + + Args: + dataset_path: Path to the dataset directory. + + Returns: + Tuple of (training_files, validation_files). + + Raises: + DatasetFormatError: If no training files can be found. + """ + dataset_path = Path(dataset_path).resolve() + + if not dataset_path.exists(): + raise DatasetFormatError(f"Dataset path does not exist: {dataset_path}") + + # If path is a file, treat it as the training file + if dataset_path.is_file(): + logger.info(f"Dataset path is a file, treating as training data: {dataset_path}") + return [dataset_path], [] + + # Discover training files + train_files = _discover_files_by_patterns(dataset_path, TRAIN_PATTERNS, TRAIN_DIRS) + + # Discover validation files + val_files = _discover_files_by_patterns(dataset_path, VAL_PATTERNS, VAL_DIRS) + + # Fallback: if no files found with patterns, check for any .jsonl files + if not train_files and not val_files: + all_jsonl = sorted(f for f in dataset_path.glob("*.jsonl") if f.is_file()) + if len(all_jsonl) == 1: + logger.info(f"Found single JSONL file, treating as training data: {all_jsonl[0]}") + train_files = all_jsonl + elif len(all_jsonl) > 1: + # Ambiguous - could be train/val or multiple training files + logger.warning( + f"Found {len(all_jsonl)} JSONL files without clear train/val naming. " + f"Treating all as training data: {[f.name for f in all_jsonl]}" + ) + train_files = all_jsonl + + if not train_files: + raise DatasetFormatError( + f"No training files found in {dataset_path}. " + f"Expected files matching patterns like train*.jsonl or a train/ directory." + ) + + logger.info(f"Discovered {len(train_files)} training file(s): {[f.name for f in train_files]}") + if val_files: + logger.info(f"Discovered {len(val_files)} validation file(s): {[f.name for f in val_files]}") + else: + logger.info("No validation files found - will auto-split from training data") + + return train_files, val_files + + +def _merge_files(files: list[Path], output_file: Path) -> int: + """ + Merge multiple JSONL files into a single file. + + Args: + files: List of files to merge. + output_file: Output file path. + + Returns: + Total number of samples (non-empty lines) in merged file. + """ + output_file.parent.mkdir(parents=True, exist_ok=True) + + # Stream each file in fixed-size chunks rather than reading it whole: dataset + # shards can be multiple GB and inp.read() would pull an entire shard into + # memory. We track whether the last chunk ended in a newline so a file that + # lacks a trailing newline doesn't get concatenated onto the next one. + chunk_size = 1024 * 1024 # 1 MiB + with open(output_file, "w", encoding="utf-8") as out: + for f in files: + ended_with_newline = True # empty file → no separator needed + with open(f, "r", encoding="utf-8") as inp: + for chunk in iter(lambda inp=inp: inp.read(chunk_size), ""): + out.write(chunk) + ended_with_newline = chunk.endswith("\n") + # Ensure a newline separates this file's last record from the next file. + if not ended_with_newline: + out.write("\n") + + return count_jsonl_samples(output_file) + + +def _create_val_split( + train_file: Path, + output_train: Path, + output_val: Path, + val_ratio: float = 0.1, + seed: int = DEFAULT_SEED, +) -> tuple[int, int]: + """ + Split a training file into train and validation sets. + + Args: + train_file: Source training file. + output_train: Output path for training split. + output_val: Output path for validation split. + val_ratio: Fraction of data to use for validation (default: 10%). + seed: Random seed for reproducible splits (default: 1111). + + Returns: + Tuple of (train_samples, validation_samples). + """ + # First pass: count non-empty rows without materializing the file (the merged + # input can be multi-GB, so we never read it all into memory). + total = count_jsonl_samples(train_file) + + # A split needs at least one sample on each side. Fail fast for files too small + # to split rather than emitting an empty train (or validation) set that crashes + # downstream config compilation. + if total < 2: + raise DatasetFormatError( + f"Training file {train_file} has {total} sample(s); at least 2 are required to " + "create a train/validation split. Provide more training data or supply a separate " + "validation file." + ) + + # Clamp to total - 1 so at least one training sample always remains, even when + # the requested ratio would otherwise consume the whole (tiny) dataset. + val_size = max(1, int(total * val_ratio)) + val_size = min(val_size, total - 1) + + # Pick which row indices go to validation. A local RNG keeps the split + # deterministic (important for multi-node) without mutating the process-wide + # random state that later training code may rely on. Holding only the chosen + # indices (not the rows) keeps memory bounded. + rng = random.Random(seed) + val_indices = set(rng.sample(range(total), val_size)) + + output_train.parent.mkdir(parents=True, exist_ok=True) + output_val.parent.mkdir(parents=True, exist_ok=True) + + # The auto-split-from-merged path passes the same path for both train_file and + # output_train. Opening output_train for write would truncate the file we still + # need to stream from `src`, yielding empty splits and destroying the merged + # train file. When they collide, write the train split to a temp file and + # atomically replace at the end. + rewrite_train_in_place = train_file.resolve() == output_train.resolve() + train_write_path = output_train.with_name(f"{output_train.name}.tmp") if rewrite_train_in_place else output_train + + # Second pass: stream each non-empty row to train/val by index — re-serializing + # to normalize JSON — without ever holding the whole dataset in memory. + train_count = 0 + val_count = 0 + with ( + open(train_file, "r", encoding="utf-8") as src, + open(train_write_path, "w", encoding="utf-8") as train_out, + open(output_val, "w", encoding="utf-8") as val_out, + ): + idx = 0 + for raw in src: + if not raw.strip(): + continue + normalized = json.dumps(json.loads(raw)) + "\n" + if idx in val_indices: + val_out.write(normalized) + val_count += 1 + else: + train_out.write(normalized) + train_count += 1 + idx += 1 + + if rewrite_train_in_place: + train_write_path.replace(output_train) + + logger.info( + f"Created validation split: {train_count} train samples, {val_count} val samples ({val_ratio:.0%} split)" + ) + + return train_count, val_count + + +def prepare_dataset( + dataset_path: Path, + output_dir: Optional[Path] = None, + val_split_ratio: float = 0.1, + seed: int = DEFAULT_SEED, +) -> PreparedDataset: + """ + Prepare dataset for training by discovering, merging, and optionally splitting files. + + This function: + 1. Discovers training and validation files using heuristics + 2. Merges multiple files into single train.jsonl and val.jsonl + 3. Auto-creates validation split if no validation files found + 4. Returns paths to the prepared files + + Args: + dataset_path: Path to the dataset directory or file. + output_dir: Directory for merged output (default: dataset_path/merged). + val_split_ratio: Fraction for auto-split if no validation data (default: 0.1). + seed: Random seed for reproducible validation splits (default: 1111). + + Returns: + PreparedDataset with paths to merged files and sample counts. + + Raises: + DatasetFormatError: If dataset cannot be prepared. + """ + dataset_path = Path(dataset_path).resolve() + + # Determine output directory + if output_dir is None: + if dataset_path.is_file(): + merged_dir = dataset_path.parent / MERGED_DIR + else: + merged_dir = dataset_path / MERGED_DIR + else: + merged_dir = Path(output_dir).resolve() + + train_output = merged_dir / TRAIN_FILE + validation_output = merged_dir / VAL_FILE + + # Discover files + train_files, val_files = discover_dataset_files(dataset_path) + + # Merge training files + if len(train_files) == 1 and not val_files: + # Single file, no validation - need to split + logger.info("Single training file with no validation data - creating split") + train_samples, validation_samples = _create_val_split( + train_files[0], + train_output, + validation_output, + val_ratio=val_split_ratio, + seed=seed, + ) + else: + # Merge training files + train_samples = _merge_files(train_files, train_output) + logger.info(f"Merged {len(train_files)} training file(s) → {train_output} ({train_samples} samples)") + + if val_files: + # Merge validation files + validation_samples = _merge_files(val_files, validation_output) + logger.info( + f"Merged {len(val_files)} validation file(s) → {validation_output} ({validation_samples} samples)" + ) + else: + # Auto-split from merged training file + logger.info("No validation files - creating split from merged training data") + # Read merged, split, re-write + train_samples, validation_samples = _create_val_split( + train_output, + train_output, + validation_output, + val_ratio=val_split_ratio, + seed=seed, + ) + + return PreparedDataset( + merged_dir=merged_dir, + train_file=train_output, + validation_file=validation_output, + train_samples=train_samples, + validation_samples=validation_samples, + ) diff --git a/services/rl/src/nmp/rl/tasks/training/datasets/schemas.py b/services/rl/src/nmp/rl/tasks/training/datasets/schemas.py new file mode 100644 index 0000000000..329dd59195 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/datasets/schemas.py @@ -0,0 +1,436 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +# ============================================================================= +# Dataset Schemas for DPO Training +# ============================================================================= +# Preference Dataset Schemas for DPO Training: +# - PreferenceDataset: Native format with context + ranked completions +# - BinaryPreferenceDataset: Simple prompt/chosen/rejected strings +# - HelpSteer3Dataset: NVIDIA HelpSteer3 format with preference scores +# - Tulu3PreferenceDataset: AllenAI Tulu3 format with message lists +# +# SFT Dataset Schemas: +# - SFTDatasetItemSchema: Standard prompt/completion format +from typing import Annotated, Any, List, Literal, Optional, Union + +from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, model_validator + +# Dataset class names from nmp.rl.tasks.training.backends.nemo_rl.preference_datasets +# These constants ensure consistency between the discriminator and Tag values +PREFERENCE_DATASET = "PreferenceDataset" +BINARY_PREFERENCE_DATASET = "BinaryPreferenceDataset" +HELPSTEER3_DATASET = "HelpSteer3" +TULU3_PREFERENCE_DATASET = "Tulu3Preference" + + +class ChatMessage(BaseModel): + """A single message in a conversation.""" + + role: str = Field(..., description="The role of the message sender (e.g., 'user', 'assistant', 'system')") + content: str = Field(..., description="The content of the message") + + +class CompletionItem(BaseModel): + """A ranked completion in a preference dataset.""" + + rank: int = Field(..., description="Rank of this completion (0 = best/chosen, higher = worse)") + completion: List[ChatMessage] = Field(..., description="The completion as a list of messages") + + +class PreferenceDatasetItemSchema(BaseModel): + """Schema for native PreferenceDataset format. + + This is the canonical format used by nemo-rl's PreferenceDataset class. + It supports multi-turn context and multiple ranked completions. + + Example: + { + "context": [{"role": "user", "content": "What is 2+2?"}], + "completions": [ + {"rank": 0, "completion": [{"role": "assistant", "content": "4"}]}, + {"rank": 1, "completion": [{"role": "assistant", "content": "5"}]} + ] + } + """ + + context: List[ChatMessage] = Field( + ..., description="The conversation context (prompt messages including previous turns)" + ) + completions: List[CompletionItem] = Field( + ..., description="List of ranked completions (rank 0 = preferred, rank 1 = rejected, etc.)" + ) + + model_config = ConfigDict(extra="allow") + + +class BinaryPreferenceDatasetItemSchema(BaseModel): + """Schema for BinaryPreferenceDataset format. + + Simple format with prompt, chosen response, and rejected response as strings. + The prompt can be either a string or a list of messages. + + Example: + { + "prompt": "What is the capital of France?", + "chosen": "The capital of France is Paris.", + "rejected": "The capital of France is London." + } + """ + + prompt: Union[str, List[ChatMessage]] = Field(..., description="The input prompt (string or list of messages)") + chosen: str = Field(..., description="The preferred/chosen response") + rejected: str = Field(..., description="The rejected/non-preferred response") + + model_config = ConfigDict(extra="allow") + + +class HelpSteer3DatasetItemSchema(BaseModel): + """Schema for NVIDIA HelpSteer3 preference dataset format. + + Uses numeric preference scores to indicate which response is preferred. + - Negative overall_preference: response1 is preferred + - Positive overall_preference: response2 is preferred + - Zero overall_preference: tie (no preference) + + Example: + { + "context": "Explain quantum computing", + "response1": "Quantum computing uses qubits...", + "response2": "Quantum computing is magic...", + "overall_preference": -2 + } + """ + + context: Union[str, List[ChatMessage]] = Field(..., description="The input context (string or list of messages)") + response1: str = Field(..., description="First response option") + response2: str = Field(..., description="Second response option") + overall_preference: int = Field( + ..., + description="Preference score: negative=response1 preferred, positive=response2 preferred, 0=tie", + ) + + model_config = ConfigDict(extra="allow") + + +class Tulu3PreferenceDatasetItemSchema(BaseModel): + """Schema for AllenAI Tulu3 preference dataset format. + + Contains full conversation histories for both chosen and rejected responses. + The last message in each list must be from the assistant role. + + Example: + { + "chosen": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi! How can I help?"} + ], + "rejected": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Go away."} + ] + } + """ + + chosen: List[ChatMessage] = Field( + ..., description="Full conversation with preferred response (last message must be assistant)" + ) + rejected: List[ChatMessage] = Field( + ..., description="Full conversation with rejected response (last message must be assistant)" + ) + + model_config = ConfigDict(extra="allow") + + +def get_preference_dataset_discriminator(v: Any) -> str: + """Determine the preference dataset schema type based on field presence. + + This discriminator function examines the fields present in the data + to determine which schema type it matches. Returns the NeMo RL dataset + class name that corresponds to the detected format: + - PreferenceDataset: Has 'context' and 'completions' fields (native format) + - HelpSteer3: Has 'overall_preference' field (HelpSteer3 format) + - Tulu3PreferenceDataset: Has 'chosen' and 'rejected' as lists of messages + - BinaryPreferenceDataset: Has 'prompt', 'chosen', 'rejected' + + Args: + v: The data to discriminate (dict or model instance) + + Returns: + NeMo RL dataset class name identifying the schema type + """ + if isinstance(v, dict): + # Native PreferenceDataset format: context + completions + if "completions" in v and "context" in v: + return PREFERENCE_DATASET + + # HelpSteer3 format: has overall_preference score + if "overall_preference" in v: + return HELPSTEER3_DATASET + + # Tulu3 format: chosen/rejected are lists of messages (must check BEFORE BinaryPreferenceDataset) + # Tulu3 data may also have 'prompt' field, so we differentiate by checking if chosen/rejected are lists + if "chosen" in v and "rejected" in v: + chosen = v.get("chosen") + if isinstance(chosen, list) and len(chosen) > 0: + # Check if it looks like a message list + if isinstance(chosen[0], dict) and "role" in chosen[0]: + return TULU3_PREFERENCE_DATASET + + # BinaryPreferenceDataset format: prompt + chosen + rejected (as strings) + if "prompt" in v and "chosen" in v and "rejected" in v: + return BINARY_PREFERENCE_DATASET + + return PREFERENCE_DATASET # Default fallback + + +# Union type for all preference dataset formats +DPOPreferenceDatasetSchemaType = Annotated[ + Union[ + Annotated[PreferenceDatasetItemSchema, Tag(PREFERENCE_DATASET)], + Annotated[BinaryPreferenceDatasetItemSchema, Tag(BINARY_PREFERENCE_DATASET)], + Annotated[HelpSteer3DatasetItemSchema, Tag(HELPSTEER3_DATASET)], + Annotated[Tulu3PreferenceDatasetItemSchema, Tag(TULU3_PREFERENCE_DATASET)], + ], + Discriminator(get_preference_dataset_discriminator), +] + + +# ============================================================================= +# SFT Dataset Schemas +# ============================================================================= +class SFTPromptTemplateDatasetItemSchema(BaseModel): + """Schema for standard SFT (Supervised Fine-Tuning) dataset format. + + The standard format has prompt and completion fields, but allows additional + fields for custom templates (e.g., {input}, {output}, {instruction}, etc.). + + Example (standard format): + { + "prompt": "What is the capital of France?", + "completion": "The capital of France is Paris." + } + + Example (custom template format): + { + "instruction": "Answer the question", + "input": "What is the capital of France?", + "output": "The capital of France is Paris." + } + """ + + model_config = ConfigDict(extra="allow") + + # Make all fields optional so custom templates can use any field names + prompt: Optional[str] = Field(None, description="The input prompt (standard format)") + completion: Optional[str] = Field(None, description="The expected completion/output (standard format)") + + +class FunctionCallDetails(BaseModel): + """Details of a function call made by a tool call. + + Example: + { + "name": "get_weather", + "arguments": {"location": "San Francisco"} + } + """ + + name: str = Field(..., description="The name of the function to call") + arguments: dict[str, Any] = Field(..., description="The arguments to pass to the function") + content_type: Optional[str] = Field(None, description="Optional content type of the function response") + + +class ToolCall(BaseModel): + """A tool call in a message.""" + + type: Literal["function"] = Field(..., description="The type of tool call (must be 'function')") + function: FunctionCallDetails = Field(..., description="Function call details including name and arguments") + + +class SFTChatMessage(BaseModel): + """A single message in an SFT chat conversation. + + Each message must have a role and at least one of: content, thinking, or tool_calls. + + Important: content and thinking are mutually exclusive within a single message. + If both are needed, they should be in separate messages (e.g., one message with + thinking followed by another message with content). + """ + + role: str = Field(..., description="The role of the message sender (e.g., 'user', 'assistant', 'system')") + content: str | None = Field(None, description="The content of the message") + thinking: str | None = Field(None, description="Thinking/reasoning content") + # min_length=1: an empty tool_calls list is non-None but renders a tool-call + # block the templates never close, producing malformed prompts. A message with + # no tool calls must use tool_calls=None, not []. + tool_calls: list[ToolCall] | None = Field(None, min_length=1, description="Tool calls made in this message") + + @staticmethod + def _schema_extra(schema: dict[str, Any]) -> None: + """Add anyOf constraint requiring at least one of content, thinking, or tool_calls.""" + schema["anyOf"] = [ + { + "required": ["content"], + "properties": {"content": {"type": "string"}}, + "not": {"required": ["thinking"]}, + }, + { + "required": ["thinking"], + "properties": {"thinking": {"type": "string"}}, + "not": {"required": ["content"]}, + }, + {"required": ["tool_calls"], "properties": {"tool_calls": {"minItems": 1}}}, + ] + + model_config = ConfigDict(extra="forbid", json_schema_extra=_schema_extra) + + @model_validator(mode="after") + def check_has_content_or_thinking_or_tool_calls(self) -> "SFTChatMessage": + """Validate that message has at least one of content, thinking, or tool_calls. + + Also enforces that content and thinking are mutually exclusive - they cannot + both be present in the same message. + """ + if self.content is None and self.thinking is None and self.tool_calls is None: + raise ValueError("Message must have at least one of: content, thinking, or tool_calls") + + if self.content is not None and self.thinking is not None: + raise ValueError("Message cannot have both content and thinking - they are mutually exclusive") + + return self + + +class FunctionParameters(BaseModel): + """Parameters schema for a function definition. + + Example: + { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + } + } + """ + + type: Literal["object"] = Field(..., description="The type of parameters (must be 'object')") + properties: dict[str, Any] = Field(..., description="The properties/arguments the function accepts") + + +class FunctionDefinitionDetails(BaseModel): + """Details of a function definition for tool calling. + + Example: + { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": {"type": "object", "properties": {...}}, + "required": ["location"] + } + """ + + name: str = Field(..., description="The name of the function") + description: str = Field(..., description="A description of what the function does") + parameters: FunctionParameters = Field(..., description="The parameters schema for the function") + required: list[str] | None = Field(None, description="List of required parameter names") + + +class ToolDefinition(BaseModel): + """A tool definition for function calling.""" + + type: Literal["function"] = Field(..., description="The type of tool (must be 'function')") + function: FunctionDefinitionDetails = Field( + ..., description="Function definition with name, description, and parameters" + ) + + +class SFTChatDatasetItemSchema(BaseModel): + """Schema for SFT chat format based on MESSAGES_SCHEMA. + + This format represents conversations with message lists and optional tool definitions. + + Example: + { + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"} + ], + "tools": [...] # optional + } + """ + + # min_length=1: templates index messages[0] immediately, so an empty list + # would pass validation and then crash during rendering. + messages: list[SFTChatMessage] = Field(..., min_length=1, description="List of messages in the conversation") + tools: list[ToolDefinition] | None = Field( + None, description="Optional tool definitions available in the conversation" + ) + + model_config = ConfigDict(extra="allow") + + +# Embedding Dataset Schemas +class EmbeddingDatasetItemSchema(BaseModel): + """Schema for embedding dataset format. + + Example: + { + "query": "What is machine learning?", + "pos_doc": "Machine learning is a branch of AI...", + "neg_doc": ["Deep learning is...", "Neural networks are..."] + } + """ + + query: str = Field(..., description="The query text") + pos_doc: str = Field(..., description="The positive document") + neg_doc: list[str] = Field(..., description="List of negative documents") + + model_config = ConfigDict(extra="allow") + + +def get_sft_dataset_discriminator(v: Any) -> str: + """Determine the SFT dataset schema type based on field presence. + + This discriminator examines the fields to determine format: + - "EmbeddingDatasetItemSchema": Has 'query', 'pos_doc', 'neg_doc' fields (embedding format) + - "SFTChatDatasetItemSchema": Has 'messages' field (chat format) + - "SFTPromptTemplateDatasetItemSchema": Has other fields (prompt template format) + + Args: + v: The data to discriminate (dict or model instance) + + Returns: + Schema type name identifying the format + """ + if isinstance(v, dict): + # Embedding format: has query, pos_doc, neg_doc fields + if "query" in v and "pos_doc" in v and "neg_doc" in v: + return "EmbeddingDatasetItemSchema" + + # Chat format: has messages array. The returned tag must match the + # SFTChatDatasetItemSchema Tag on the union below. + if "messages" in v: + return "SFTChatDatasetItemSchema" + + # Prompt template format: has prompt/completion or custom fields + return "SFTPromptTemplateDatasetItemSchema" + + return "SFTPromptTemplateDatasetItemSchema" # Default fallback + + +# Union type for all SFT dataset formats +SFTDatasetSchemaType = Annotated[ + Union[ + Annotated[SFTPromptTemplateDatasetItemSchema, Tag(str(SFTPromptTemplateDatasetItemSchema.__name__))], + Annotated[SFTChatDatasetItemSchema, Tag(str(SFTChatDatasetItemSchema.__name__))], + Annotated[EmbeddingDatasetItemSchema, Tag(str(EmbeddingDatasetItemSchema.__name__))], + ], + Discriminator(get_sft_dataset_discriminator), +] diff --git a/services/rl/src/nmp/rl/tasks/training/datasets/validation.py b/services/rl/src/nmp/rl/tasks/training/datasets/validation.py new file mode 100644 index 0000000000..00fe79e2ad --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/datasets/validation.py @@ -0,0 +1,382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import json +import logging +import os +import re +from pathlib import Path +from typing import Any, Callable, Optional + +import jsonschema +from jsonschema import exceptions +from nmp.rl.entities.values import FinetuningType, TrainingType +from nmp.rl.tasks.training.datasets.preparation import DatasetFormatError +from nmp.rl.tasks.training.datasets.schemas import ( + DPOPreferenceDatasetSchemaType, + SFTDatasetSchemaType, + get_preference_dataset_discriminator, +) + +logger = logging.getLogger(__name__) + + +def DPO_SCHEMA(_: str | None = None) -> dict: + """Generate JSON schema for DPO preference datasets. + + Uses the DPOPreferenceDatasetSchemaType union which supports: + - PreferenceDataset: Native format with context + ranked completions + - BinaryPreferenceDataset: Simple prompt/chosen/rejected strings + - HelpSteer3Dataset: NVIDIA HelpSteer3 format with preference scores + - Tulu3PreferenceDataset: AllenAI Tulu3 format with message lists + """ + from pydantic import TypeAdapter + + # Create TypeAdapter for the DPO union type to generate JSON schema + adapter = TypeAdapter(DPOPreferenceDatasetSchemaType) + schema = adapter.json_schema() + + # Add JSON schema metadata + schema["$schema"] = "https://json-schema.org/draft/2020-12/schema" + + return schema + + +def SFT_SCHEMA(prompt_template: str | None = None): + """Generate JSON schema for SFT datasets. + + Uses the SFTDatasetSchemaType union which supports: + - SFTPromptTemplateDatasetItemSchema: Flexible prompt template format + - SFTChatDatasetItemSchema: Chat format with messages and tools + + Args: + prompt_template: Optional template string with placeholders like "{input} {output}". + If None or empty string, defaults to standard prompt/completion format. + Ignored for chat format detection. + + Returns: + JSON schema dict with required fields based on the format. + """ + from pydantic import TypeAdapter + + # Determine required fields for prompt template format + if prompt_template is not None and prompt_template != "": + # Extract placeholders from template + found_keys = re.findall(r"{(.*?)}", prompt_template) + + # TODO: Are we constrained by len == 2? + # Check for duplicates + if len(found_keys) != len(set(found_keys)): + duplicates = [key for key in found_keys if found_keys.count(key) > 1] + unique_duplicates = list(dict.fromkeys(duplicates)) + raise ValueError( + f"Prompt template contains duplicate placeholders: {unique_duplicates}. " + f"Each placeholder should appear only once." + ) + + prompt_template_keys = found_keys + else: + prompt_template_keys = ["prompt", "completion"] + + # Create TypeAdapter for the SFT union type to generate base JSON schema + adapter = TypeAdapter(SFTDatasetSchemaType) + schema = adapter.json_schema() + + # Add JSON schema metadata + schema["$schema"] = "https://json-schema.org/draft/2020-12/schema" + schema["title"] = "SFT Schema" + + # Update the prompt template sub-schema with required fields from prompt_template_keys + # The schema structure has $defs with the actual schemas, and oneOf/anyOf with $ref pointers + if "$defs" in schema: + # Update the SFTPromptTemplateDatasetItemSchema in $defs + if "SFTPromptTemplateDatasetItemSchema" in schema["$defs"]: + template_schema = schema["$defs"]["SFTPromptTemplateDatasetItemSchema"] + # Add template fields as required properties + if "properties" not in template_schema: + template_schema["properties"] = {} + for key in prompt_template_keys: + template_schema["properties"][key] = {"type": "string"} + template_schema["required"] = prompt_template_keys + template_schema["additionalProperties"] = True + return schema + + +# The RL backend trains DPO today; only the DPO schema is wired into validation. +# SFT_SCHEMA is retained for parity/headroom but is not registered here because +# the RL TrainingType enum has no SFT member. +SCHEMAS: dict[str, Callable[[str | None], dict]] = { + TrainingType.DPO.value: DPO_SCHEMA, +} + + +class DatasetValidator: + """Validator for training datasets. + + This class encapsulates dataset validation logic and avoids parameter drilling + by storing configuration as instance attributes. + + Example usage from dpo_config.py after prepare_dataset(): + ```python + from nmp.rl.tasks.training.datasets.preparation import prepare_dataset + from nmp.rl.tasks.training.datasets.validation import DatasetValidator + + # Prepare datasets + prepared = prepare_dataset( + dataset_path=Path(customizer_config.dataset.path), + output_dir=workspace_dir / "dataset", + ) + + # Validate the prepared datasets + validator = DatasetValidator( + training_type=customizer_config.training.training_type, + finetuning_type=customizer_config.training.finetuning_type, + prompt_template=customizer_config.dataset.prompt_template, + ) + validator.validate_dataset(str(prepared.train_file)) + validator.validate_dataset(str(prepared.validation_file)) + ``` + """ + + def __init__( + self, + training_type: TrainingType, + finetuning_type: Optional[FinetuningType] = None, + *, + prompt_template: str | None = None, + ): + """Initialize validator with training configuration. + + Args: + training_type: The type of training (DPO, etc.) + finetuning_type: Optional finetuning type (LoRA, all_weights, etc.) + prompt_template: Optional prompt template for datasets + """ + self.training_type = training_type + self.finetuning_type = finetuning_type + self.prompt_template = prompt_template + + def _validate_json_object(self, obj: dict, schema: dict[str, Any]) -> None: + """Validate a JSON object against a schema. + + Args: + obj: The JSON object to validate + schema: The JSON schema to validate against + + Raises: + TypeError: If validation fails + """ + try: + jsonschema.validate(instance=obj, schema=schema) + except exceptions.ValidationError as e: + logger.debug(f"Dataset Schema Validation failed: {str(e)}") + raise TypeError(f"Dataset Schema Validation failed: {e.message}") + except Exception as e: + logger.debug(f"Dataset Schema Validation failed: {str(e)}") + raise TypeError(f"Dataset Schema Validation failed: {e}") + + def detect_dataset_schema(self, file_path: str) -> str: + """Detect the dataset schema from the first line of the file. + + Args: + file_path: Path to the dataset file + + Returns: + Schema name (e.g., 'dpo') + + Raises: + DatasetFormatError: If file format is invalid or doesn't match any schema + """ + first_line = _first_nonempty_line(file_path) + if first_line is None: + raise DatasetFormatError(f"{file_path} has no non-empty rows") + + try: + obj: dict[str, Any] = json.loads(first_line) + except Exception as e: + # Log identifiers only — the raw row can contain customer training data. + logger.debug(f"{file_path}: first row is not valid JSON: {e}") + raise DatasetFormatError(f"{file_path} has an entry which is not valid JSON: {e}") + + for schema_name, schema_factory in SCHEMAS.items(): + try: + validation_schema = schema_factory(self.prompt_template) + self._validate_json_object(obj, validation_schema) + except Exception as e: + logger.debug(f"Parsed jsonl line does not conform to schema {schema_name}. Error: {e}") + else: + logger.debug(f"Parsed jsonl line conforms to schema {schema_name}.") + return schema_name + + raise DatasetFormatError("Dataset does not match any supported format") + + def validate_dataset(self, file_path: str, dataset_type: Optional[str] = None) -> None: + """Validate a single dataset file. + + Args: + file_path: Path to the dataset file + dataset_type: Optional dataset type to validate against. If None, uses training type from config + + Raises: + DatasetFormatError: If dataset is empty or validation fails + """ + # Use provided dataset_type or fall back to training type from config + if dataset_type is None: + dataset_type = self.training_type.value + + schema_factory = SCHEMAS.get(dataset_type) + if not schema_factory: + # Fail loudly: a typo or an unwired training type would otherwise skip + # validation entirely and let a malformed dataset reach training. + raise DatasetFormatError(f"Unsupported dataset_type for validation: {dataset_type}") + + if os.path.getsize(file_path) == 0: + raise DatasetFormatError(f"{file_path} is empty") + + validation_schema = schema_factory(self.prompt_template) + is_dpo = dataset_type == TrainingType.DPO.value + expected_dpo_schema: str | None = None + validated_rows = 0 + + # Validate each line in the JSONL file. Log identifiers (path + row number) + # only — never the raw line/object, which can contain customer training data. + with open(file_path, "r", encoding="utf-8") as jsonl_file: + for line_number, raw_line in enumerate(jsonl_file, start=1): + line = raw_line.strip() + if not line: + continue + + try: + obj: dict[str, Any] = json.loads(line) + except Exception as e: + logger.debug(f"{file_path}:{line_number} is not valid JSON: {e}") + raise DatasetFormatError(f"{file_path} line {line_number} is not valid JSON: {e}") + + # Reject files that mix multiple DPO schemas: detect_dpo_schema_name() + # later selects one concrete loader from the first row, so a mixed + # file would silently be handed to the wrong NeMo-RL dataset class. + if is_dpo: + row_schema = get_preference_dataset_discriminator(obj) + if expected_dpo_schema is None: + expected_dpo_schema = row_schema + elif row_schema != expected_dpo_schema: + raise DatasetFormatError( + f"{file_path} mixes DPO dataset schemas: expected {expected_dpo_schema}, " + f"got {row_schema} on line {line_number}" + ) + + try: + self._validate_json_object(obj, validation_schema) + except Exception as e: + logger.debug(f"{file_path}:{line_number} does not conform to the expected schema: {e}") + raise DatasetFormatError( + f"{file_path} line {line_number} does not conform to the expected schema: {e}" + ) + validated_rows += 1 + + # A whitespace-only file would otherwise pass silently (no rows validated). + if validated_rows == 0: + raise DatasetFormatError(f"{file_path} has no non-empty rows to validate") + + +def _first_nonempty_line(file_path: str | Path) -> str | None: + """Return the first non-blank line (stripped) of a file, or None if there is none.""" + with open(file_path, "r", encoding="utf-8") as f: + for line in f: + stripped = line.strip() + if stripped: + return stripped + return None + + +def detect_dpo_schema_name(file_path: str | Path) -> str: + """Detect the DPO preference dataset schema from the first line of the file. + + This function reads the first line of a JSONL dataset file and determines + which preference dataset schema it matches. It's designed to be called after + prepare_dataset() to dynamically determine the correct NeMo RL dataset class. + + For DPO training, it detects one of: + - PreferenceDataset: Native format with context + ranked completions + - BinaryPreferenceDataset: Simple prompt/chosen_response/rejected_response + - HelpSteer3: NVIDIA HelpSteer3 format with preference scores + - Tulu3Preference: AllenAI Tulu3 format with message lists + + Args: + file_path: Path to the dataset file (JSONL format) + + Returns: + The NeMo RL dataset class name (e.g., "BinaryPreferenceDataset", "HelpSteer3") + + Raises: + DatasetFormatError: If the file is empty or not valid JSON + """ + file_path = Path(file_path) + + if not file_path.exists(): + raise DatasetFormatError(f"Dataset file not found: {file_path}") + + # Read the first non-empty row so a leading blank line doesn't break detection. + first_line = _first_nonempty_line(file_path) + if first_line is None: + raise DatasetFormatError(f"Dataset file has no content: {file_path}") + + # Parse as JSON + try: + obj: dict[str, Any] = json.loads(first_line) + except json.JSONDecodeError as e: + raise DatasetFormatError(f"First row of {file_path} is not valid JSON: {e}") + + # Use the discriminator function to detect the schema type (returns NeMo RL class name directly) + dataset_name = get_preference_dataset_discriminator(obj) + logger.debug(f"Detected DPO preference dataset: {dataset_name} for {file_path}") + + logger.info(f"Detected dataset schema '{dataset_name}' for {file_path}") + return dataset_name + + +# Backward compatibility: provide standalone functions that create a validator instance +def detect_dataset_schema( + file_path: str, + training_type: TrainingType, + *, + finetuning_type: Optional[FinetuningType] = None, + prompt_template: str | None = None, +) -> str: + """Detect the dataset schema from the first line of the file.""" + validator = DatasetValidator(training_type, finetuning_type, prompt_template=prompt_template) + return validator.detect_dataset_schema(file_path) + + +def validate_dataset( + file_path: str, + training_type: TrainingType, + *, + dataset_type: Optional[str] = None, + finetuning_type: Optional[FinetuningType] = None, + prompt_template: str | None = None, +) -> None: + """Validate a single dataset file.""" + validator = DatasetValidator(training_type, finetuning_type, prompt_template=prompt_template) + validator.validate_dataset(file_path, dataset_type) + + +def validate_datasets( + file_names: list[str], + training_type: TrainingType, + *, + dataset_type: Optional[str] = None, + finetuning_type: Optional[FinetuningType] = None, + prompt_template: str | None = None, +) -> None: + """Validate a list of dataset files.""" + validator = DatasetValidator(training_type, finetuning_type, prompt_template=prompt_template) + for file_name in file_names: + validator.validate_dataset(file_name, dataset_type) diff --git a/services/rl/src/nmp/rl/tasks/training/distributed.py b/services/rl/src/nmp/rl/tasks/training/distributed.py new file mode 100644 index 0000000000..1a51dc6a22 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/distributed.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Distributed training coordination utilities. + +Provides role detection and file-based barrier synchronization for multi-node +training where multiple pods/containers run the same entry point. +""" + +import logging +import os +import time +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path + +logger = logging.getLogger(__name__) + +# Environment variables for distributed training injected by Volcano's pytorch plugin. +# Do not confuse these with the same env vars injected by torchrun. +# Here, WORLD_SIZE refers to number of nodes, while torchrun's WORLD_SIZE is the number of GPUs. +# RANK refers to the rank of the node, while torchrun's RANK is the global rank of the GPU. +RANK_ENVVAR = "RANK" +WORLD_SIZE_ENVVAR = "WORLD_SIZE" + +# Marker file the coordinator writes when it fails before signaling a barrier. +# Workers poll for it so they can abort promptly instead of waiting out the timeout. +COORDINATOR_FAILURE_MARKER = "coordinator.failed" + + +class CoordinatorFailedError(RuntimeError): + """Raised in a worker when the coordinator has published a failure marker.""" + + +class DistributedRole(Enum): + """Role of this node in distributed training.""" + + COORDINATOR = "coordinator" # Rank 0 - runs all phases + WORKER = "worker" # Rank > 0 - only participates in training + + +@dataclass +class DistributedContext: + """ + Distributed training context with file-based barrier synchronization. + + In multi-node training, all pods run the same entry point. This context + provides: + - Role detection (coordinator vs worker) based on RANK + - File-based barriers for cross-pod synchronization + + File barriers work by: + - Coordinator creates marker files to signal phase completion + - Workers poll for marker files before proceeding + - All ranks can sync via mutual signal-and-wait + + Attributes: + role: Whether this node is coordinator (rank 0) or worker + rank: This node's rank in the distributed job + world_size: Total number of nodes participating + barrier_dir: Directory for barrier marker files (on shared storage). + Must be provided by caller for multi-node; None for single-node. + """ + + role: DistributedRole + rank: int + world_size: int + barrier_dir: Path + _barrier_timeout: float = field(default=600.0, repr=False) + _poll_interval: float = field(default=0.5, repr=False) + + @classmethod + def from_env(cls, barrier_dir: Path) -> "DistributedContext": + """ + Create distributed context from environment variables. + + The caller is responsible for constructing the barrier_dir path, + including any task-specific namespacing for pause/resume support. + + Args: + barrier_dir: Directory for barrier files (on shared storage). + Caller should namespace this by task ID for pause/resume support. + + Environment Variables: + RANK: This node's rank (default: 0) + WORLD_SIZE: Total number of nodes (default: 1) + + Returns: + Configured DistributedContext + """ + rank = int(os.environ.get(RANK_ENVVAR, "0")) + world_size = int(os.environ.get(WORLD_SIZE_ENVVAR, "1")) + + role = DistributedRole.COORDINATOR if rank == 0 else DistributedRole.WORKER + + # Setup barrier directory if distributed + if world_size > 1: + barrier_dir.mkdir(parents=True, exist_ok=True) + + # Coordinator clears stale marker files from previous task runs + # (e.g., after pause/resume or retry). We unlink individual markers + # rather than rmtree the directory: rmtree can race with a worker that + # has already created the directory or a live marker, deleting state + # out from under it and deadlocking the barrier. + if role == DistributedRole.COORDINATOR: + logger.info(f"Cleaning up stale barrier markers from previous run: {barrier_dir}") + for stale_marker in list(barrier_dir.glob("*.ready")) + list( + barrier_dir.glob(COORDINATOR_FAILURE_MARKER) + ): + try: + stale_marker.unlink() + except OSError as e: + logger.warning(f"Failed to remove stale barrier marker {stale_marker}: {e}") + + ctx = cls( + role=role, + rank=rank, + world_size=world_size, + barrier_dir=barrier_dir, + ) + + logger.info( + f"Distributed context: rank={rank}, world_size={world_size}, " + f"role={role.value}, barriers={'enabled' if ctx.is_distributed else 'disabled'}" + ) + + return ctx + + @property + def is_coordinator(self) -> bool: + """True if this is the coordinator node (rank 0).""" + return self.role == DistributedRole.COORDINATOR + + @property + def is_distributed(self) -> bool: + """True if running in multi-node mode.""" + return self.world_size > 1 + + # --- Barrier Implementation --- + + def _marker_path(self, barrier_name: str, rank: int) -> Path: + """Get path to barrier marker file for a specific rank.""" + return self.barrier_dir / f"{barrier_name}.rank{rank}.ready" + + def _failure_marker_path(self) -> Path: + """Get path to the coordinator failure marker file.""" + return self.barrier_dir / COORDINATOR_FAILURE_MARKER + + def signal_failure(self) -> None: + """ + Publish a coordinator failure marker (coordinator only). + + Workers blocked in :meth:`wait_for_coordinator` / :meth:`wait_all` poll + for this marker and abort with :class:`CoordinatorFailedError` instead of + stranding on the barrier until the timeout expires. + """ + if not self.is_distributed or not self.is_coordinator: + return + + try: + self._failure_marker_path().touch() + logger.info("Published coordinator failure marker to release waiting workers") + except OSError as e: + logger.warning(f"Failed to write coordinator failure marker: {e}") + + def signal(self, barrier_name: str) -> None: + """ + Signal that this rank has reached a synchronization point. + + Creates a marker file indicating this rank is ready. + + Args: + barrier_name: Name of the barrier (should be unique per sync point) + """ + if not self.is_distributed: + return + + marker = self._marker_path(barrier_name, self.rank) + marker.touch() + logger.debug(f"Barrier signal: {barrier_name} (rank {self.rank})") + + def wait_for_coordinator(self, barrier_name: str, timeout: float | None = None) -> None: + """ + Wait for the coordinator (rank 0) to signal. + + Used by workers to wait for coordinator to complete a phase. + + Args: + barrier_name: Name of the barrier to wait for + timeout: Override default timeout (seconds) + + Raises: + CoordinatorFailedError: If the coordinator published a failure marker + TimeoutError: If coordinator doesn't signal within timeout + """ + if not self.is_distributed: + return + + if self.is_coordinator: + # Coordinator doesn't wait for itself + return + + timeout = self._barrier_timeout if timeout is None else timeout + marker = self._marker_path(barrier_name, rank=0) + failure_marker = self._failure_marker_path() + start = time.time() + + logger.debug(f"Waiting for coordinator at barrier: {barrier_name}") + + while time.time() - start < timeout: + if marker.exists(): + logger.debug(f"Coordinator signaled barrier: {barrier_name}") + return + if failure_marker.exists(): + raise CoordinatorFailedError(f"Coordinator reported failure while waiting at barrier '{barrier_name}'") + time.sleep(self._poll_interval) + + raise TimeoutError(f"Timeout waiting for coordinator at barrier '{barrier_name}' after {timeout}s") + + def wait_all(self, barrier_name: str, timeout: float | None = None) -> None: + """ + Wait for all ranks to reach this barrier. + + All ranks must call signal() before any rank proceeds. + + Args: + barrier_name: Name of the barrier + timeout: Override default timeout (seconds) + + Raises: + CoordinatorFailedError: If the coordinator published a failure marker + TimeoutError: If not all ranks signal within timeout + """ + if not self.is_distributed: + return + + timeout = self._barrier_timeout if timeout is None else timeout + failure_marker = self._failure_marker_path() + start = time.time() + + logger.debug(f"Waiting for all ranks at barrier: {barrier_name}") + + while time.time() - start < timeout: + ready_count = sum(1 for r in range(self.world_size) if self._marker_path(barrier_name, r).exists()) + if ready_count >= self.world_size: + logger.debug(f"All ranks reached barrier: {barrier_name}") + return + # A non-coordinator rank should bail out if the coordinator died; the + # coordinator itself never waits on its own failure marker. + if not self.is_coordinator and failure_marker.exists(): + raise CoordinatorFailedError(f"Coordinator reported failure while waiting at barrier '{barrier_name}'") + time.sleep(self._poll_interval) + + # Report which ranks are missing for debugging + missing = [r for r in range(self.world_size) if not self._marker_path(barrier_name, r).exists()] + raise TimeoutError(f"Timeout at barrier '{barrier_name}' after {timeout}s. Missing ranks: {missing}") + + def sync_point(self, barrier_name: str, timeout: float | None = None) -> None: + """ + Synchronization point where all ranks must arrive before any proceed. + + Combines signal() and wait_all() - this rank signals and then waits + for all other ranks. + + Args: + barrier_name: Name of the sync point + timeout: Override default timeout (seconds) + """ + self.signal(barrier_name) + self.wait_all(barrier_name, timeout) + + def cleanup_barrier(self, barrier_name: str) -> None: + """ + Clean up barrier marker files (coordinator only). + + Call after all ranks have passed the barrier. + + Args: + barrier_name: Name of the barrier to clean up + """ + if not self.is_distributed or not self.is_coordinator: + return + + for r in range(self.world_size): + marker = self._marker_path(barrier_name, r) + try: + if marker.exists(): + marker.unlink() + except OSError as e: + logger.warning(f"Failed to clean up barrier marker {marker}: {e}") diff --git a/services/rl/src/nmp/rl/tasks/training/errors/converter.py b/services/rl/src/nmp/rl/tasks/training/errors/converter.py new file mode 100644 index 0000000000..d6c0f097de --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/errors/converter.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import subprocess +from pathlib import Path + +from nmp.common.errors import ExceptionConverter, RulesLoader + +from .exceptions import ( + EXCEPTION_REGISTRY, + CustomizerTrainingError, + ErrorDetails, + InternalError, + default_exception_handler, +) + +logger = logging.getLogger(__name__) + +# Path to the error rules YAML file (relative to this module) +_ERROR_RULES_PATH = Path(__file__).parent / "error_rules.yaml" + +# Additional modules to search for exception types not in the registry +# subprocess.TimeoutExpired is used for training timeout detection +_FALLBACK_MODULES = [subprocess] + +# Module-level singleton converter +_converter: ExceptionConverter | None = None + + +def _load_converter() -> ExceptionConverter: + """Load the converter from YAML rules.""" + logger.debug(f"Loading Customizer error rules from: {_ERROR_RULES_PATH}") + + converter = RulesLoader.from_yaml( + _ERROR_RULES_PATH, + exception_registry=EXCEPTION_REGISTRY, + default_handler=default_exception_handler, + fallback_exception=InternalError, + fallback_modules=_FALLBACK_MODULES, + ) + + logger.info(f"Loaded {converter.rule_count} Customizer error mapping rules") + return converter + + +def get_error_converter() -> ExceptionConverter: + """ + Get the singleton ExceptionConverter for Customizer training errors. + + The converter is created once on first access and reused for the module's lifetime. + It loads rules from error_rules.yaml and uses InternalError as fallback. + + Returns: + Configured ExceptionConverter ready to convert exceptions. + + Raises: + FileNotFoundError: If error_rules.yaml is not found. + ValueError: If rules file has invalid syntax. + """ + global _converter + if _converter is None: + _converter = _load_converter() + return _converter + + +def create_error_details(exception: Exception) -> ErrorDetails: + """ + Create error_details dict for Jobs service reporting. + + Converts the exception to a CustomizerTrainingError and returns + a dict suitable for passing to progress_reporter.report_error(). + + If the exception is already a CustomizerTrainingError, returns its + details directly without re-conversion. + + Uses the library's fallback mechanism (InternalError) for unmatched exceptions. + + Args: + exception: The exception to convert. + + Returns: + ErrorDetails with 'message', 'type', and 'detail' keys. + """ + # If already a CustomizerTrainingError, return its details directly + if isinstance(exception, CustomizerTrainingError): + return exception.to_error_details() + + # Convert using the library - fallback_exception=InternalError handles unmatched. + # get_error_converter() is inside the try so a converter-loading failure (e.g. a + # missing/invalid rules file) still yields a structured InternalError dict rather + # than propagating out of this last-resort error reporter. + 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), + ) + return exc.to_error_details() + + +__all__ = [ + "get_error_converter", + "create_error_details", +] diff --git a/services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml b/services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml new file mode 100644 index 0000000000..4615023796 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/errors/error_rules.yaml @@ -0,0 +1,645 @@ +# This file defines rules for converting low-level training exceptions +# into user-friendly CustomizerTrainingError subclasses. +# Rules for all backends are present in the same yaml file. +# +# Rules are evaluated in order; first match wins. +# +# Rule structure: +# - : # When to match (pick ONE) +# exception: # Exception class from EXCEPTION_REGISTRY +# error_details: # Optional user-friendly message + + +rules: + # =========================================================================== + # 1. TRAINING TIMEOUT (subprocess.TimeoutExpired) + # All backends + # =========================================================================== + + - type: TimeoutExpired # subprocess.TimeoutExpired from fallback_modules + exception: TrainingTimeoutError + error_details: "Training exceeded the maximum allowed time limit. To reduce training time: 1) Reduce max_steps or epochs, 2) Use a smaller dataset, 3) Use a smaller model, 4) Use LoRA/PEFT instead of all_weights fine-tuning (LoRA trains faster), or 5) Increase batch_size to process more samples per step (if GPU memory allows). If you need longer training times, contact your administrator to adjust the job timeout limits." + + # =========================================================================== + # 2. DATASET FORMAT ERRORS (400) + # =========================================================================== + + # --- Automodel --- + # Unsupported role in chat messages + - regex: "Unsupported role in messages: \\w+" + exception: DatasetFormatError + error_details: "Your dataset contains chat messages with an invalid role. Each message in a conversation must have a 'role' field with one of the following values: 'system' (for system prompts), 'user' (for user inputs), 'assistant' (for model responses), or 'tool' (for tool/function outputs). Please check your dataset and ensure all messages use valid roles." + + # --- NeMo-RL --- + # Text type error + - regex: "^text must be a string or a list of strings, got .+$" + exception: DatasetFormatError + error_details: "The 'text' field in your dataset has an invalid type. For NeMo-RL training (DPO/GRPO), the text field must be either a single string or a list of strings. Please check your dataset format and ensure the text field contains the correct data type." + + # Prompt file not found + - regex: "^Prompt file .+ not found$" + exception: DatasetFormatError + error_details: "The prompt template file specified in your training dataset configuration does not exist. Prompt templates define how your dataset samples are formatted for training. Please verify the prompt file path is correct and the file is accessible at the specified location." + + # --- Automodel --- + # Empty dataset + - regex: "^no sample to consume: \\d+$" + exception: DatasetFormatError + error_details: "Your dataset is empty or contains zero valid samples after filtering. This can happen if: 1) The dataset file is empty, 2) All samples were filtered out due to format issues, or 3) The dataset path is incorrect. Please verify your dataset contains valid training samples." + + # All samples consumed + - regex: "^no samples left to consume: \\d+, \\d+$" + exception: DatasetFormatError + error_details: "All samples in your dataset have been consumed before completing the requested number of training steps. This happens when your dataset is too small for the configured epochs or max_steps. Please either: 1) Add more samples to your dataset, 2) Reduce the number of epochs, or 3) Reduce max_steps." + + # Error loading example + - regex: "Error while loading example \\d+ from dataset .+" + exception: DatasetFormatError + error_details: "Failed to load a specific sample from your dataset. This typically indicates a malformed sample that doesn't match the expected format. Please check your dataset for: 1) Missing required fields, 2) Invalid JSON formatting, 3) Incorrect data types for fields. The error message includes the sample index to help you locate the problematic entry." + + # =========================================================================== + # 3. MODEL NOT FOUND ERRORS (404) + # Megatron Bridge + # =========================================================================== + + # Checkpoint file not found (input model checkpoint for training) + - regex: "^Checkpoint file not found: .+$" + exception: ModelNotFoundError + error_details: "The input model checkpoint file could not be found. Please verify the base model path is correct and accessible. This checkpoint is used as the starting point for training." + + # No checkpoints found for resume (output checkpoint directory empty) + - regex: "There were no checkpoints found in checkpoint_dir.*Cannot resume" + exception: ModelNotFoundError + error_details: "The output checkpoint directory is empty. Cannot resume training because no previous training checkpoints were found. Ensure a prior training run completed successfully and saved checkpoints." + + # Nemotron model missing HF source + - regex: "Nemotron Super models expect HF source code to exist at .+" + exception: ModelNotFoundError + error_details: "The Nemotron Super model checkpoint is missing the required HuggingFace source code directory (nemotron_src/). This directory must be present inside the model checkpoint. Please ensure you are using a complete Nemotron Super model checkpoint that includes the HuggingFace source files." + + # =========================================================================== + # 4. MODEL LOAD ERRORS (500) + # =========================================================================== + + # --- Automodel --- + # Model weights swap failure + - contains: "_apply(): Couldn't swap" + exception: ModelLoadError + error_details: "Failed to load the base model: weights could not be applied to a model layer. The base model checkpoint may be corrupted, incomplete, or incompatible with the selected training configuration." + + # Model patching failure + - exact: "Failed to patch model" + exception: ModelLoadError + error_details: "Failed to apply optimizations to the base model. The base model architecture may not be supported for the selected training configuration. Try using a different model or training method." + + # Method signature mismatch + - starts_with: "Signature mismatch:" + exception: ModelLoadError + error_details: "The base model has an incompatible method signature. This typically indicates a version mismatch between the base model and the training framework. Please verify you are using a supported model version." + + # Missing lm_head.weight + - exact: "lm_head.weight not found in model" + exception: ModelLoadError + error_details: "The base model is missing the language model head (lm_head.weight). The base model checkpoint may be corrupted, incomplete, or not a valid language model. Please verify the base model is a complete, valid language model checkpoint." + + # --- NeMo-RL --- + # vLLM not installed + - contains: "vLLM is not installed" + exception: ModelLoadError + error_details: "vLLM is not installed in the training environment. This is an issue with the training environment setup, please contact the administrator to raise an issue with the NeMo Platform team." + + # Missing generation output keys + - regex: "^Missing required keys for GenerationOutputSpec: .+$" + exception: ModelLoadError + error_details: "The base model's generation output is missing required fields. The base model may not be compatible with the selected training method (e.g., GRPO). Please verify you are using a supported model for this training type." + + # Missing score output keys + - regex: "^Missing required keys for ScoreOutputSpec: .+$" + exception: ModelLoadError + error_details: "The base model's score output is missing required fields. The base model may not be compatible with the selected training method. Please verify you are using a supported model for this training type." + + # Pretrained run config not found (Megatron HF-to-mcore conversion) + - contains: "Pretrained run config not found at" + exception: ModelLoadError + error_details: "The pretrained model configuration file was not found after Megatron checkpoint conversion. This usually means the HuggingFace-to-Megatron conversion on the head node saved to a directory not accessible by this worker node. This is an infrastructure issue - please ensure shared storage is properly mounted across all nodes, or contact your administrator." + + # --- Megatron Bridge --- + # Shape mismatch for parameter + - regex: "^Shape mismatch for parameter .+: target shape .+ vs source shape .+$" + exception: ModelLoadError + error_details: "The base model parameter shape does not match the checkpoint. The base model checkpoint may be from a different model architecture or an incompatible version. Please ensure the base model matches the expected architecture for this training configuration." + + # Shape mismatch for buffer + - regex: "^Shape mismatch for buffer .+: .+ vs .+$" + exception: ModelLoadError + error_details: "The base model buffer shape does not match the checkpoint. The base model checkpoint may be corrupted, incomplete, or from an incompatible model version. Please verify the base model checkpoint is valid and complete." + + # =========================================================================== + # 5. TRAINING CONFIG ERRORS - PARALLELISM (400) + # =========================================================================== + + # --- Automodel --- + # Pipeline parallelism: tied embeddings not supported + - all_keywords: ["not compatible with pipeline parallelism", "tie_word_embeddings"] + exception: TrainingConfigError + error_details: "The base model has tied embeddings (tie_word_embeddings=True) which is not compatible with pipeline parallelism. Try using a different parallelism configuration or a model without tied embeddings." + + # Pipeline parallelism: encoder-decoder models not supported + - all_keywords: ["not compatible with pipeline parallelism", "Encoder-Decoder"] + exception: TrainingConfigError + error_details: "The base model is an encoder-decoder architecture (like T5 or BART) which is not supported with pipeline parallelism. Please use a decoder-only base model, or disable pipeline parallelism in your training configuration." + + # PP batch size / microbatch validation + - contains: "pp_batch_size // pp_microbatch_size must be >= pp_size" + exception: TrainingConfigError + error_details: "Pipeline parallelism requires: batch_size >= pipeline_parallel_size. The current batch_size is too small to fill all pipeline stages. Either increase batch_size or reduce pipeline_parallel_size." + + # Context parallelism: SDPA not supported + - contains: "Model does not support SDPA required for context parallelism" + exception: TrainingConfigError + error_details: "The base model does not support scaled dot-product attention (SDPA) which is required for context parallelism. Please set context_parallel_size=1 to disable context parallelism." + + # --- NeMo-RL --- + # Megatron and DTensor both enabled + - exact: "Configure either Megatron (policy.megatron_cfg.enabled=true) or DTensor (policy.dtensor_cfg.enabled=true), not both." + exception: TrainingConfigError + error_details: "Internal configuration error: both Megatron and DTensor training backends are enabled, but only one can be active at a time. This is an issue with the training environment setup, please contact the administrator." + + # Neither Megatron nor DTensor enabled + - contains: "Please either set policy.megatron_cfg.enabled=true" + exception: TrainingConfigError + error_details: "Internal configuration error: no training backend is enabled. The training environment requires either Megatron or DTensor backend to be active. This is an issue with the training environment setup, please contact the administrator." + + # World size insufficient for parallelism + - regex: "^World size \\(\\d+\\) is insufficient for the parallelism configuration" + exception: TrainingConfigError + error_details: "Not enough GPUs available for the requested parallelism settings. The total number of GPUs must be at least pipeline_parallel_size * context_parallel_size * tensor_parallel_size. Either reduce parallelism settings or request more GPUs." + + # World size not divisible by parallelism + - regex: "^World size \\(\\d+\\) must be divisible by PP \\* CP \\* TP" + exception: TrainingConfigError + error_details: "The total number of GPUs must be evenly divisible by (pipeline_parallel_size * context_parallel_size * tensor_parallel_size). For example, with PP=2, CP=1, TP=2, you need 4, 8, 12, etc. GPUs. Please adjust your parallelism settings or cluster size." + + # DTensor world size mismatch + - regex: "^World size\\(\\d+\\) must equal to dp_size\\(\\d+\\) \\* tp_size\\(\\d+\\) \\* cp_size\\(\\d+\\) to use DTensor$" + exception: TrainingConfigError + error_details: "The total number of GPUs (world_size) does not match the product of data_parallel_size * tensor_parallel_size * context_parallel_size for the DTensor backend. Please adjust your parallelism settings so they are consistent with the available GPU count." + + # Dynamic batching with PP > 1 + - contains: "Dynamic batching is only supported for single pipeline parallel stage" + exception: TrainingConfigError + error_details: "Dynamic batching is only supported when pipeline_parallel_size=1. With pipeline parallelism (PP > 1), the model is split across GPU stages which requires fixed batch sizes. Please either set pipeline_parallel_size=1 or disable dynamic batching." + + # Dynamic batching exclusive of sequence packing + - contains: "Dynamic Batching is exclusive of Sequence Packing" + exception: TrainingConfigError + error_details: "Dynamic batching and sequence packing cannot be used together. Please disable one of them: either set dynamic_batching=false or set sequence_packing_enabled=false." + + # Sequence packing not supported for VLM models + - contains: "Sequence packing is not supported for VLM models" + exception: TrainingConfigError + error_details: "Sequence packing is not supported for Vision-Language Models (VLMs). Please set sequence_packing_enabled=false when training VLM models." + + # Context parallel not supported for sequence packing (DTensor) + - exact: "Context parallel is not supported for sequence packing. Refer to https://github.com/NVIDIA/NeMo-RL/blob/main/docs/model-quirks.md#context-parallel-with-fsdp2 for more details." + exception: TrainingConfigError + error_details: "Context parallelism cannot be used with sequence packing in the DTensor backend. Please either set context_parallel_size=1 to disable context parallelism, or set sequence_packing_enabled=false to disable sequence packing." + + # Context parallel not supported for Gemma3 + - contains: "Context parallel is not supported for Gemma3ForCausalLM" + exception: TrainingConfigError + error_details: "Context parallelism is not supported for Gemma3 models due to limitations in the PyTorch context parallel implementation. Please set context_parallel_size=1 when training Gemma3 models." + + # Context parallel not supported for VLM models + - contains: "Context parallel is yet not supported for VLM models" + exception: TrainingConfigError + error_details: "Context parallelism is not yet supported for Vision-Language Models (VLMs). Please set context_parallel_size=1 when training VLM models." + + # Context parallelism requires sequence packing (Megatron) + - contains: "Context Parallelism (CP>1) requires sequence packing to be enabled" + exception: TrainingConfigError + error_details: "When using the Megatron backend with context_parallel_size > 1, sequence packing must be enabled. Please either enable sequence packing (sequence_packing_enabled=true) or reduce context_parallel_size to 1." + + # Reward models not supported with Megatron backend + - contains: "Reward models are not yet supported with the Megatron backend" + exception: TrainingConfigError + error_details: "Reward models are not yet supported with the Megatron training backend. This is a current limitation of the framework. Please use the DTensor backend for reward model training, or contact your administrator for alternative configurations." + + # Dynamic sampling max batches reached + - contains: "Dynamic sampling has reached the maximum allowed number of batches" + exception: TrainingConfigError + error_details: "Dynamic sampling exceeded the maximum number of generation batches allowed per training step. This means the training data or reward signal is too challenging for the model to produce enough valid samples. Consider: 1) Simplifying your dataset, 2) Adjusting num_prompts_per_step or num_generations_per_prompt, 3) Checking that your reward function is not too strict." + + # Batch size not divisible by DP + - regex: "Configuration error: \\(num_prompts_per_step \\* num_generations_per_prompt\\) = \\d+ must be divisible by data_parallel size \\d+" + exception: TrainingConfigError + error_details: "The effective batch size (num_prompts_per_step * num_generations_per_prompt) must be evenly divisible by the number of data parallel workers. Please adjust num_prompts_per_step or num_generations_per_prompt so their product divides evenly." + + # =========================================================================== + # 6. TRAINING CONFIG ERRORS - DPO/GRPO (400) + # NeMo-RL + # =========================================================================== + + # Dynamic batching with DPO + - contains: "Dynamic batching is currently not supported with DPO" + exception: TrainingConfigError + error_details: "DPO (Direct Preference Optimization) training does not support dynamic batching. This is an internal configuration issue with the training environment, please contact the administrator." + + # Sequence packing with DPO + - contains: "Sequence packing is currently not supported with DPO" + exception: TrainingConfigError + error_details: "DPO (Direct Preference Optimization) training does not support sequence packing. Please set sequence_packing_enabled=false in your training request." + + # GRPO requires generation config + - contains: "A generation config in the PolicyConfig is required for GRPO" + exception: TrainingConfigError + error_details: "GRPO (Group Relative Policy Optimization) requires a generation configuration to produce responses during training. This is an internal configuration issue with the training environment, please contact the administrator." + + # Validation dataset required + - exact: "Validation dataset is required if validation is enabled" + exception: TrainingConfigError + error_details: "Validation is enabled for this training job, but no validation dataset was provided. Please provide a validation dataset in your training request, or disable validation." + + # Non-colocated inference with Megatron + - contains: "Non-colocated inference is not supported for Megatron generation backends" + exception: TrainingConfigError + error_details: "The current training configuration uses Megatron for generation, which does not support the required inference mode. This is an internal configuration issue with the training environment, please contact the administrator." + + # Async GRPO requires vLLM async + - contains: "Async GRPO requires vLLM backend with vllm_cfg.async_engine=True" + exception: TrainingConfigError + error_details: "Async GRPO training requires the vLLM backend with async engine enabled, but the current configuration does not have this set. This is an internal configuration issue with the training environment, please contact the administrator." + + # Async GRPO requires importance sampling + - contains: "Importance sampling correction must be enabled for async GRPO" + exception: TrainingConfigError + error_details: "Async GRPO training requires importance sampling correction to handle off-policy samples and ensure stable training. This is an internal configuration issue with the training environment, please contact the administrator." + + # Async GRPO doesn't support colocated inference + - contains: "Colocated inference is not supported for async GRPO" + exception: TrainingConfigError + error_details: "Async GRPO training does not support colocated inference (running training and generation on the same GPUs). This is an internal configuration issue with the training environment, please contact the administrator." + + # top_k sampling threshold (vLLM V1 engine limitation) + - contains: "top_k sampling with values <" + exception: TrainingConfigError + error_details: "The top_k value is too low for the vLLM V1 engine. The vLLM V1 engine does not return logprobs after top_k filtering, so very low top_k values produce inaccurate logprob computations. Please increase top_k or remove the top_k constraint." + + # top_p sampling threshold (vLLM V1 engine limitation) + - contains: "top_p sampling with values <" + exception: TrainingConfigError + error_details: "The top_p value is too low for the vLLM V1 engine. The vLLM V1 engine does not return logprobs after top_p filtering, so very low top_p values produce inaccurate logprob computations. Please increase top_p or remove the top_p constraint." + + # MoE aux loss not supported + - contains: "MoE aux loss is currently not supported" + exception: TrainingConfigError + error_details: "Mixture-of-Experts (MoE) auxiliary loss is not currently supported due to a known bug in Megatron-LM. Please disable the MoE auxiliary loss in your training configuration." + + # =========================================================================== + # 7. TRAINING CONFIG ERRORS - PEFT/LORA (400) + # Automodel + # =========================================================================== + + # Triton not installed + - contains: "triton is not installed" + exception: TrainingConfigError + error_details: "The Triton library, which is required for optimized LoRA kernel operations, is not installed in the training environment. This is an issue with the training environment setup, please contact the administrator to ensure Triton is properly installed." + + # LoRA dimensions mismatch + - contains: "Incompatible X and LoRA A dimensions" + exception: TrainingConfigError + error_details: "The LoRA adapter dimensions are incompatible with the base model's layer dimensions. This can happen if you are trying to apply a pre-trained LoRA adapter that was created for a different model architecture. Please ensure the LoRA configuration (lora_dim/rank) is compatible with the base model you are fine-tuning." + + # =========================================================================== + # 8. TRAINING CONFIG ERRORS - PACKING (400) + # NeMo-RL + # =========================================================================== + + # Sequence too long for packing + - regex: "^Sequence length \\d+ exceeds bin capacity \\d+$" + exception: TrainingConfigError + error_details: "When sequence packing is enabled, one or more sequences in your dataset exceed the maximum sequence length (max_seq_length). Sequence packing combines multiple shorter sequences into a single training sample, but each individual sequence must fit within max_seq_length. Please either increase max_seq_length to accommodate longer sequences, or preprocess your dataset to truncate or remove sequences that are too long." + + # Not enough sequences for packing + - regex: "^Cannot create \\d+ bins with only \\d+ sequences" + exception: TrainingConfigError + error_details: "When sequence packing is enabled, the packing algorithm needs enough sequences to efficiently fill the training batches. Your dataset does not have enough sequences for the current batch configuration. Please either add more samples to your dataset, reduce the batch_size, or disable sequence packing by setting sequence_packing_enabled=false." + + # =========================================================================== + # 9. ENVIRONMENT ERRORS (400) + # NeMo-RL + # =========================================================================== + + # Unable to find compatible environment + - regex: "^Unable to find compatible environment - .+$" + exception: TrainingEnvironmentError + error_details: "The specified GRPO environment name is not recognized. GRPO (Group Relative Policy Optimization) requires a valid environment that defines how to evaluate model responses. Please check the environment name in your training request and ensure it matches one of the supported environments for your use case." + + # GRPO environment required + - exact: "hyperparameters.environment is required for GRPO, but it is not set" + exception: TrainingEnvironmentError + error_details: "GRPO (Group Relative Policy Optimization) training requires an environment configuration to evaluate model responses and compute rewards. Please specify the environment in your training request's hyperparameters. The environment determines how the model's generated responses will be scored during reinforcement learning." + + # No environment for task type + - regex: "^No environment found for task type: .+$" + exception: TrainingEnvironmentError + error_details: "No GRPO environment is registered for the specified task type. The environment defines how model responses are evaluated during reinforcement learning. This may indicate an unsupported task type or a misconfiguration. Please verify your task type is supported for GRPO training." + + # =========================================================================== + # 10. CHECKPOINT ERRORS (500) + # =========================================================================== + + # --- Automodel --- + # Checkpoint directory already exists + - regex: "Checkpoint directory .* already exists" + exception: CheckpointError + error_details: "The output checkpoint directory already exists from a previous training run. This typically happens when a previous training job failed or was cancelled but left partial checkpoint files behind. Please use a clean output directory, if you do not have access to remove the existing checkpoint directory, contact your administrator." + + # Global plan validation failure + - exact: "Failed to validate global plan" + exception: CheckpointError + error_details: "Checkpoint validation failed during distributed checkpoint loading. This occurs when the 'global plan' (which coordinates how model weights are distributed across GPUs) cannot be validated. Common causes include: 1) Corrupted checkpoint metadata files, 2) Mismatch between the number of GPUs used when saving vs loading the checkpoint, or 3) Interrupted checkpoint save operation. Please ensure the checkpoint is complete and you are using the same GPU topology as when the checkpoint was saved." + + # Missing key in checkpoint + - starts_with: "Missing key in checkpoint state_dict:" + exception: CheckpointError + error_details: "The checkpoint is missing one or more required model weights. This typically indicates that the checkpoint file is corrupted, incomplete (possibly from an interrupted save), or was created from a different model architecture than the one being loaded. Please verify the checkpoint is complete and matches the expected model architecture." + + # MoE expert weights missing + - contains: "Expert weights missing from checkpoint" + exception: CheckpointError + error_details: "The checkpoint for this Mixture-of-Experts (MoE) model is missing one or more expert weights. MoE models have multiple 'expert' sub-networks, and all expert weights must be present in the checkpoint. This typically indicates the checkpoint is corrupted or was saved incorrectly. Please use a complete, valid MoE checkpoint." + + # --- NeMo-RL --- + # Checkpoint file corrupted (JSONDecodeError) + - type_name: JSONDecodeError + exception: CheckpointError + error_details: "The checkpoint metadata file (training_info.json) is corrupted and cannot be parsed. This file stores training progress information like the current step and loss values. The checkpoint may have been saved incompletely or the file was corrupted during storage." + + # Distributed process group not initialized for checkpoint save + - exact: "Distributed process group is not initialized. Cannot save checkpoint." + exception: CheckpointError + error_details: "Cannot save checkpoint because the distributed process group is not initialized. This typically occurs when the training cluster encountered communication issues before checkpoint saving could complete. This is a transient infrastructure issue - please try running your training job again." + + # Megatron core state not initialized for checkpoint save + - exact: "Megatron core state or model is not initialized. Cannot save checkpoint." + exception: CheckpointError + error_details: "Cannot save checkpoint because the Megatron model state is not initialized. This typically occurs when the model failed to load or initialize correctly before training could produce a checkpoint. Please verify the base model is valid and try again." + + # HF checkpoint already exists + - regex: "^HF checkpoint already exists at .+\\. Delete it to run or set overwrite=True\\.$" + exception: CheckpointError + error_details: "The HuggingFace checkpoint output directory already exists from a previous training run or conversion. This typically happens when a previous training job left partial output behind. Please use a clean output directory, or contact your administrator to remove the existing checkpoint." + + # =========================================================================== + # 11. CUDA/GPU ERRORS (500) + # =========================================================================== + + # --- NeMo-RL --- + # Disk space exhausted - occurs in Ray cluster workers during RL training + # Ray stores session logs in /tmp/ray/session_*/logs/ which can fill up ephemeral node storage + - contains: "No space left on device" + exception: DistributedError + error_details: "Disk space exhausted on the node's ephemeral storage (/tmp). During reinforcement learning training (DPO/GRPO), Ray stores session logs and temporary files in /tmp/ray/ which can fill up the node's local disk. This is separate from the PVC used for checkpoints and datasets. This is typically a transient infrastructure issue - please try running your training job again, or contact your administrator to ensure adequate ephemeral storage is configured for the cluster nodes." + + # CUDA out of memory - catch by type name + - type_name: OutOfMemoryError + exception: CudaError + error_details: "GPU out of memory. To reduce memory usage: 1) Lower batch_size, 2) Reduce max_seq_length, 3) Use LoRA/PEFT instead of all_weights fine-tuning, or 4) Use a model with fewer parameters." + + # CUDA OOM - catch by message pattern + - contains: "CUDA out of memory" + exception: CudaError + error_details: "GPU out of memory. To reduce memory usage: 1) Lower batch_size, 2) Reduce max_seq_length, 3) Use LoRA/PEFT instead of all_weights fine-tuning, or 4) Use a model with fewer parameters." + + # NOTE: there is intentionally no bare "out of memory" catch-all here. It would + # also match host/Ray/system RAM exhaustion (e.g. the kernel OOM killer or a + # Python MemoryError) and misclassify those as a GPU CudaError, giving the user + # GPU-specific remediation advice for a non-GPU problem. The CUDA-specific rules + # above (type_name OutOfMemoryError + "CUDA out of memory") cover real GPU OOMs; + # anything else falls through to the generic handlers / fallback. + + # General CUDA errors + - and: + - any_keywords: ["CUDA", "cuda"] + - any_keywords: ["error", "Error", "failed", "Failed"] + exception: CudaError + error_details: "A GPU/CUDA error occurred. Please check GPU availability, ensure the GPU is not being used by another process, and try again." + + # =========================================================================== + # 12. DISTRIBUTED ERRORS (500) + # =========================================================================== + + # --- Automodel --- + # torch.distributed not available + - exact: "torch.distributed not available" + exception: DistributedError + error_details: "The PyTorch distributed package is not available in the training environment. Distributed training requires PyTorch to be built with distributed support enabled. This is an issue with the training environment setup, please contact the administrator to ensure the correct PyTorch version is installed." + + # torch.distributed not initialized + - exact: "expected torch.distributed to be initialized" + exception: DistributedError + error_details: "PyTorch distributed training was not properly initialized before the training process started. This typically happens when the training script is not launched correctly with the distributed launcher (torchrun). This is an issue with the training environment setup, please contact the administrator." + + # Distributed timeout - check for TimeoutError in cause chain + - cause: + type_name: TimeoutError + recursive: true + exception: DistributedError + error_details: "A distributed training operation timed out while waiting for communication between GPUs or nodes. This can happen when: 1) One or more GPU workers crashed or became unresponsive, 2) Network connectivity issues between nodes, 3) Uneven workload causing some GPUs to wait too long for others. This may be a transient issue - please try running your training job again. If the problem persists, contact your administrator." + + # NCCL errors + - any_keywords: ["NCCL", "nccl"] + exception: DistributedError + error_details: "An NCCL (NVIDIA Collective Communications Library) error occurred during GPU-to-GPU communication. NCCL is used to synchronize data between GPUs during distributed training. Common causes include: 1) Network connectivity issues between GPU nodes, 2) GPU hardware problems, 3) Incompatible NCCL versions, or 4) Memory pressure on GPUs. This may be a transient issue - please try running your training job again. If the problem persists, contact your administrator." + + # c10d errors + - contains: "c10d" + exception: DistributedError + error_details: "A PyTorch distributed communication error occurred (c10d is PyTorch's distributed communication backend). This indicates a failure in the inter-process or inter-node communication during distributed training. This may be caused by network issues, process crashes, or resource exhaustion. Please try running your training job again. If the problem persists, contact your administrator." + + # --- NeMo-RL --- + # Not enough GPUs + - and: + - type_name: ResourceInsufficientError + - contains: "Not enough GPUs available" + exception: DistributedError + error_details: "The training cluster does not have enough GPUs available for your requested configuration. Your training job requires more GPUs than are currently available in the cluster. Try reducing the parallelism settings (tensor_parallel_size, pipeline_parallel_size) to require fewer GPUs." + + # Not enough CPUs + - and: + - type_name: ResourceInsufficientError + - contains: "Not enough CPUs available" + exception: DistributedError + error_details: "The training cluster does not have enough CPUs available for your requested configuration. CPUs are needed for data loading and preprocessing alongside GPU training." + + # Maximum retries reached + - and: + - type_name: ResourceInsufficientError + - contains: "Maximum number of retries reached" + exception: DistributedError + error_details: "Failed to allocate cluster resources after multiple retry attempts. This is typically a transient issue - please wait a few minutes and try submitting your training job again. If the problem persists, contact your administrator to check cluster health." + + # Placement group timeout + - contains: "Timed out waiting for placement groups to be ready" + exception: DistributedError + error_details: "Timed out while waiting for Ray placement groups to be allocated. Placement groups are used to co-locate GPU workers on the same nodes for efficient communication. This typically happens when the cluster is under heavy load and cannot allocate the required resources in time. Please try submitting your training job again. If the problem persists, contact your administrator." + + # No valid placement groups + - contains: "No valid placement groups found" + exception: DistributedError + error_details: "No valid Ray placement groups could be found for the training job. This indicates a problem with the distributed training cluster configuration or resource availability. This is an infrastructure issue - please contact your administrator to investigate the cluster setup." + + # Workers per node mismatch + - regex: "^workers_per_node list length \\(\\d+\\) must match" + exception: DistributedError + error_details: "The workers-per-node configuration does not match the number of placement groups allocated. This indicates an internal mismatch in the distributed training setup. This is an infrastructure issue - please contact your administrator." + + # Missing sharding annotations + - exact: "Sharding annotations must be provided to use sharded data distribution" + exception: DistributedError + error_details: "The training configuration requires sharded data distribution but sharding annotations are not provided. Sharding annotations specify how data should be distributed across workers for efficient parallel processing. This is an internal configuration issue - please contact your administrator." + + # =========================================================================== + # 13. GENERATION ERRORS (500) + # NeMo-RL + # =========================================================================== + + # Weight update failed during refit + - regex: "^Updating weights for the generation policy failed during refit" + exception: GenerationError + error_details: "Failed to update the vLLM generation model weights from the training policy during the 'refit' step. In GRPO training, the generation model periodically syncs weights from the training model. This failure may be caused by: 1) CUDA IPC (Inter-Process Communication) issues between training and generation workers, 2) NCCL communication errors, or 3) Memory pressure on GPUs. This is typically a transient issue - please try running your training job again." + + # generate_text with async_engine + - contains: "generate_text cannot be used with async_engine=True" + exception: GenerationError + error_details: "A synchronous generation method was called on an async vLLM engine. When async_engine is enabled, you must use async methods (e.g., generate_text_async). This is an internal configuration issue with the training environment, please contact the administrator." + + # update_weights_via_ipc with async_engine + - contains: "cannot be used with async_engine=True" + exception: GenerationError + error_details: "A synchronous method was called on an async vLLM engine. When async_engine is enabled, all vLLM operations must use their async variants. This is an internal configuration issue with the training environment, please contact the administrator." + + # Error in sample rollout + - regex: "^Error in sample \\d+ rollout: .+$" + exception: GenerationError + error_details: "An error occurred while generating a response (rollout) for one of the training samples during GRPO training. Rollouts are the model-generated responses used to compute rewards and policy gradients. This may be caused by: 1) Invalid input data in the sample, 2) Generation parameters causing issues (e.g., max_tokens too low), or 3) vLLM backend errors. Check your dataset for problematic samples." + + # Async generation not enabled + - contains: "Async generation is not enabled" + exception: GenerationError + error_details: "Async generation was requested but the vLLM engine is not configured with async_engine=True. Async generation allows overlapping training and generation for better throughput. This is an internal configuration issue with the training environment, please contact the administrator." + + # NeMo-Gym/Penguin requires async vLLM + - contains: "you must use vllm generation backend with" + exception: GenerationError + error_details: "The NeMo-Gym (Penguin) environment requires the vLLM generation backend with async_engine enabled. NeMo-Gym provides advanced RL training features that depend on async generation. This is an internal configuration issue with the training environment, please contact the administrator." + + # NeMo-Gym/Penguin requires HTTP server + - contains: "expose the vllm server via" + exception: GenerationError + error_details: "The NeMo-Gym (Penguin) environment requires the vLLM server to be exposed via HTTP (expose_http_server: true). This allows the environment to communicate with the generation model through an HTTP API. This is an internal configuration issue with the training environment, please contact the administrator." + + # NeMo-Gym/Penguin incompatible with reasoning parser + - contains: "Please do not use a reasoning parser in vLLM" + exception: GenerationError + error_details: "The NeMo-Gym (Penguin) environment is incompatible with vLLM's reasoning parser. NeMo-Gym handles all data processing including reasoning traces itself, so having a reasoning parser in vLLM would cause conflicts. This is an internal configuration issue with the training environment, please contact the administrator." + + # No placement groups available for vLLM + - exact: "No placement groups available in the cluster" + exception: GenerationError + error_details: "No Ray placement groups are available for vLLM generation workers. This means the cluster could not allocate the required GPU resources for the generation component of training. This is typically a resource availability issue - please try again or contact your administrator to check cluster capacity." + + # Unable to allocate vLLM worker groups + - contains: "Unable to allocate any worker groups with the available resources" + exception: GenerationError + error_details: "Could not allocate any vLLM worker groups with the available cluster resources. The generation component of DPO/GRPO training requires dedicated GPU resources for vLLM inference workers. Please ensure the cluster has enough GPUs, or reduce the generation parallelism settings." + + # Placement group contains no bundles + - exact: "Placement group contains no bundles" + exception: GenerationError + error_details: "A Ray placement group allocated for vLLM generation workers contains no resource bundles. This indicates an issue with cluster resource allocation. This is an infrastructure issue - please contact your administrator." + + # Failed to retrieve bundle/node mapping from placement group + - contains: "Failed to retrieve bundle/node mapping from placement group" + exception: GenerationError + error_details: "Could not retrieve the bundle-to-node mapping from the Ray placement group for vLLM workers. This indicates an issue with the distributed training cluster setup. This is an infrastructure issue - please contact your administrator." + + # No output received for generation request + - regex: "^No output received for request .+$" + exception: GenerationError + error_details: "The vLLM async generation engine did not produce any output for a generation request. This can happen when: 1) The generation request timed out, 2) The vLLM worker encountered an internal error, or 3) GPU memory was exhausted during generation. This is typically a transient issue - please try running your training job again." + + # =========================================================================== + # 14. INTERNAL ERRORS (500) + # =========================================================================== + + # --- Automodel Pipeline Parallelism Errors --- + # Pipeline parallelism: first stage missing inputs + - exact: "You must provide either input_ids or inputs_embeds" + exception: InternalError + error_details: "Pipeline parallelism internal error: the first pipeline stage did not receive input data (input_ids or inputs_embeds). This is an internal configuration issue with how the model is split across pipeline stages, please reach out to the NeMo Platform team." + + # Pipeline parallelism: intermediate stage missing embeddings + - exact: "inputs_embeds must be provided for pipeline stages without embed_tokens" + exception: InternalError + error_details: "Pipeline parallelism internal error: an intermediate pipeline stage did not receive embeddings from the previous stage. In pipeline parallelism, each stage processes a portion of the model layers and passes activations to the next stage. This error indicates the inter-stage communication failed, and is an internal training configuration issue, please reach out to the NeMo Platform team." + + # --- Automodel MoE (Mixture of Experts) Errors --- + # MoE: only 1D mesh supported (occurs when TP+EP are both > 1) + - exact: "We only support 1D mesh for MoE" + exception: ParallelismConfigError + error_details: "MoE (Mixture of Experts) models do not support combining tensor parallelism with expert parallelism. When using expert_model_parallel_size > 1, you must set tensor_parallel_size=1. Please update your parallelism configuration to disable tensor parallelism for MoE training." + + # MoE: DTensor placement error (checkpoint/parallelism mismatch) + - contains: "has unsupported DTensor placement" + exception: ParallelismConfigError + error_details: "MoE (Mixture of Experts) model checkpoint has an incompatible tensor distribution for the current expert parallelism settings. This typically occurs when the base model checkpoint was saved with different expert_model_parallel_size than what you're using for training. Please ensure your expert_model_parallel_size matches how the base model was originally distributed, or use a checkpoint that was saved without expert parallelism (expert_model_parallel_size=1)." + + # --- Automodel Fused Optimization Errors --- + # FusedLinearCrossEntropy configuration + - contains: "FusedLinearCrossEntropy requires the model to output hidden states" + exception: InternalError + error_details: "The fused linear cross-entropy optimization requires the model to output hidden states, but the model is configured to only output logits. FusedLinearCrossEntropy is a memory optimization that combines the final linear projection and loss computation. This is an internal configuration issue, contact the NeMo Platform team." + + # --- NeMo-RL Async GRPO Errors --- + # Stale trajectories in replay buffer + - regex: "^Found \\d+ trajectories older than min_valid_version \\d+$" + exception: InternalError + error_details: "The async GRPO replay buffer contains stale trajectories that are older than the minimum valid version. In async GRPO, trajectories are generated asynchronously and stored in a replay buffer. Stale trajectories can cause training instability because they were generated by an outdated policy. This indicates a synchronization issue between generation and training workers. Please contact the administrator." + + # --- NeMo-RL Tensor Processing Errors --- + # Tensor dimension mismatch + - regex: "^tensors for .+ must have same number of dimensions" + exception: InternalError + error_details: "Tensors being processed have mismatched dimensions during internal batching. This is an internal data processing issue that should not occur with valid datasets. Please contact the NeMo Platform team with your dataset format details." + + # Tensor dtype mismatch + - contains: "expected consistent types but got:" + exception: InternalError + error_details: "Tensors being processed have inconsistent data types (dtypes) during internal batching. This is an internal data processing issue that should not occur with valid datasets. Please contact the NeMo Platform team." + + # Tensors on different devices + - contains: "expected tensors on the same device but got:" + exception: InternalError + error_details: "Tensors are located on different devices during internal processing. This is an internal distributed training issue. Please contact the NeMo Platform team." + + # --- Automodel Configuration Errors --- + # Config instantiation failure (from ConfigNode.instantiate()) + # This prints a detailed error with "Instantiation failed for `func_name`" + - contains: "Instantiation failed for" + exception: InternalError + error_details: "Failed to instantiate a training configuration component. The training system uses a configuration tree where each node can instantiate Python objects (like optimizers, schedulers, or model components). This error means one of these instantiations failed, possibly due to invalid parameters or missing dependencies. Please contact the administrator." + + # Model compilation failure + - contains: "Model compilation failed" + exception: InternalError + error_details: "PyTorch model compilation (torch.compile) failed. Model compilation is an optional optimization that can speed up training by compiling the model graph. Training will fall back to eager mode and continue without compilation. If this error persists, it may indicate an incompatibility between the model architecture and PyTorch's compiler. Please contact the administrator if training fails." + + # --- General Training Process Errors --- + # Training subprocess error (generic fallback when no specific error was parsed) + # Matches both parser ("Training failed with exit code: X") and train.py ("Training subprocess returned with error code: X") + - regex: "^Training (failed with exit code|subprocess returned with (?:error )?code):? \\d+.*" + exception: InternalError + error_details: "The training process exited with a non-zero exit code, but no specific error message could be extracted from the training logs. This is a generic failure that can have many causes. Please check the full training logs for more details, and contact the administrator if you cannot determine the cause." + diff --git a/services/rl/src/nmp/rl/tasks/training/errors/exceptions.py b/services/rl/src/nmp/rl/tasks/training/errors/exceptions.py new file mode 100644 index 0000000000..6a98527051 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/errors/exceptions.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Custom exceptions for Customizer training errors. + +These exceptions provide user-friendly error messages for errors that may occur +during training with various backends: +- Automodel +- NeMo-RL +- Megatron Bridge +""" + +from dataclasses import dataclass +from typing import TypedDict + + +def format_exception_string(exc: BaseException) -> str: + """Format an exception as ``TypeName: message`` matching Python's traceback style. + + This is the canonical format used throughout the error-handling pipeline: + - ``ray_bootstrap`` writes it into the driver output buffer so the parser + can extract exceptions that occurred outside the subprocess. + - ``default_exception_handler`` uses it for the ``detail`` field reported + to the Jobs service. + - The parser's ``_EXCEPTION_RE`` regex is designed to match this format + when reading subprocess output. + """ + return f"{type(exc).__name__}: {exc}" + + +class ErrorDetails(TypedDict): + """Error details dict for Jobs service reporting.""" + + message: str + type: str + detail: str | None + + +@dataclass +class CustomizerTrainingError(Exception): + """ + Base exception for Customizer training errors. + + Attributes: + message: User-friendly error message shown to the user. + detail: Technical details about the original error (for debugging). + user_message: Class-level default message used as fallback when the YAML rule + does not specify an `error_details` field. Subclasses override this. + """ + + message: str + detail: str | None = None + + # Default user-facing message - subclasses override this. + # Used as fallback when YAML rule omits `error_details` field. + # See default_exception_handler() for usage. + user_message: str = "An error occurred during training." + + def __post_init__(self): + # Call Exception.__init__ with the message + super().__init__(self.message) + + def __str__(self) -> str: + return self.message + + def to_error_details(self) -> ErrorDetails: + """Convert to error_details dict for Jobs service reporting.""" + return ErrorDetails( + message=self.message, + type=type(self).__name__, + detail=self.detail, + ) + + +# ============================================================================= +# CLIENT ERRORS (400) +# ============================================================================= + + +@dataclass +class DatasetFormatError(CustomizerTrainingError): + """ + Dataset has invalid format or schema. + + Raised when: + - Dataset sample has unsupported role (not system/user/assistant/tool) + - Dataset is empty or has zero valid samples + - Text input is not a string or list of strings + - Required field missing from dataset sample + - Prompt file does not exist + """ + + user_message: str = "Dataset format error. Please check your dataset matches the expected schema." + + +@dataclass +class TrainingConfigError(CustomizerTrainingError): + """ + Invalid training configuration. + + Raised when: + - Model incompatible with pipeline parallelism (tied embeddings, encoder-decoder) + - PP batch/microbatch configuration invalid + - Model doesn't support SDPA for context parallelism + - Triton not installed for optimized LoRA kernels + - LoRA adapter dimensions mismatch + - DPO with dynamic batching or sequence packing + - GRPO missing generation config or validation dataset + - Async GRPO configuration errors + - Batch size not divisible by data parallel size + - World size insufficient for parallelism configuration + """ + + user_message: str = ( + "Training configuration error. Please check your parallelism settings " + "(tensor_parallel_size, pipeline_parallel_size, expert_model_parallel_size), " + "batch settings (batch_size, micro_batch_size), or training type configuration." + ) + + +@dataclass +class TrainingEnvironmentError(CustomizerTrainingError): + """ + Invalid environment configuration for GRPO. + + Raised when: + - GRPO environment name is not recognized + - GRPO environment not configured + - No environment found for task type + """ + + user_message: str = "Environment configuration error. Please check your GRPO environment settings." + + +@dataclass +class ParallelismConfigError(CustomizerTrainingError): + """ + Invalid parallelism configuration for MoE models. + + Raised when: + - MoE model uses tensor parallelism with expert parallelism (only 1D mesh supported) + - DTensor placement incompatible with expert parallelism settings + - Checkpoint parallelism settings don't match training configuration + """ + + user_message: str = ( + "Parallelism configuration error for Mixture-of-Experts (MoE) model. " + "MoE models do not support combining tensor_parallel_size > 1 with expert_model_parallel_size > 1. " + "To fix: either set tensor_parallel_size=1 when using expert parallelism, " + "or set expert_model_parallel_size=1 when using tensor parallelism." + ) + + +# ============================================================================= +# NOT FOUND ERRORS (404) +# ============================================================================= + + +@dataclass +class ModelNotFoundError(CustomizerTrainingError): + """ + Model or checkpoint path doesn't exist. + + Raised when: + - The specified checkpoint path does not exist + - The checkpoint directory is empty when resuming + - Nemotron model missing required HF source code + """ + + user_message: str = ( + "Model or checkpoint not found. The specified model path does not exist or is inaccessible. " + "Please verify the model identifier is correct and the model was successfully downloaded." + ) + + +# ============================================================================= +# SERVER ERRORS (500) +# ============================================================================= + + +@dataclass +class ModelLoadError(CustomizerTrainingError): + """ + Failed to load or initialize model. + + Raised when: + - Model weights could not be applied to a layer (corruption) + - Model optimizations/patches failed + - Method signature mismatch during patching + - Missing lm_head.weight in model + - vLLM library not installed + - Shape mismatch for model parameters or buffers + - Generation output missing required fields + """ + + user_message: str = ( + "Failed to load the model. This can happen when: " + "1) The model checkpoint is corrupted or incomplete, " + "2) The model architecture is incompatible with the training configuration, " + "3) There is a version mismatch between the model and the training framework. " + "Please verify the model checkpoint is valid and complete." + ) + + +@dataclass +class CheckpointError(CustomizerTrainingError): + """ + Checkpoint save or load failure. + + Raised when: + - Checkpoint directory already exists + - Failed to validate global plan (distributed checkpoint corruption) + - Missing key in checkpoint state_dict + - Expert weights missing from MoE checkpoint + - Training interrupted during checkpoint save + - Parallelism settings don't match checkpoint + - Model export or upload failed + """ + + user_message: str = ( + "Checkpoint save or load failed. This can happen when: " + "1) The checkpoint is corrupted or was saved incompletely (e.g., training was interrupted), " + "2) Disk space is insufficient for saving checkpoints, " + "3) The base model checkpoint is incompatible with the current training configuration." + ) + + +@dataclass +class CudaError(CustomizerTrainingError): + """ + GPU/CUDA runtime error. + + Raised when: + - GPU out of memory (OOM) + - General CUDA runtime errors + """ + + user_message: str = ( + "GPU memory exhausted. To reduce memory usage: " + "1) Reduce batch_size or micro_batch_size, " + "2) Reduce max_seq_length, " + "3) Use LoRA fine-tuning instead of full fine-tuning, " + "4) Increase tensor_parallel_size to distribute the model across more GPUs." + ) + + +@dataclass +class DistributedError(CustomizerTrainingError): + """ + Distributed training or Ray cluster failure. + + Raised when: + - torch.distributed not available + - torch.distributed not initialized + - Distributed operation timeout + - NCCL communication errors + - Ray cluster resource insufficiency + - Placement group allocation failure + """ + + user_message: str = "Distributed training error. Please check cluster resources and try again." + + +@dataclass +class GenerationError(CustomizerTrainingError): + """ + vLLM generation/inference failure. + + Raised when: + - Failed to update vLLM weights from training policy + - Sync method called on async engine + - Error during rollout for a sample + - Async generation called without async engine + - Penguin requires async vLLM + """ + + user_message: str = ( + "Generation error during reinforcement learning training. " + "DPO and GRPO training generate model responses during the training loop to compute rewards. " + "This error indicates the generation step failed, which may be caused by vLLM backend issues " + "or incompatible generation settings." + ) + + +@dataclass +class TrainingTimeoutError(CustomizerTrainingError): + """ + Training exceeded time limit. + + Raised when: + - Training subprocess exceeded configured timeout + """ + + user_message: str = ( + "Training exceeded the maximum allowed time limit. " + "To reduce training time: reduce epochs or max_steps, use a smaller dataset, " + "use a smaller model, or use LoRA fine-tuning instead of full fine-tuning. " + "Contact your administrator if you need longer training time limits." + ) + + +@dataclass +class InternalError(CustomizerTrainingError): + """ + Unexpected internal error. + + Raised when: + - Pipeline stage missing input_ids or inputs_embeds + - MoE device mesh configuration error + - DTensor placement error for expert parallelism + - FusedLinearCrossEntropy configuration error + - Tensor dimension/dtype/device mismatch + - Logger misconfiguration + - Any unmatched error (fallback) + """ + + user_message: str = ( + "An unexpected internal error occurred during training. " + "This is typically caused by framework-level issues such as tensor misconfigurations, " + "device mesh errors, or internal pipeline failures. " + "Please try running your job again. If the issue persists, contact your administrator " + "with the job ID and error details for further investigation." + ) + + +@dataclass +class GenericTrainingError(CustomizerTrainingError): + """ + Fallback when error classification is ambiguous. + + Used when multiple error rules match the same exception, + making classification unreliable. + """ + + user_message: str = ( + "Training failed due to an error that could not be precisely categorized. " + "Please review the error details for more information. " + "If the issue persists, try adjusting your training configuration." + ) + + +# ============================================================================= +# EXCEPTION REGISTRY +# ============================================================================= + +# Maps exception class names (strings in YAML) to actual Python classes +EXCEPTION_REGISTRY: dict[str, type[Exception]] = { + # Base + "CustomizerTrainingError": CustomizerTrainingError, + # Client errors (400) + "DatasetFormatError": DatasetFormatError, + "TrainingConfigError": TrainingConfigError, + "TrainingEnvironmentError": TrainingEnvironmentError, + "ParallelismConfigError": ParallelismConfigError, + # Not found (404) + "ModelNotFoundError": ModelNotFoundError, + # Server errors (500) + "ModelLoadError": ModelLoadError, + "CheckpointError": CheckpointError, + "CudaError": CudaError, + "DistributedError": DistributedError, + "GenerationError": GenerationError, + "TrainingTimeoutError": TrainingTimeoutError, + "InternalError": InternalError, + "GenericTrainingError": GenericTrainingError, +} + + +# ============================================================================= +# DEFAULT EXCEPTION HANDLER +# ============================================================================= + + +def default_exception_handler( + exception_class: type[Exception], + original_exception: Exception, + error_details: str | None, +) -> Exception: + """ + Default handler for creating Customizer training exceptions. + + This handler is used by RulesLoader when: + 1. A rule matches but doesn't have a custom handler + 2. No rule matches and fallback_exception is set + + Args: + exception_class: The exception class to create (from EXCEPTION_REGISTRY) + original_exception: The original exception that was caught + error_details: User-friendly message from the rule's error_details field, + or None if not specified + + Returns: + A new instance of exception_class with appropriate message and detail + """ + # Get the default user message from the class if no error_details provided + if issubclass(exception_class, CustomizerTrainingError): + user_message = error_details or exception_class.user_message + # For InternalError fallback (no matching rule), include the original error + # in the message so users get actionable information instead of a vague message + if exception_class is InternalError and error_details is None: + user_message = f"{user_message} ({format_exception_string(original_exception)})" + return exception_class( + message=user_message, + detail=format_exception_string(original_exception), + ) + else: + # For non-CustomizerTrainingError classes (shouldn't happen, but be safe) + return exception_class(error_details or str(original_exception)) + + +__all__ = [ + "CheckpointError", + "CudaError", + "CustomizerTrainingError", + "DatasetFormatError", + "DistributedError", + "ErrorDetails", + "EXCEPTION_REGISTRY", + "format_exception_string", + "GenerationError", + "GenericTrainingError", + "InternalError", + "ModelLoadError", + "ModelNotFoundError", + "ParallelismConfigError", + "TrainingConfigError", + "TrainingEnvironmentError", + "TrainingTimeoutError", + "default_exception_handler", +] diff --git a/services/rl/src/nmp/rl/tasks/training/errors/parser.py b/services/rl/src/nmp/rl/tasks/training/errors/parser.py new file mode 100644 index 0000000000..66d2df17a9 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/errors/parser.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Error parser for subprocess output. + +This module provides utilities to parse and extract meaningful error messages from +training subprocess output (stdout/stderr). It should be used by all training backends +(Automodel, NeMo-RL, Megatron Bridge) to capture errors for classification. + +The extracted error messages are then matched against YAML rules by the +error converter to produce user-friendly error messages. +""" + +import re +import subprocess +import sys +from collections import deque +from dataclasses import dataclass + +# Number of recent output lines to keep for error parsing +MAX_OUTPUT_LINES = 500 + +# Patterns that indicate an error line (case-insensitive search) +# These match Python exception types and common error patterns from training libraries +ERROR_INDICATORS = [ + # Python exception type names (appear as "ExceptionType: message") + "runtimeerror", + "valueerror", + "assertionerror", + "importerror", + "attributeerror", + "keyerror", + "typeerror", + "filenotfounderror", + "permissionerror", + "oserror", + "ioerror", + # Generic error patterns + "error:", + "exception:", + "traceback", + # Automodel-specific patterns + "instantiation failed", # From ConfigNode.instantiate() + "model compilation failed", # From compile_utils.py + # NeMo-RL patterns + "ray error", + "actor died", + "worker crashed", + # Megatron Bridge patterns + "nemo error", + "lightning error", + # CUDA/GPU patterns + "cuda out of memory", + "out of memory", + "oom", + "cuda error", + "cublas error", + "cudnn error", + # Distributed training patterns + "nccl", + "gloo", + "distributed", + "mpi error", + # General failure patterns + "failed", + "failure", + "abort", + "killed", + "segmentation fault", + "signal", +] + +# Regex to detect Python exception lines and extract the type name (group 1) +# and message (group 2) as separate captures. The type name may be module +# qualified (e.g. "subprocess.TimeoutExpired", "torch.cuda.OutOfMemoryError"); +# _extract_exception strips the module prefix down to the bare class name. +# The suffix set deliberately goes beyond Error/Exception so non-"*Error" +# exceptions like TimeoutExpired / KeyboardInterrupt / SystemExit are captured. +_EXCEPTION_RE = re.compile( + r"\b((?:[A-Za-z_]\w*\.)*[A-Z]\w*(?:Error|Exception|Expired|Interrupt|Exit)):\s*(.*)", +) + +# Wrapper exceptions from distributed training - skip these to find root cause +WRAPPER_EXCEPTION_PATTERNS = [ + "childfailederror", # torch.distributed wrapper + "torch.distributed.elastic", # torch elastic wrapper + "multiprocessing.errors", # multiprocessing wrapper +] + + +@dataclass(frozen=True) +class ParsedError: + """Error extracted from subprocess output. + + Preserves both the original exception type name (as printed in the + traceback) and the message, so callers can reconstruct a typed + exception for the converter's type-based matchers. + """ + + exception_type: str + message: str + + def to_exception(self) -> Exception: + """Reconstruct an exception that preserves the original type name. + + Dynamically creates an exception class whose ``__name__`` matches + the original type (e.g. ``ValueError``, ``ResourceInsufficientError``) + so that ``type_name`` YAML matchers can match it. The class inherits + from ``RuntimeError`` so that standard ``except Exception`` handling + works without needing the real library class to be importable. + """ + exc_class = type(self.exception_type, (RuntimeError,), {}) + return exc_class(self.message) + + +def _clean_line(line: str) -> str: + """Remove common prefixes like [rank0]: from distributed output.""" + line = re.sub(r"^\[rank\d+\]:\s*", "", line.strip()) + return line.strip() + + +def _is_wrapper_exception(line: str) -> bool: + """Check if this is a wrapper exception that should be skipped.""" + line_lower = line.lower() + return any(pattern in line_lower for pattern in WRAPPER_EXCEPTION_PATTERNS) + + +def _extract_exception(line: str) -> ParsedError | None: + """ + Extract the exception type and message from a subprocess output line. + + Examples: + >>> _extract_exception("[rank0]: ValueError: invalid input") + ParsedError(exception_type='ValueError', message='invalid input') + >>> _extract_exception("torch.cuda.OutOfMemoryError: CUDA OOM") + ParsedError(exception_type='OutOfMemoryError', message='CUDA OOM') + >>> _extract_exception("subprocess.TimeoutExpired: Command timed out") + ParsedError(exception_type='TimeoutExpired', message='Command timed out') + >>> _extract_exception(" File 'train.py', line 42") + None + >>> _extract_exception("ChildFailedError: worker 0 failed") + None # Wrapper exception, skipped + + Returns None for non-exception lines and wrapper exceptions. + """ + if _is_wrapper_exception(line): + return None + + match = _EXCEPTION_RE.search(line) + if match: + # Strip any module prefix (e.g. "subprocess.TimeoutExpired" -> "TimeoutExpired") + # so the bare class name matches the converter's type_name rules. + exc_type = match.group(1).strip().rsplit(".", 1)[-1] + message = match.group(2).strip() if match.group(2) else "" + return ParsedError( + exception_type=exc_type, + message=message or exc_type, + ) + + return None + + +def parse_error_from_output(output_lines: deque, returncode: int) -> ParsedError: + """ + Parse subprocess output and extract a structured error. + + Searches the captured output for Python exception lines and returns a + ``ParsedError`` preserving both the exception type name and message. + Callers use ``result.to_exception()`` to reconstruct a typed exception + that works with both message-based *and* type-based YAML matchers. + + Strategy: + 1. Find the LAST Python exception line (e.g., "ValueError: message") + 2. Extract the type name and message separately + 3. Deduplicate across distributed ranks + + Args: + output_lines: Rolling buffer of recent output lines. + returncode: Process exit code. + + Returns: + ParsedError with exception_type and message. + """ + if not output_lines: + return ParsedError("RuntimeError", f"Training failed with exit code: {returncode}") + + lines = list(output_lines) + + # Search backwards for exception lines and collect unique ones + # (distributed training often prints the same error multiple times) + found: list[ParsedError] = [] + seen_messages: set[str] = set() + + for i in range(len(lines) - 1, -1, -1): + parsed = _extract_exception(lines[i]) + if parsed and parsed.message not in seen_messages: + seen_messages.add(parsed.message) + found.append(parsed) + if len(found) >= 3: + break + + if found: + return found[0] + + # Fallback: search for any error-related lines + error_lines: list[str] = [] + for line in reversed(lines): + line_lower = line.lower() + is_error_line = any(indicator in line_lower for indicator in ERROR_INDICATORS) + if is_error_line: + cleaned = _clean_line(line) + if cleaned and cleaned not in error_lines: + error_lines.insert(0, cleaned) + if len(error_lines) > 10: + break + + if error_lines: + return ParsedError("RuntimeError", "\n".join(error_lines[-10:])) + + # Last resort: return last N lines of output + last_lines = [_clean_line(line) for line in lines[-10:]] + message = f"Training failed with exit code {returncode}. Last output:\n" + "\n".join(last_lines) + return ParsedError("RuntimeError", message) + + +def read_subprocess_output(proc: subprocess.Popen, buffer: deque) -> None: + """ + Read subprocess output, stream to console, and capture in buffer. + + This function is designed to run in a daemon thread alongside a subprocess, + reading its stdout line-by-line, printing to console in real-time, and + storing lines in a rolling buffer for later error extraction. + + Args: + proc: The subprocess.Popen object with stdout=PIPE. + buffer: A deque with maxlen to store recent output lines. + """ + if proc.stdout is None: + return + + try: + for line in iter(proc.stdout.readline, ""): + if not line: + break + # Stream to console + sys.stdout.write(line) + sys.stdout.flush() + # Capture in rolling buffer + buffer.append(line.rstrip("\n")) + except (ValueError, OSError): + # Process closed or pipe broken + pass + + +__all__ = [ + "ERROR_INDICATORS", + "MAX_OUTPUT_LINES", + "ParsedError", + "parse_error_from_output", + "read_subprocess_output", +] diff --git a/services/rl/src/nmp/rl/tasks/training/integrations.py b/services/rl/src/nmp/rl/tasks/training/integrations.py new file mode 100644 index 0000000000..a9c67ba23b --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/integrations.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Shared integration config helpers for training backends.""" + +import logging +import os +from pathlib import Path +from typing import Any + +from nmp.customization_common.service.context import NMPJobContext +from nmp.rl.app.jobs.training.schemas import TrainingStepConfig + +logger = logging.getLogger(__name__) + + +def _resolve_with_fallback( + primary: str | None, + fallback: str | None, + default: str, + field_label: str | None = None, +) -> str: + """Pick the first truthy value from *primary* → *fallback* → *default*. + + When *field_label* is given and neither *primary* nor *fallback* is set, + a warning is logged so operators know a hardcoded default is in use. + """ + if field_label and not (primary or fallback): + logger.warning(f"{field_label} is not set; using fallback '{default}'.") + return primary or fallback or default + + +def build_mlflow_config( + customizer_config: TrainingStepConfig, + job_ctx: NMPJobContext, + framework: str, +) -> dict[str, Any] | None: + """Build MLflow config shared across training backends. + The resulting dict is passed to MLflow logging setup by a backend. + + Run naming strategy (same as WandB): + - run_name uses job_id (stable across pause/resume) + - task_id is added to tags for granular execution tracking + + Missing tracking URI disables integration with a warning. + """ + user_config = customizer_config.integrations.mlflow + if not user_config: + return None + + # User-provided tracking URI takes precedence over environment variable + tracking_uri = user_config.tracking_uri or os.environ.get("MLFLOW_TRACKING_URI") + if not tracking_uri: + logger.warning( + "MLflow integration is configured but no tracking URI is set " + "(MLFLOW_TRACKING_URI env var and integrations.mlflow.tracking_uri in job POST request are empty); " + "MLflow integration will be disabled." + ) + return None + + tags: dict[str, str] = { + "service": "rl", + "framework": framework, + } + if job_ctx.workspace: + tags["workspace"] = job_ctx.workspace + if job_ctx.job_id: + tags["job"] = job_ctx.job_id + if job_ctx.task: + tags["task"] = job_ctx.task + if customizer_config.model.name: + tags["model_name"] = customizer_config.model.name + + # User-provided tags override defaults above + if user_config.tags: + tags.update(user_config.tags) + if user_config.description: + # MLflow run description is stored in the reserved `mlflow.note.content` tag. + # See: https://mlflow.org/docs/latest/ml/tracking/#how-to-include-additional-description-texts-about-the-run + tags["mlflow.note.content"] = user_config.description + + experiment_name = _resolve_with_fallback( + user_config.experiment_name, + customizer_config.output_model, + "default-experiment", + field_label="MLflow experiment_name", + ) + run_name = _resolve_with_fallback( + user_config.run_name, + job_ctx.job_id, + "default-run", + field_label="MLflow run_name", + ) + + mlflow_config: dict[str, Any] = { + "tracking_uri": tracking_uri, + "experiment_name": experiment_name, + "run_name": run_name, + "tags": tags, + } + + return mlflow_config + + +def build_wandb_config( + customizer_config: TrainingStepConfig, + job_ctx: NMPJobContext, + framework: str, +) -> dict[str, Any] | None: + """Build WandB config shared across training backends. + + The resulting dict is passed to wandb.init() as kwargs. + See: https://docs.wandb.ai/ref/python/init + + TODO: Add pause/resume support: + - 'name' and 'id' use job_id (stable across pause/resume) + - 'resume="allow"' enables continuing runs after pause/resume + """ + user_config = customizer_config.integrations.wandb + if not user_config: + return None + + wandb_api_key = os.environ.get("WANDB_API_KEY") + if not user_config.base_url and not wandb_api_key: + logger.warning("WandB API key is not set and no base_url is provided, skipping WandB integration") + return None + + # Note: This is semantically different from job_ctx.workspace. + # This is the workspace for training artifacts. + run_dir = Path(customizer_config.workspace_path) / "wandb" + + tags: list[str] = ["service:rl", f"framework:{framework}"] + if job_ctx.workspace: + tags.append(f"workspace:{job_ctx.workspace}") + if job_ctx.job_id: + tags.append(f"job:{job_ctx.job_id}") + if job_ctx.task: + tags.append(f"task:{job_ctx.task}") + if customizer_config.model.name: + tags.append(f"model:{customizer_config.model.name}") + # User-provided tags are appended (can override tags above) + if user_config.tags: + tags.extend(user_config.tags) + + wandb_config: dict[str, Any] = { + "project": _resolve_with_fallback(user_config.project, customizer_config.output_model, "default-project"), + "name": _resolve_with_fallback(user_config.name, job_ctx.job_id, "default-run"), + "dir": str(run_dir), + "tags": tags, + } + if user_config.entity: + wandb_config["entity"] = user_config.entity + if user_config.notes: + wandb_config["notes"] = user_config.notes + if user_config.base_url: + # For self-hosted W&B servers, base_url is passed via the settings dict + # (wandb.init accepts settings as Union[Settings, Dict[str, Any], None]). + logger.info(f"Using self-hosted W&B server: {user_config.base_url}") + wandb_config["settings"] = {"base_url": user_config.base_url} + + return wandb_config diff --git a/services/rl/src/nmp/rl/tasks/training/protocol.py b/services/rl/src/nmp/rl/tasks/training/protocol.py new file mode 100644 index 0000000000..80516c0e2e --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/protocol.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional, Protocol, runtime_checkable + +from nmp.customization_common.training.progress import JobsServiceProgressReporter +from nmp.rl.app.jobs.training.schemas import ( + CheckpointInfo, + TrainingMetrics, + TrainingStepConfig, +) +from nmp.rl.app.jobs.training.schemas import TrainingBackend as TrainingBackendEnum + + +@dataclass +class LibraryConfig: + """ + Library-specific configuration ready for training. + + We track both the config dict and the path, and let the consumer decide how to use them. + """ + + config_dict: dict[str, Any] # Library-specific config dict + config_path: Path # Path to the config file (managed by runner) + + +@runtime_checkable +class SupportsPreprocessing(Protocol): + """Protocol for backends that need pre-training preprocessing. + + Backends that implement this protocol will have their `run_preprocessing` + method called before config compilation. Use this for operations like + model format conversion that must happen before training. + """ + + def run_preprocessing( + self, + customizer_config: TrainingStepConfig, + ) -> None: + """Run pre-training conversions (e.g., model format conversion). + + Called before config compilation on the coordinator node only. + + Args: + customizer_config: Standardized training configuration + """ + ... + + +@runtime_checkable +class TrainingBackend(Protocol): + """ + Interface for training backends (Strategy Pattern). + + Each backend (e.g. nemo_rl) implements this interface. + Backends are responsible for: + + 1. Compiling library-specific configuration (pure transformation) + 2. Executing training using library-specific wrappers/recipes + 3. Processing checkpoints to standard output format + + Note: Pre-training conversions are optional. Backends that need them + should also implement `SupportsPreprocessing`. + """ + + @property + def backend_type(self) -> TrainingBackendEnum: + """Backend type identifier.""" + ... + + def compile_config( + self, + config: TrainingStepConfig, + workspace_dir: Path, + ) -> dict[str, Any]: + """ + Transform standardized config to library-specific config. + + This is a pure transformation - no file I/O. The runner handles + writing the config to disk. + + Called by the coordinator node only. + + Args: + config: Standardized training configuration + workspace_dir: Directory for training artifacts (for paths in config) + + Returns: + Library-specific config dict ready to be written as YAML + """ + ... + + def execute_training( + self, + customizer_config: TrainingStepConfig, + library_config: LibraryConfig, + progress: JobsServiceProgressReporter, + ) -> TrainingMetrics: + """ + Execute training using library-specific wrappers. + """ + ... + + def find_best_checkpoint( + self, + workspace_dir: Path, + customizer_config: TrainingStepConfig, + library_config: Optional[LibraryConfig] = None, + ) -> Path: + """ + Find the best checkpoint after training. + """ + ... + + def process_checkpoint( + self, + checkpoint_path: Path, + output_path: Path, + customizer_config: TrainingStepConfig, + library_config: LibraryConfig | None = None, + ) -> CheckpointInfo: + """ + Process checkpoint to standard output format. + + Args: + checkpoint_path: Path to the checkpoint directory + output_path: Where to write the processed checkpoint + customizer_config: Training configuration + library_config: Library-specific config (contains resolved chat template, etc.) + """ + ... diff --git a/services/rl/src/nmp/rl/tasks/training/runner.py b/services/rl/src/nmp/rl/tasks/training/runner.py new file mode 100644 index 0000000000..e9573dc68d --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/runner.py @@ -0,0 +1,321 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Training runner with distributed coordination support. + +Orchestrates training execution across single-node and multi-node environments, +using file-based barriers for cross-pod synchronization. +""" + +import json +import logging +import random +import time +from enum import Enum +from pathlib import Path +from types import TracebackType +from typing import Self + +import yaml +from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.training.progress import JobsServiceProgressReporter +from nmp.rl.app.constants import DEFAULT_TRAINING_RESULT_FILE_NAME, SERVICE_NAME +from nmp.rl.app.jobs.training.schemas import ( + GPUInfo, + TrainingMetrics, + TrainingResult, + TrainingStepConfig, +) +from nmp.rl.app.jobs.training.schemas import TrainingBackend as TrainingBackendEnum + +from .distributed import DistributedContext +from .errors.converter import create_error_details +from .protocol import LibraryConfig, SupportsPreprocessing, TrainingBackend +from .utils import get_gpu_info + + +# Custom YAML representer to serialize Enum values as their string values +def _enum_representer(dumper: yaml.Dumper, data: Enum) -> yaml.Node: + """Represent Enum as its value (string) rather than a Python object tag.""" + return dumper.represent_str(str(data.value)) + + +yaml.add_representer(Enum, _enum_representer) +# Also add for all Enum subclasses +yaml.add_multi_representer(Enum, _enum_representer) + +logger = logging.getLogger(__name__) + +# Barrier names for distributed synchronization +BARRIER_CONFIG_READY = "config_ready" +BARRIER_TRAINING_COMPLETE = "training_complete" +BARRIER_PREPROCESSING_COMPLETE = "preprocessing_complete" + + +class TrainingRunner: + """ + Orchestrates training execution across single-node and multi-node environments. + + Initializes from environment variables and coordinates training phases: + - Config compilation: Coordinator only, workers wait + - Training: All ranks participate (via torchrun) + - Post-processing: Coordinator only, workers exit after training sync + + Usage: + with TrainingRunner() as runner: + result = runner.run() + """ + + def __init__(self, backend: TrainingBackend | None = None) -> None: + """Initialize the runner from environment variables.""" + self._job_ctx = NMPJobContext.from_env() + + self._config = self._load_config(self._job_ctx.config_path) + self._progress = JobsServiceProgressReporter(self._job_ctx, SERVICE_NAME) + self._dist_ctx = DistributedContext.from_env(self._get_barrier_dir()) + self._backend = backend or self._load_backend(self._config.backend) + # workspace_path and output_path are absolute paths from the config + self._workspace_path = Path(self._config.workspace_path) + self._output_path = Path(self._config.output_path) + + def __enter__(self) -> Self: + """Context manager entry.""" + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Context manager exit - ensures progress reporter is closed.""" + self.close() + + def close(self) -> None: + """Clean up resources (progress reporter).""" + self._progress.close() + + # --- Training execution --- + + def run(self) -> TrainingResult: + """ + Execute training with distributed coordination. + + Phases: + 1. Config compilation (coordinator, workers wait) + 2. Training (all ranks via torchrun) + 3. Sync point (all ranks) — workers return success and exit here + 4. Post-processing (coordinator only) + 5. Result writing (coordinator only) + + Returns: + TrainingResult with success/failure and metrics + """ + # Set global seed as first layer of defense for reproducibility + random.seed(self._config.seed) + logger.info(f"Global random seed set to {self._config.seed}") + + start_time = time.time() + gpu_info = get_gpu_info() + result = TrainingResult(success=False, error_message="No result") + + try: + # === Phase 0: Pre-training conversions (coordinator only) === + self._preprocessing_phase() + + # === Phase 1: Config compilation (coordinator, workers wait) === + library_config = self._compile_config_phase() + + # === Phase 2: Training (all ranks) === + metrics = self._training_phase(library_config) + + # === Phase 3: Sync after training === + self._dist_ctx.sync_point(BARRIER_TRAINING_COMPLETE) + + # === Phase 4: Post-processing (coordinator only, workers exit) === + result = self._postprocess_phase(gpu_info, metrics, start_time, library_config) + + except Exception as e: + logger.exception(f"Training failed: {e}") + # Convert exception to user-friendly error details using error mapping rules + error_details = create_error_details(e) + result = TrainingResult( + success=False, + error_message=error_details.get("message", str(e)), + gpu_info=gpu_info, + training_duration_seconds=time.time() - start_time, + ) + if self._dist_ctx.is_coordinator: + self._progress.report_error(error_details) + # Publish a failure marker so workers blocked on a coordinator + # barrier exit promptly instead of waiting out the full timeout. + self._dist_ctx.signal_failure() + finally: + # === Phase 5: Write result (coordinator only) === + self._write_result(result) + + # Returning outside `finally` so an uncaught BaseException (e.g. + # KeyboardInterrupt) propagates instead of being swallowed by the return. + return result + + # --- Helper methods --- + def _load_backend(self, backend_type: TrainingBackendEnum) -> TrainingBackend: + """Load the backend for the given backend type.""" + if backend_type == TrainingBackendEnum.NEMO_RL: + from .backends.nemo_rl.backend import NemoRLBackend + + return NemoRLBackend(self._job_ctx) + + raise ValueError(f"Unknown backend type: {backend_type}") + + def _get_barrier_dir(self) -> Path: + """Get the barrier directory for distributed coordination.""" + return self._job_ctx.storage_path / self._job_ctx.attempt_id / "distributed" / "barriers" + + def _load_config(self, config_path: Path) -> TrainingStepConfig: + """Load the training step config.""" + with open(config_path) as f: + config = TrainingStepConfig.model_validate(json.load(f)) + return config + + def _get_library_config_path(self) -> Path: + """ + Get the path for the library-specific config file. + + We define it here and pass it to the backend so that the backend can read it as-is without constructing paths. + """ + return self._workspace_path / f"{self._backend.backend_type.value}_config.yaml" + + def _preprocessing_phase(self) -> None: + """ + Run pre-training conversions if the backend supports them. + + Coordinator runs conversions (e.g., model format conversion), workers skip. + This phase runs before config compilation so that compiled configs can + reference converted artifacts. + + Only backends implementing SupportsPreprocessing will have conversions run. + + !!! Important !!! + Only coordinator runs _preprocessing_phase. Workers wait for coordinator to finish. + Any changes to the configs would affect only coordinator, so avoid any config changes. + """ + if self._dist_ctx.is_coordinator: + if isinstance(self._backend, SupportsPreprocessing): + self._progress.report_running("conversions") + self._backend.run_preprocessing(self._config) + logger.info("Pre-training conversions complete") + # Always release workers, even if no conversions are needed + self._dist_ctx.signal(BARRIER_PREPROCESSING_COMPLETE) + else: + self._dist_ctx.wait_for_coordinator(BARRIER_PREPROCESSING_COMPLETE) + + def _compile_config_phase(self) -> LibraryConfig: + """ + Compile library-specific config. + + Coordinator compiles config and writes to disk, then signals. + Workers wait for signal, then load the config file. + + The runner handles all file I/O; backend just compiles. + """ + config_path = self._get_library_config_path() + + if self._dist_ctx.is_coordinator: + self._progress.report_running("compiling_config") + + # Backend compiles config (pure transformation, no I/O) + config_dict = self._backend.compile_config(self._config, self._workspace_path) + + # Runner writes config to disk + config_path.parent.mkdir(parents=True, exist_ok=True) + with open(config_path, "w") as f: + yaml.dump(config_dict, f, default_flow_style=False) + + logger.info(f"Library config written to: {config_path}") + self._dist_ctx.signal(BARRIER_CONFIG_READY) + + return LibraryConfig(config_dict=config_dict, config_path=config_path) + else: + self._dist_ctx.wait_for_coordinator(BARRIER_CONFIG_READY) + return self._load_library_config(config_path) + + def _load_library_config(self, config_path: Path) -> LibraryConfig: + """Load library config from disk (used by workers).""" + if not config_path.exists(): + raise FileNotFoundError( + f"Library config not found at {config_path}. Coordinator may not have written it yet." + ) + + with open(config_path) as f: + config_dict = yaml.safe_load(f) + + logger.info(f"Loaded library config from: {config_path}") + return LibraryConfig(config_dict=config_dict, config_path=config_path) + + def _training_phase(self, library_config: LibraryConfig) -> TrainingMetrics: + """ + Execute training on all ranks. + + Training itself is distributed via Ray, which handles inter-process coordination internally. + """ + return self._backend.execute_training( + self._config, + library_config, + self._progress, + ) + + def _postprocess_phase( + self, + gpu_info: GPUInfo | None, + metrics: TrainingMetrics, + start_time: float, + library_config: LibraryConfig, + ) -> TrainingResult: + """ + Process checkpoint and create result. + + Workers return immediately with a minimal success result. They have no + post-training responsibilities, so letting them exit avoids barrier + timeouts that would cause Volcano to kill the coordinator mid-copy + because checkpoint copies for large models can take more than 600s + which is the default barrier timeout + + The coordinator finds the best checkpoint, copies/processes it to the + output path, and reports completion. + """ + if not self._dist_ctx.is_coordinator: + return TrainingResult( + success=True, + gpu_info=gpu_info, + training_duration_seconds=time.time() - start_time, + ) + + self._progress.report_running("processing_checkpoint") + checkpoint_path = self._backend.find_best_checkpoint(self._workspace_path, self._config, library_config) + checkpoint_info = self._backend.process_checkpoint( + checkpoint_path, self._output_path, self._config, library_config + ) + + result = TrainingResult( + success=True, + checkpoint=checkpoint_info, + gpu_info=gpu_info, + metrics=metrics, + training_duration_seconds=time.time() - start_time, + ) + + self._progress.report_completed("Training completed") + return result + + def _write_result(self, result: TrainingResult) -> None: + """Write result for downstream tasks.""" + if not self._dist_ctx.is_coordinator: + return + + result_path = self._workspace_path / DEFAULT_TRAINING_RESULT_FILE_NAME + result_path.parent.mkdir(parents=True, exist_ok=True) + with open(result_path, "w") as f: + f.write(result.model_dump_json(indent=2)) + logger.info(f"Result written to: {result_path}") diff --git a/services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja b/services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja new file mode 100644 index 0000000000..e074cba578 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/llama-3.1-instruct.jinja @@ -0,0 +1,61 @@ +{{- bos_token }} +{%- if not date_string is defined %} + {%- if strftime_now is defined %} + {%- set date_string = strftime_now("%d %b %Y") %} + {%- else %} + {%- set date_string = "26 Jul 2024" %} + {%- endif %} +{%- endif %} +{%- set loop_messages = messages %} +{%- if tools is not none and tool_choice is not none %} + {{- '<|start_header_id|>system<|end_header_id|>\n\n' }} + {{- "Environment: ipython\n\n" }} + {{- "Cutting Knowledge Date: December 2023\n" }} + {{- "Today Date: " + date_string + "\n\n" }} + {{- "You are a helpful assistant.\n" }} + {{- '<|eot_id|>' }} + {{- '<|start_header_id|>user<|end_header_id|>\n\n' }} + {{- 'You have access to the following functions to supplement your existing knowledge:\n\n' }} + {%- 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' }} + {%- endfor %} + {{- 'Think very carefully before calling functions.\n' }} + {{- 'Only call them if they are relevant to the prompt.\n' }} + {{- 'If you choose to call a function ONLY reply in the following format with no natural language surrounding it:\n\n' }} + {{- '{"example_name": "example_value"}\n\n' }} + {{- 'Reminder:\n' }} + {{- '- Function calls MUST follow the specified format, start with \n' }} + {{- '- Required parameters MUST be specified\n' }} + {{- '- Only call one function at a time\n' }} + {{- '- Put the entire function call reply on one line\n' }} + {{- '- Do not call functions if they are not relevant to the prompt' }} + {{- '<|eot_id|>' }} +{%- endif %} +{%- for message in loop_messages %} + {%- if message['role'] in ['ipython', 'tool'] %} + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }} + {{- "[stdout]" + message['content'] | trim + "[/stdout]\n<|eot_id|>" }} + {%- elif message['role'] == 'assistant'%} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} + {%- if message.get('tool_calls') is not none %} + {%- set tool_call = message['tool_calls'][0] %} + {%- generation %} + {{- '<|python_tag|>' + tool_call.function.arguments | tojson + '\n<|eot_id|>' }} + {%- endgeneration %} + {%- else %} + {%- generation %} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endgeneration %} + {%- endif %} + {%- else %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja b/services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja new file mode 100644 index 0000000000..e074cba578 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/llama-3.2-instruct.jinja @@ -0,0 +1,61 @@ +{{- bos_token }} +{%- if not date_string is defined %} + {%- if strftime_now is defined %} + {%- set date_string = strftime_now("%d %b %Y") %} + {%- else %} + {%- set date_string = "26 Jul 2024" %} + {%- endif %} +{%- endif %} +{%- set loop_messages = messages %} +{%- if tools is not none and tool_choice is not none %} + {{- '<|start_header_id|>system<|end_header_id|>\n\n' }} + {{- "Environment: ipython\n\n" }} + {{- "Cutting Knowledge Date: December 2023\n" }} + {{- "Today Date: " + date_string + "\n\n" }} + {{- "You are a helpful assistant.\n" }} + {{- '<|eot_id|>' }} + {{- '<|start_header_id|>user<|end_header_id|>\n\n' }} + {{- 'You have access to the following functions to supplement your existing knowledge:\n\n' }} + {%- 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' }} + {%- endfor %} + {{- 'Think very carefully before calling functions.\n' }} + {{- 'Only call them if they are relevant to the prompt.\n' }} + {{- 'If you choose to call a function ONLY reply in the following format with no natural language surrounding it:\n\n' }} + {{- '{"example_name": "example_value"}\n\n' }} + {{- 'Reminder:\n' }} + {{- '- Function calls MUST follow the specified format, start with \n' }} + {{- '- Required parameters MUST be specified\n' }} + {{- '- Only call one function at a time\n' }} + {{- '- Put the entire function call reply on one line\n' }} + {{- '- Do not call functions if they are not relevant to the prompt' }} + {{- '<|eot_id|>' }} +{%- endif %} +{%- for message in loop_messages %} + {%- if message['role'] in ['ipython', 'tool'] %} + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }} + {{- "[stdout]" + message['content'] | trim + "[/stdout]\n<|eot_id|>" }} + {%- elif message['role'] == 'assistant'%} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} + {%- if message.get('tool_calls') is not none %} + {%- set tool_call = message['tool_calls'][0] %} + {%- generation %} + {{- '<|python_tag|>' + tool_call.function.arguments | tojson + '\n<|eot_id|>' }} + {%- endgeneration %} + {%- else %} + {%- generation %} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endgeneration %} + {%- endif %} + {%- else %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja b/services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja new file mode 100644 index 0000000000..a0ba6017e1 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/llama-3.3-instruct.jinja @@ -0,0 +1,61 @@ +{{- bos_token }} +{%- if not date_string is defined %} + {%- if strftime_now is defined %} + {%- set date_string = strftime_now("%d %b %Y") %} + {%- else %} + {%- set date_string = "26 Jul 2024" %} + {%- endif %} +{%- endif %} +{%- set loop_messages = messages %} +{%- if tools is not none and tool_choice is not none %} + {{- '<|start_header_id|>system<|end_header_id|>\n\n' }} + {{- "Environment: ipython\n\n" }} + {{- "Cutting Knowledge Date: December 2023\n" }} + {{- "Today Date: " + date_string + "\n\n" }} + {{- "You are a helpful assistant.\n" }} + {{- '<|eot_id|>' }} + {{- '<|start_header_id|>user<|end_header_id|>\n\n' }} + {{- 'You have access to the following functions to supplement your existing knowledge:\n\n' }} + {%- 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' }} + {%- endfor %} + {{- 'Think very carefully before calling functions.\n' }} + {{- 'Only call them if they are relevant to the prompt.\n' }} + {{- 'If you choose to call a function ONLY reply in the following format with no natural language surrounding it:\n\n' }} + {{- '{"example_name": "example_value"}\n\n' }} + {{- 'Reminder:\n' }} + {{- '- Function calls MUST follow the specified format, start with \n' }} + {{- '- Required parameters MUST be specified\n' }} + {{- '- Only call one function at a time\n' }} + {{- '- Put the entire function call reply on one line\n' }} + {{- '- Do not call functions if they are not relevant to the prompt' }} + {{- '<|eot_id|>' }} +{%- endif %} +{%- for message in loop_messages %} + {%- if message['role'] in ['ipython', 'tool'] %} + {{- "<|start_header_id|>ipython<|end_header_id|>\n\n" }} + {{- "[stdout]" + message['content'] | trim + "[/stdout]\n<|eot_id|>" }} + {%- elif message['role'] == 'assistant'%} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} + {%- if message.get('tool_calls') is not none %} + {%- set tool_call = message['tool_calls'][0] %} + {%- generation %} + {{- '' + tool_call.function.arguments | tojson + '\n<|eot_id|>' }} + {%- endgeneration %} + {%- else %} + {%- generation %} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endgeneration %} + {%- endif %} + {%- else %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' }} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinja b/services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinja new file mode 100644 index 0000000000..00cfd85e48 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.1.jinja @@ -0,0 +1,51 @@ +{%- if messages[0]['role'] == 'system' %} + {%- set system_message = messages[0]['content'] | trim %} + {%- set messages = messages[1:] %} +{%- else %} + {%- set system_message = '' %} +{%- endif %} +{%- if tools is not none %} + {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>' + '\n\n' + system_message }} + {{- '\n\n' if system_message else '' }} + {{- '[' }} + {%- for t in tools %} + {{- (t.function if t.function is defined else t) | tojson() }} + {{- ', ' if not loop.last else '' }} + {%- endfor %} + {{- ']' }} + {{- '<|eot_id|>' }} +{%- else %} + {{- '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n' + system_message + '<|eot_id|>' }} +{%- endif %} +{%- for message in messages %} + {%- if (message['role'] in ['user', 'tool']) != (loop.index0 % 2 == 0) %} + {{- raise_exception('Conversation roles must alternate between user/tool and assistant') }} + {%- elif message['role'] == 'user' %} + {{- '<|start_header_id|>user<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }} + {%- elif message['role'] == 'tool' %} + {%- set tool_response = '[' + message['content'] | trim + ']' %} + {{- '<|start_header_id|>user<|end_header_id|>\n\n' + tool_response + '<|eot_id|>' }} + {%- elif message['role'] == 'assistant' and message.get('tool_calls') is not none %} + {%- set tool_calls = message['tool_calls'] %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} + {%- generation %} + {{- '['}} + {%- for tool_call in tool_calls %} + {{- '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}' }} + {%- if not loop.last %} + {{- ', ' }} + {%- else %} + {{- ']<|eot_id|>' }} + {%- endif %} + {%- endfor %} + {%- endgeneration %} + {%- elif message['role'] == 'assistant' %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} + {%- generation %} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endgeneration %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja b/services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja new file mode 100644 index 0000000000..7530a8c87e --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/nemotron-3.3.jinja @@ -0,0 +1,21 @@ +{{- bos_token }} +{%- if messages[0]['role'] == 'system' %} + {%- set system_message = messages[0]['content']|trim %} + {%- set messages = messages[1:] %} +{%- else %} + {%- set system_message = '' %} +{%- endif %} +{{- '<|start_header_id|>system<|end_header_id|>\n\n' }} +{{- system_message }} +{{- '<|eot_id|>' }} +{%- for message in messages %} + {%- if message['role'] == 'assistant' and '' in message['content'] %} + {%- set content = message['content'].split('')[-1].lstrip() %} + {%- else %} + {%- set content = message['content'] %} + {%- endif %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' + content | trim + '<|eot_id|>' }} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinja b/services/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinja new file mode 100644 index 0000000000..1deec26343 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/nemotron-super-3.3.jinja @@ -0,0 +1,82 @@ +{{- bos_token }} +{%- set ns = namespace(p='', has_tools=False) %} +{%- if tools is not none and tool_choice is not none %} + {%- set ns.has_tools = True %} + {%- set ns.p = ns.p + 'You are an expert in composing functions. You are given a question and a set of possible functions. ' %} + {%- set ns.p = ns.p + 'Based on the question, you will need to make one or more function/tool calls to achieve the purpose. ' %} + {%- set ns.p = ns.p + 'If none of the function can be used, point it out. ' %} + {%- set ns.p = ns.p + 'If the given question lacks the parameters required by the function, also point it out. ' %} + {%- set ns.p = ns.p + 'You should only return the function call in tools call sections. ' %} + {%- set ns.p = ns.p + 'Here is a list of functions in JSON format that you can invoke.\n' %} + {%- set ns.p = ns.p + '[' %} + {%- for tool in tools %} + {%- set function = tool.function %} + {%- set keys = function.keys() | reject('equalto', 'return') | list %} + {%- set ns.p = ns.p + '{"type": "function", "function": {' %} + {%- for key in keys %} + {%- set val = function[key] %} + {%- if val is string %} + {%- set ns.p = ns.p + '"' + key + '": "' + val + '"' %} + {%- else %} + {%- set ns.p = ns.p + '"' + key + '": ' + val|tojson %} + {%- endif %} + {%- if not loop.last %} + {%- set ns.p = ns.p + ', ' %} + {%- endif %} + {%- endfor %} + {%- set ns.p = ns.p + '}}' %} + {%- if not loop.last %} + {%- set ns.p = ns.p + ', ' %} + {%- endif %} + {%- endfor %} + {%- set ns.p = ns.p + ']\n' %} + {%- set ns.p = ns.p + 'If you decide to invoke any of the function(s), put it in the JSON TOOL CALLING format of ' %} + {%- set ns.p = ns.p + '[{"name": "func_name1", "arguments": {"params_name1": "params_value1", "params_name2": "params_value2"}}, ' %} + {%- set ns.p = ns.p + '{"name": "func_name2", "arguments": {"params_name1": "params_value1", "params_name2": "params_value2"}}] ' %} + {%- set ns.p = ns.p + '\n' %} + {%- set ns.p = ns.p + 'You SHOULD NOT include any other information in the response. REMEMBER TO USE JSON TOOL CALLING FORMAT.\n\n' %} +{%- endif %} +{%- for message in messages %} + {%- if message['role'] == 'user' %} + {%- if ns.has_tools %} + {%- if add_generation_prompt and loop.index0 == ((messages | length) - 1) %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + ns.p + (message['content'] | trim) + '<|eot_id|>' }} + {%- elif not add_generation_prompt and loop.index0 == ((messages | length) - 2) %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + ns.p + (message['content'] | trim) + '<|eot_id|>' }} + {%- else %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + (message['content'] | trim) + '<|eot_id|>' }} + {%- endif %} + {%- else %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + (message['content'] | trim) + '<|eot_id|>' }} + {%- endif %} + {%- elif message['role'] in ['ipython', 'tool'] %} + {{- '<|start_header_id|>user<|end_header_id|>\n\n' }} + {{- 'Here are the results from the tool:' + (message['content'] | trim) + '<|eot_id|>' }} + {%- elif message['role'] == 'assistant' %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} + {%- generation %} + {%- if message.get('tool_calls') is not none %} + {{- '[' }} + {%- for tool_call in message['tool_calls'] %} + {{- '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson }} + {%- if tool_call.get('id') is not none %} + {{- ', "id": "' + tool_call.id + '"' }} + {%- endif %} + {{- '}' }} + {%- if not loop.last %} + {{- ', ' }} + {%- endif %} + {%- endfor %} + {{- ']' }} + {{- '<|eot_id|>' }} + {%- else %} + {{- message['content'] | trim + '<|eot_id|>' }} + {%- endif %} + {%- endgeneration %} + {%- else %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + (message['content'] | trim) + '<|eot_id|>' }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/templates/phi-4.jinja b/services/rl/src/nmp/rl/tasks/training/templates/phi-4.jinja new file mode 100644 index 0000000000..33a466f884 --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/templates/phi-4.jinja @@ -0,0 +1,15 @@ +{%- for message in messages %} + {%- if (message['role'] == 'system') %} + {{- '<|im_start|>system<|im_sep|>' + message['content'] + '<|im_end|>'}} + {%- elif (message['role'] == 'user') %} + {{-'<|im_start|>user<|im_sep|>' + message['content'] + '<|im_end|>'}} + {%- elif (message['role'] == 'assistant') %} + {{- '<|im_start|>assistant<|im_sep|>' }} + {%- generation %} + {{- message['content'] + '<|im_end|>'}} + {%- endgeneration %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant<|im_sep|>' }} +{%- endif %} diff --git a/services/rl/src/nmp/rl/tasks/training/utils.py b/services/rl/src/nmp/rl/tasks/training/utils.py new file mode 100644 index 0000000000..650c0ff33c --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/utils.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +import os + +from nmp.common.jobs.constants import NEMO_JOB_ID_ENVVAR +from nmp.rl.app.jobs.training.schemas import GPUInfo + +logger = logging.getLogger(__name__) + + +def _get_architecture_name(major: int, minor: int) -> str: + """Map CUDA compute capability to architecture name. + + https://developer.nvidia.com/cuda-gpus + """ + if major == 3: + return "Kepler" + if major == 5: + return "Maxwell" + if major == 6: + return "Pascal" + if major == 7: + # 7.0/7.2 = Volta, 7.5 = Turing + if minor >= 5: + return "Turing" + return "Volta" + if major == 8: + return "Ampere" + if major == 9: + return "Hopper" + if major == 10: + return "Blackwell" + return f"Unknown (sm_{major}{minor})" + + +def get_gpu_info() -> GPUInfo | None: + """Capture GPU architecture information.""" + try: + import torch + + if not torch.cuda.is_available(): + return None + + device_id = torch.cuda.current_device() + props = torch.cuda.get_device_properties(device_id) + major, minor = torch.cuda.get_device_capability(device_id) + + return GPUInfo( + architecture=_get_architecture_name(major, minor), + device_name=props.name, + memory_gb=props.total_memory / (1024**3), + cuda_version=str(torch.version.cuda), + ) + except Exception as e: + logger.warning(f"Failed to capture GPU info: {e}") + return None + + +def generate_torchrun_flags_from_env() -> list[str]: + """Generate torchrun flags for distributed training.""" + # These values are typically injected by the Volcano/PyTorch operator + # or the Core Jobs Service when using DistributedGPUExecutionProvider. + master_addr = os.environ.get("MASTER_ADDR", "localhost") + master_port = os.environ.get("MASTER_PORT", "23456") # Default to port from volcano_job.py + node_rank = os.environ.get("NODE_RANK", os.environ.get("RANK", "0")) + num_nodes = os.environ.get("WORLD_SIZE", "1") + gpus_per_node = os.environ.get("GPUS_PER_NODE") + if gpus_per_node is None: + try: + import torch + + # device_count() returns 0 when CUDA is unavailable; torchrun's + # --nproc_per_node must be >= 1, so treat <= 0 as "unknown" and fall + # back to 1 rather than emitting "0". + device_count = torch.cuda.device_count() + gpus_per_node = str(device_count) if device_count > 0 else "1" + except Exception as e: + logger.warning(f"Failed to determine number of GPUs: {e}, using default of 1") + gpus_per_node = "1" + + return [ + "--nnodes", + num_nodes, + "--nproc_per_node", + gpus_per_node, + "--node_rank", + node_rank, + "--rdzv_id", + os.environ.get(NEMO_JOB_ID_ENVVAR, "customizer-rdzv"), + "--rdzv_backend", + "c10d", + "--rdzv_endpoint", + f"{master_addr}:{master_port}", + ] diff --git a/services/rl/tests/test_compiler.py b/services/rl/tests/test_compiler.py new file mode 100644 index 0000000000..ad5fa4c9c9 --- /dev/null +++ b/services/rl/tests/test_compiler.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compiler tests: public-spec → TrainingStepConfig mapping, executor selection, +and the 4-step PlatformJobSpec shape.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any +from unittest.mock import AsyncMock, Mock + +import pytest +from nemo_platform import AsyncNeMoPlatform +from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.integrations import IntegrationsSpec, MlflowIntegration, WandbIntegration +from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError +from nmp.common.entities.utils import get_random_id +from nmp.rl.app.jobs.compiler import ( + _build_training_step, + _build_training_step_config, + platform_job_config_compiler, +) +from nmp.rl.app.jobs.training.schemas import OptimizerType, TrainingType +from nmp.rl.entities.values import FinetuningType +from nmp.rl.schemas import DPOTraining, OutputResponse, ParallelismParams, RlJobOutput + + +def _make_model_entity(fileset: str | None = "default/base-model") -> ModelEntity: + return ModelEntity( + id=get_random_id("model"), + workspace="default", + name="base-model", + fileset=fileset, + trust_remote_code=False, + finetuning_type=None, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + +def _make_job_output( + training: DPOTraining | None = None, + integrations: IntegrationsSpec | None = None, +) -> RlJobOutput: + return RlJobOutput( + model="default/base-model", + dataset="default/prefs", + training=training or DPOTraining(), + integrations=integrations, + output=OutputResponse(name="my-dpo", type="model", fileset="my-dpo-fs"), + ) + + +# Job specs/steps/executors/containers are all TypedDicts (plain dicts). +def _container(step: dict[str, Any]) -> dict[str, Any]: + return step["executor"]["container"] + + +def _provider(step: dict[str, Any]) -> str: + return step["executor"]["provider"] + + +@pytest.fixture +def mock_sdk() -> Mock: + return Mock(spec=AsyncNeMoPlatform) + + +# --------------------------------------------------------------------------- # +# _build_training_step_config: public DPOTraining → internal TrainingStepConfig +# --------------------------------------------------------------------------- # + + +def test_training_step_config_maps_exposed_knobs() -> None: + t = DPOTraining( + optimizer_type=OptimizerType.ADAM_WITH_FLAT_LR, + adam_eps=3e-7, + activation_checkpointing=True, + keep_top_k=5, + val_at_end=True, + ref_policy_kl_penalty=0.2, + max_grad_norm=2.0, + ) + sc = _build_training_step_config(_make_job_output(t), trust_remote_code=True) + + # Optimizer knobs. + assert sc.optimizer.optimizer_type is OptimizerType.ADAM_WITH_FLAT_LR + assert sc.optimizer.eps == 3e-7 + # Memory / checkpoint / validation knobs. + assert sc.parallelism.activation_checkpointing is True + assert sc.schedule.keep_top_k == 5 + assert sc.schedule.val_at_end is True + # DPO hyperparameters + passthrough. + assert sc.training.training_type is TrainingType.DPO + assert sc.training.finetuning_type is FinetuningType.ALL_WEIGHTS + assert sc.training.dpo is not None + assert sc.training.dpo.ref_policy_kl_penalty == 0.2 + assert sc.training.dpo.max_grad_norm == 2.0 + assert sc.model.trust_remote_code is True + + +def test_training_step_config_maps_integrations() -> None: + """job_spec.integrations must reach the step config; otherwise W&B/MLflow are + silently disabled because the driver's builders read customizer_config.integrations.""" + integrations = IntegrationsSpec( + wandb=WandbIntegration( + project="proj", name="run", entity="team", tags=["t1"], notes="n", base_url="https://wandb.example" + ), + mlflow=MlflowIntegration( + experiment_name="exp", name="mlrun", tags={"k": "v"}, description="d", tracking_uri="http://mlflow:5000" + ), + ) + sc = _build_training_step_config(_make_job_output(integrations=integrations), trust_remote_code=False) + + assert sc.integrations.wandb is not None + assert sc.integrations.wandb.project == "proj" + assert sc.integrations.wandb.name == "run" + assert sc.integrations.wandb.entity == "team" + assert sc.integrations.wandb.base_url == "https://wandb.example" + + assert sc.integrations.mlflow is not None + assert sc.integrations.mlflow.experiment_name == "exp" + # public MLflow `name` maps to the step config's `run_name` + assert sc.integrations.mlflow.run_name == "mlrun" + assert sc.integrations.mlflow.tracking_uri == "http://mlflow:5000" + assert sc.integrations.mlflow.tags == {"k": "v"} + + +def test_training_step_config_no_integrations_is_empty() -> None: + sc = _build_training_step_config(_make_job_output(), trust_remote_code=False) + assert sc.integrations.wandb is None + assert sc.integrations.mlflow is None + + +def test_training_step_config_defaults_match_prior_hardcodes() -> None: + sc = _build_training_step_config(_make_job_output(), trust_remote_code=False) + assert sc.optimizer.optimizer_type is None + assert sc.optimizer.eps == 1e-5 + assert sc.parallelism.activation_checkpointing is False + assert sc.schedule.keep_top_k == 1 + # val_at_end defaults True → final checkpoint carries val metrics for best-checkpoint selection. + assert sc.schedule.val_at_end is True + + +# --------------------------------------------------------------------------- # +# _build_training_step: executor selection by topology +# --------------------------------------------------------------------------- # + + +def test_single_node_uses_gpu_executor() -> None: + job = _make_job_output(DPOTraining(parallelism=ParallelismParams(num_nodes=1, num_gpus_per_node=1))) + step = _build_training_step(job, [], trust_remote_code=False, profile=None) + assert step["name"] == "dpo-training" + assert _provider(step) == "gpu" + assert _container(step)["command"] == ["-m", "nmp.rl.tasks.training"] + + +def test_multi_node_requires_shared_storage(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("nmp.rl.app.jobs.compiler.config.multinode_shared_storage_path", None, raising=False) + job = _make_job_output(DPOTraining(parallelism=ParallelismParams(num_nodes=2, num_gpus_per_node=2))) + with pytest.raises(PlatformJobCompilationError, match="shared filesystem"): + _build_training_step(job, [], trust_remote_code=False, profile=None) + + +def test_multi_node_uses_distributed_executor_with_shared_storage(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("nmp.rl.app.jobs.compiler.config.multinode_shared_storage_path", "/shared", raising=False) + job = _make_job_output(DPOTraining(parallelism=ParallelismParams(num_nodes=2, num_gpus_per_node=2))) + step = _build_training_step(job, [], trust_remote_code=False, profile=None) + assert _provider(step) == "gpu_distributed" + + # BASE_LOG_DIR is injected so Ray can coordinate the cross-node barrier. + def _env_value(env: Any) -> Any: + return env["value"] if isinstance(env, dict) else getattr(env, "value", None) + + assert any(_env_value(env) == "/shared" for env in step["environment"]) + + +def test_explicit_profile_overrides_default() -> None: + job = _make_job_output() + step = _build_training_step(job, [], trust_remote_code=False, profile="custom-gpu") + assert step["executor"]["profile"] == "custom-gpu" + + +# --------------------------------------------------------------------------- # +# platform_job_config_compiler: full 4-step spec +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_compiler_emits_four_steps(monkeypatch: pytest.MonkeyPatch, mock_sdk: Mock) -> None: + monkeypatch.setattr( + "nmp.rl.app.jobs.compiler.fetch_model_entity", + AsyncMock(return_value=_make_model_entity()), + ) + spec = await platform_job_config_compiler("default", _make_job_output(), mock_sdk) + + steps = spec["steps"] + names = [s["name"] for s in steps] + assert names == ["model-and-dataset-download", "dpo-training", "model-upload", "model-entity-creation"] + + # CPU task steps share the lighter tasks image; the GPU step uses the training image. + assert "nmp-rl-tasks" in _container(steps[0])["image"] + assert "nmp-rl-training" in _container(steps[1])["image"] + assert "nmp-rl-tasks" in _container(steps[2])["image"] + assert _container(steps[0])["command"] == ["-m", "nmp.rl.tasks.file_io"] + assert _container(steps[3])["command"] == ["-m", "nmp.rl.tasks.model_entity"] + + +@pytest.mark.asyncio +async def test_compiler_rejects_model_without_fileset(monkeypatch: pytest.MonkeyPatch, mock_sdk: Mock) -> None: + monkeypatch.setattr( + "nmp.rl.app.jobs.compiler.fetch_model_entity", + AsyncMock(return_value=_make_model_entity(fileset=None)), + ) + with pytest.raises(PlatformJobCompilationError, match="has no fileset"): + await platform_job_config_compiler("default", _make_job_output(), mock_sdk) diff --git a/services/rl/tests/test_dpo_config.py b/services/rl/tests/test_dpo_config.py new file mode 100644 index 0000000000..a8f2ae5d2e --- /dev/null +++ b/services/rl/tests/test_dpo_config.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the pure config-builder helpers in dpo_config. + +These cover the optimizer/scheduler/precision/data/logger builders and the inert +Megatron block — i.e. everything except ``compile_dpo_config`` itself, which needs +a real on-disk dataset to prepare and validate. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from nmp.customization_common.service.context import NMPJobContext +from nmp.rl.app.jobs.training.schemas import ( + DPOConfig, + ModelConfig, + OptimizerType, + TrainingStepConfig, + TrainingType, +) +from nmp.rl.tasks.training.backends.nemo_rl import dpo_config +from nmp.rl.tasks.training.backends.nemo_rl.dpo_config import ( + _adapt_precision, + _build_data_config, + _build_logger_config, + _build_optimizer_config, + _build_scheduler_config, + _megatron_cfg_disabled, +) +from nmp.rl.tasks.training.datasets.preparation import PreparedDataset + + +def _make_step_config( + *, + optimizer: TrainingStepConfig.OptimizerConfig | None = None, + schedule: TrainingStepConfig.ScheduleConfig | None = None, + parallelism: TrainingStepConfig.ParallelismConfig | None = None, + max_seq_length: int = 1024, +) -> TrainingStepConfig: + return TrainingStepConfig( + model=ModelConfig(path="/model", max_seq_length=max_seq_length), + dataset=TrainingStepConfig.DatasetConfig(path="/data"), + training=TrainingStepConfig.TrainingConfig(training_type=TrainingType.DPO, dpo=DPOConfig()), + schedule=schedule or TrainingStepConfig.ScheduleConfig(), + batch=TrainingStepConfig.BatchConfig(), + optimizer=optimizer or TrainingStepConfig.OptimizerConfig(), + parallelism=parallelism or TrainingStepConfig.ParallelismConfig(), + output_model="out", + ) + + +def _job_ctx(tmp_path: Path) -> NMPJobContext: + return NMPJobContext( + workspace="default", + job_id="rl-test", + attempt_id="attempt-1", + step="dpo-training", + task="task-1", + jobs_url=None, + files_url=None, + storage_path=tmp_path, + config_path=tmp_path / "config.json", + ) + + +# --------------------------------------------------------------------------- # +# _adapt_precision +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "value,expected", + [ + ("bf16", "bfloat16"), + ("bf16-mixed", "bfloat16"), + ("fp16", "float16"), + ("fp32", "float32"), + (None, "bfloat16"), + ("nonsense", "bfloat16"), # unknown → safe default + ], +) +def test_adapt_precision(value: str | None, expected: str) -> None: + assert _adapt_precision(value) == expected + + +# --------------------------------------------------------------------------- # +# _build_optimizer_config +# --------------------------------------------------------------------------- # + + +def test_optimizer_config_adamw_default() -> None: + opt = _build_optimizer_config(_make_step_config()) # optimizer_type None → AdamW + assert opt["name"] == "torch.optim.AdamW" + + +@pytest.mark.parametrize( + "opt_type", + [OptimizerType.ADAM_WITH_COSINE_ANNEALING, OptimizerType.ADAM_WITH_FLAT_LR], +) +def test_optimizer_config_adam_variants(opt_type: OptimizerType) -> None: + cfg = _make_step_config(optimizer=TrainingStepConfig.OptimizerConfig(optimizer_type=opt_type)) + assert _build_optimizer_config(cfg)["name"] == "torch.optim.Adam" + + +def test_optimizer_config_passes_through_kwargs() -> None: + optimizer = TrainingStepConfig.OptimizerConfig( + learning_rate=2e-5, weight_decay=0.05, beta1=0.8, beta2=0.95, eps=3e-7 + ) + kwargs = _build_optimizer_config(_make_step_config(optimizer=optimizer))["kwargs"] + assert kwargs["lr"] == 2e-5 + assert kwargs["weight_decay"] == 0.05 + assert kwargs["betas"] == [0.8, 0.95] + assert kwargs["eps"] == 3e-7 # the configurable knob actually flows through + + +# --------------------------------------------------------------------------- # +# _build_scheduler_config +# --------------------------------------------------------------------------- # + + +def test_scheduler_cosine_is_warmup_then_decay_chain() -> None: + cfg = _make_step_config( + optimizer=TrainingStepConfig.OptimizerConfig( + optimizer_type=OptimizerType.ADAMW_WITH_COSINE_ANNEALING, warmup_steps=10 + ) + ) + sched = _build_scheduler_config(cfg, num_steps=100) + assert isinstance(sched, list) + assert sched[0]["name"] == "torch.optim.lr_scheduler.LinearLR" + assert sched[1]["name"] == "torch.optim.lr_scheduler.CosineAnnealingLR" + assert sched[2]["milestones"] == [10] + + +def test_scheduler_flat_is_constant_lr() -> None: + cfg = _make_step_config( + optimizer=TrainingStepConfig.OptimizerConfig(optimizer_type=OptimizerType.ADAMW_WITH_FLAT_LR) + ) + sched = _build_scheduler_config(cfg, num_steps=100) + assert isinstance(sched, dict) + assert sched["name"] == "torch.optim.lr_scheduler.ConstantLR" + + +# --------------------------------------------------------------------------- # +# _megatron_cfg_disabled (inert block, must still be fully populated) +# --------------------------------------------------------------------------- # + + +def test_megatron_cfg_is_inert_but_complete() -> None: + mc = _megatron_cfg_disabled(precision="bfloat16", max_grad_norm=2.5) + assert mc["enabled"] is False + assert mc["pipeline_dtype"] == "bfloat16" # tracks policy.precision + assert mc["optimizer"]["clip_grad"] == 2.5 # tracks policy.max_grad_norm + # All required sub-blocks present so NeMo-RL's schema validates. + for key in ("peft", "optimizer", "scheduler", "distributed_data_parallel_config", "fp8_cfg"): + assert key in mc + assert mc["fp8_cfg"]["enabled"] is False + assert mc["peft"]["enabled"] is False + + +# --------------------------------------------------------------------------- # +# _build_data_config +# --------------------------------------------------------------------------- # + + +def test_data_config_binary_preference(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(dpo_config, "detect_dpo_schema_name", lambda _path: "BinaryPreferenceDataset") + prepared = PreparedDataset( + merged_dir=tmp_path, + train_file=tmp_path / "training.jsonl", + validation_file=tmp_path / "validation.jsonl", + train_samples=10, + validation_samples=2, + ) + data = _build_data_config(_make_step_config(max_seq_length=512), prepared) + + assert data["max_input_seq_length"] == 512 + assert data["shuffle"] is False # deterministic ordering is an intentional override + for split in ("train", "validation"): + assert data[split]["dataset_name"] == "BinaryPreferenceDataset" + assert data[split]["prompt_key"] == "prompt" + assert data[split]["chosen_key"] == "chosen" + assert data[split]["rejected_key"] == "rejected" + assert data["train"]["data_path"] == str(prepared.train_file) + + +# --------------------------------------------------------------------------- # +# _build_logger_config +# --------------------------------------------------------------------------- # + + +def test_logger_config_has_all_subsections_when_integrations_disabled(tmp_path: Path) -> None: + cfg = _build_logger_config(_make_step_config(), _job_ctx(tmp_path), tmp_path) + + assert cfg["wandb_enabled"] is False + assert cfg["mlflow_enabled"] is False + assert cfg["monitor_gpus"] is False + assert cfg["log_dir"].endswith("logs") + # Every backend subsection is present even when disabled (NeMo-RL expects them). + for key in ("wandb", "swanlab", "tensorboard", "mlflow", "gpu_monitoring"): + assert key in cfg diff --git a/services/rl/tests/test_preparation.py b/services/rl/tests/test_preparation.py new file mode 100644 index 0000000000..2f4fc6a7a0 --- /dev/null +++ b/services/rl/tests/test_preparation.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for dataset preparation (train/validation auto-split).""" + +import json +from pathlib import Path + +import pytest +from nmp.rl.tasks.training.datasets.preparation import DatasetFormatError, _create_val_split + + +def _write_jsonl(path: Path, rows: list[dict]) -> None: + path.write_text("".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8") + + +def _read_jsonl(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + + +def test_create_val_split_in_place_does_not_truncate_source(tmp_path: Path) -> None: + """train_file == output_train (auto-split-from-merged) must not destroy the source. + + Regression: a streaming split that opens output_train for write before reading + truncates the file when both paths are the same, yielding empty splits. + """ + rows = [{"prompt": f"q{i}", "chosen": f"good{i}", "rejected": f"bad{i}"} for i in range(10)] + train = tmp_path / "train.jsonl" + val = tmp_path / "validation.jsonl" + _write_jsonl(train, rows) + + # Same path for source and train output — the auto-split-from-merged case. + train_n, val_n = _create_val_split(train, train, val, val_ratio=0.2, seed=1234) + + assert train_n + val_n == len(rows) + assert train_n > 0 and val_n > 0, "in-place split produced an empty side (source was truncated)" + assert len(_read_jsonl(train)) == train_n + assert len(_read_jsonl(val)) == val_n + # No leftover temp file from the atomic replace. + assert not (tmp_path / "train.jsonl.tmp").exists() + + +def test_create_val_split_distinct_paths(tmp_path: Path) -> None: + """Distinct source/output paths split correctly and leave the source intact.""" + rows = [{"prompt": f"q{i}", "chosen": f"good{i}", "rejected": f"bad{i}"} for i in range(10)] + src = tmp_path / "source.jsonl" + train_out = tmp_path / "merged" / "train.jsonl" + val_out = tmp_path / "merged" / "validation.jsonl" + _write_jsonl(src, rows) + + train_n, val_n = _create_val_split(src, train_out, val_out, val_ratio=0.1, seed=1234) + + assert train_n + val_n == len(rows) + assert train_n > 0 and val_n > 0 + assert len(_read_jsonl(src)) == len(rows), "source file must be left untouched" + + +def test_create_val_split_is_deterministic(tmp_path: Path) -> None: + """Same seed → same split, and the local RNG doesn't depend on global state.""" + rows = [{"prompt": f"q{i}", "chosen": f"good{i}", "rejected": f"bad{i}"} for i in range(20)] + src = tmp_path / "source.jsonl" + _write_jsonl(src, rows) + + def split_prompts(suffix: str) -> list[str]: + val = tmp_path / f"val{suffix}.jsonl" + _create_val_split(src, tmp_path / f"train{suffix}.jsonl", val, val_ratio=0.25, seed=99) + return sorted(r["prompt"] for r in _read_jsonl(val)) + + assert split_prompts("a") == split_prompts("b") + + +def test_create_val_split_rejects_too_small(tmp_path: Path) -> None: + """Fewer than 2 rows cannot be split into non-empty train + validation.""" + src = tmp_path / "train.jsonl" + _write_jsonl(src, [{"prompt": "q", "chosen": "a", "rejected": "b"}]) + + with pytest.raises(DatasetFormatError): + _create_val_split(src, src, tmp_path / "validation.jsonl") diff --git a/services/rl/tests/test_schemas.py b/services/rl/tests/test_schemas.py new file mode 100644 index 0000000000..8d9e9fe2fc --- /dev/null +++ b/services/rl/tests/test_schemas.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public NeMo-RL schema tests: field defaults and the cross-field validators.""" + +from __future__ import annotations + +import pytest +from nmp.customization_common.schemas.values import OutputNameType +from nmp.rl.app.jobs.training.schemas import OptimizerType +from nmp.rl.schemas import DPOTraining, OutputResponse, ParallelismParams, RlJobOutput + + +def _make_output(name: str = "out", out_type: OutputNameType = OutputNameType.MODEL) -> OutputResponse: + return OutputResponse(name=name, type=out_type, fileset=f"{name}-fs") + + +def _make_job_output(training: DPOTraining, out_type: OutputNameType = OutputNameType.MODEL) -> RlJobOutput: + return RlJobOutput( + model="default/base", + dataset="default/prefs", + training=training, + output=_make_output(out_type=out_type), + ) + + +def test_dpo_training_defaults_preserve_prior_behavior() -> None: + """The newly exposed knobs default to the values the compiler used to hardcode.""" + t = DPOTraining() + assert t.type == "dpo" + # Newly exposed configurability. + assert t.optimizer_type is None # → AdamW + cosine annealing + assert t.adam_eps == 1e-5 + assert t.activation_checkpointing is False + assert t.keep_top_k == 1 + # val_at_end defaults True so the final checkpoint carries validation metrics + # and best-checkpoint selection works (otherwise NeMo-RL falls back to latest). + assert t.val_at_end is True + # Existing DPO hyperparameters. + assert t.ref_policy_kl_penalty == 0.05 + assert t.sft_loss_weight == 0.0 + + +def test_dpo_training_accepts_overrides() -> None: + t = DPOTraining( + optimizer_type=OptimizerType.ADAM_WITH_FLAT_LR, + adam_eps=1e-8, + activation_checkpointing=True, + keep_top_k=3, + val_at_end=True, + ) + assert t.optimizer_type is OptimizerType.ADAM_WITH_FLAT_LR + assert t.adam_eps == 1e-8 + assert t.activation_checkpointing is True + assert t.keep_top_k == 3 + assert t.val_at_end is True + + +@pytest.mark.parametrize("bad", [0.0, -1e-5]) +def test_adam_eps_must_be_positive(bad: float) -> None: + with pytest.raises(ValueError): + DPOTraining(adam_eps=bad) + + +def test_keep_top_k_must_be_positive() -> None: + with pytest.raises(ValueError): + DPOTraining(keep_top_k=0) + + +def test_validate_for_training_accepts_consistent_single_gpu() -> None: + # 1 GPU, no model parallelism, gb divisible by micro*dp → no error. + job = _make_job_output(DPOTraining(batch_size=32, micro_batch_size=1)) + job.validate_for_training() + + +def test_validate_for_training_rejects_indivisible_model_parallel() -> None: + # total_gpus=1 but tensor_parallel_size=2 → 1 % 2 != 0. + job = _make_job_output( + DPOTraining(parallelism=ParallelismParams(num_gpus_per_node=1, tensor_parallel_size=2)), + ) + with pytest.raises(ValueError, match="must be divisible by tensor_parallel_size"): + job.validate_for_training() + + +def test_validate_for_training_rejects_indivisible_batch() -> None: + # total_gpus=2, mp=1 → data_parallel=2; batch_size=3 not divisible by micro(1)*dp(2). + job = _make_job_output( + DPOTraining( + parallelism=ParallelismParams(num_gpus_per_node=2), + batch_size=3, + micro_batch_size=1, + ), + ) + with pytest.raises(ValueError, match="batch_size"): + job.validate_for_training() + + +def test_dpo_output_must_be_full_weight_model() -> None: + # DPO is full-weight; an adapter output is rejected at construction time. + with pytest.raises(ValueError, match="full-weight model"): + _make_job_output(DPOTraining(), out_type=OutputNameType.ADAPTER) diff --git a/uv.lock b/uv.lock index 91c87960a8..4c5112e274 100644 --- a/uv.lock +++ b/uv.lock @@ -41,6 +41,7 @@ members = [ "nemo-platform-plugin", "nemo-platform-sdk", "nemo-platform-sdk-tools", + "nemo-rl-plugin", "nemo-safe-synthesizer-plugin", "nemo-switchyard", "nemo-unsloth-plugin", @@ -63,6 +64,7 @@ members = [ "nmp-platform", "nmp-platform-runner", "nmp-platform-seed", + "nmp-rl", "nmp-secrets", "nmp-studio", "nmp-testing", @@ -5943,6 +5945,51 @@ test = [ { name = "pytest-mock", specifier = ">=3.14.1" }, ] +[[package]] +name = "nemo-rl-plugin" +version = "0.1.0" +source = { editable = "plugins/nemo-rl" } +dependencies = [ + { name = "nemo-platform", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nmp-customization-common", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nmp-rl", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-settings", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "typer", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.dev-dependencies] +dev = [ + { name = "fastapi", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "httpx", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-customizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pytest", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pytest-asyncio", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "ruff", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "nemo-platform", editable = "packages/nemo_platform" }, + { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "nmp-customization-common", editable = "packages/nmp_customization_common" }, + { name = "nmp-rl", editable = "services/rl" }, + { name = "pydantic", specifier = ">=2.10.6" }, + { name = "pydantic-settings", specifier = ">=2.6.1" }, + { name = "typer", specifier = ">=0.12.5" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "fastapi", specifier = ">=0.115.0" }, + { name = "httpx", specifier = ">=0.27.0" }, + { name = "nemo-customizer-plugin", editable = "plugins/nemo-customizer" }, + { name = "pytest", specifier = ">=8.3.4" }, + { name = "pytest-asyncio", specifier = ">=0.25.3" }, + { name = "ruff", specifier = ">=0.11.8" }, +] + [[package]] name = "nemo-safe-synthesizer" version = "0.1.2" @@ -6193,6 +6240,7 @@ core-services = [ { name = "nemo-guardrails-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform", extra = ["services"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-rl-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-safe-synthesizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-switchyard", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-unsloth-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6286,6 +6334,7 @@ enabled-plugins = [ { name = "nemo-deployments-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-evaluator-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-guardrails-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-rl-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-safe-synthesizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-switchyard", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-unsloth-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6302,6 +6351,7 @@ functional-services = [ { name = "nemo-guardrails-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform", extra = ["services"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-rl-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-safe-synthesizer-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-switchyard", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "nemo-unsloth-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -6393,6 +6443,7 @@ core-services = [ { name = "nemo-platform", editable = "packages/nemo_platform" }, { name = "nemo-platform", extras = ["services"], editable = "packages/nemo_platform" }, { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "nemo-rl-plugin", editable = "plugins/nemo-rl" }, { name = "nemo-safe-synthesizer-plugin", editable = "plugins/nemo-safe-synthesizer" }, { name = "nemo-switchyard", editable = "plugins/nemo-switchyard" }, { name = "nemo-unsloth-plugin", editable = "plugins/nemo-unsloth" }, @@ -6488,6 +6539,7 @@ enabled-plugins = [ { name = "nemo-deployments-plugin", editable = "plugins/nemo-deployments" }, { name = "nemo-evaluator-plugin", editable = "plugins/nemo-evaluator" }, { name = "nemo-guardrails-plugin", editable = "plugins/nemo-guardrails" }, + { name = "nemo-rl-plugin", editable = "plugins/nemo-rl" }, { name = "nemo-safe-synthesizer-plugin", editable = "plugins/nemo-safe-synthesizer" }, { name = "nemo-switchyard", editable = "plugins/nemo-switchyard" }, { name = "nemo-unsloth-plugin", editable = "plugins/nemo-unsloth" }, @@ -6505,6 +6557,7 @@ functional-services = [ { name = "nemo-platform", editable = "packages/nemo_platform" }, { name = "nemo-platform", extras = ["services"], editable = "packages/nemo_platform" }, { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "nemo-rl-plugin", editable = "plugins/nemo-rl" }, { name = "nemo-safe-synthesizer-plugin", editable = "plugins/nemo-safe-synthesizer" }, { name = "nemo-switchyard", editable = "plugins/nemo-switchyard" }, { name = "nemo-unsloth-plugin", editable = "plugins/nemo-unsloth" }, @@ -7449,6 +7502,54 @@ dev = [ { name = "pytest-asyncio", specifier = ">=0.24.0" }, ] +[[package]] +name = "nmp-rl" +version = "0.1.0" +source = { editable = "services/rl" } +dependencies = [ + { name = "httpx", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-platform-plugin", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nemo-platform-sdk", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nmp-common", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "nmp-customization-common", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic", extra = ["email"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pydantic-settings", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "tenacity", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.optional-dependencies] +integrations = [ + { name = "mlflow-skinny", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "wandb", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "pytest-asyncio", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.27.0" }, + { name = "mlflow-skinny", marker = "extra == 'integrations'" }, + { name = "nemo-platform-plugin", editable = "packages/nemo_platform_plugin" }, + { name = "nemo-platform-sdk", editable = "sdk/python/nemo-platform" }, + { name = "nmp-common", editable = "packages/nmp_common" }, + { name = "nmp-customization-common", editable = "packages/nmp_customization_common" }, + { name = "pydantic", specifier = ">=2.10.6" }, + { name = "pydantic-settings", specifier = ">=2.6.1" }, + { name = "tenacity", specifier = ">=8.5.0" }, + { name = "wandb", marker = "extra == 'integrations'", specifier = ">=0.25.1" }, +] +provides-extras = ["integrations"] + +[package.metadata.requires-dev] +dev = [ + { name = "pytest", specifier = ">=8.3.4" }, + { name = "pytest-asyncio", specifier = ">=0.25.3" }, +] + [[package]] name = "nmp-secrets" version = "0.1.0"