From 3b097b608966a75e757372acece9c5a7660bee81 Mon Sep 17 00:00:00 2001 From: Pritam Gundecha Date: Sun, 17 May 2026 19:26:07 -0700 Subject: [PATCH] v1.1.1: air-gapped Slurm support, dual-mode runtime, GRPO multi-env hardening - Air-gapped Slurm support: pre-built nemo-skills/nemo-rl/vLLM/vllm-grpo containers with offline env vars (HF_*_OFFLINE, UV_OFFLINE, TIKTOKEN_*); reproducible Dockerfiles + build instructions - Dual-mode runtime: nvflow/lib/runtime.py centralises venv/python resolution so the same code path runs in airgap and dev modes; checkpoint converter, vLLM serving, training, and SDG scripts all use it - sbatch arg propagation: nvflow/lib/sbatch.py + tests plumb extra_sbatch_args through every Slurm submission - GRPO eval per-environment: outputs split by env (equivalence_llm_judge, finance_sec_search); restored sequence_packing / logprob_chunk_size / make_sequence_length_divisible_by for the equivalence_llm_judge training policy - Eval / SDG / SFT fixes: skip tiktoken download when cache env vars pre-configured, eval base/demo config refinements, removed unused nemotron_nano_9b SFT config - CI: lightweight unit-tests workflow on both GitLab and GitHub. tests/test_sbatch.py reads source files via Path(...).read_text() instead of importing modules that transitively need nemo_skills / typer; both runners invoke `uv run --no-sync pytest tests/ -v --tb=short` so the --no-deps install is authoritative - Documentation: INSTALL.md rewrite, dockerfiles/README + docker_instructions, finance troubleshooting guide, cluster-configuration docs Signed-off-by: Pritam Gundecha --- .github/workflows/unit-tests.yml | 5 +- .gitlab-ci.yml | 5 +- INSTALL.md | 314 ++++++++------ README.md | 15 +- cluster_configs/containers.yaml | 44 +- cluster_configs/template-slurm.yaml | 62 ++- dockerfiles/Dockerfile.nemo-rl | 131 ++++++ dockerfiles/Dockerfile.nemo-skills | 138 ++++++ dockerfiles/Dockerfile.vllm | 56 +++ dockerfiles/Dockerfile.vllm-grpo | 55 +++ dockerfiles/README.md | 52 +++ dockerfiles/docker_instructions.md | 392 ++++++++++++++++++ docs/architecture/ARCHITECTURE.md | 42 +- docs/cluster-configuration.md | 57 ++- docs/recipes/finance/troubleshooting.md | 83 ++++ nvflow/core/workflow_runner.py | 7 + nvflow/lib/rl/rollout.py | 12 +- nvflow/lib/rl/verify.py | 1 + nvflow/lib/runtime.py | 60 +++ nvflow/lib/sbatch.py | 116 ++++++ nvflow/recipes/finance/stages/rl/training.py | 10 +- .../finance/stages/sdg/create_seed_data.py | 4 - .../finance/stages/sdg/generate_questions.py | 5 +- .../stages/sdg/map_questions_to_context.py | 1 - nvflow/recipes/finance/stages/sft/training.py | 6 +- .../utils/evaluation/checkpoint_converter.py | 49 ++- .../workflows/download_sec_filings.yaml | 18 - .../recipes/finance/workflows/eval/base.yaml | 8 +- .../recipes/finance/workflows/eval/demo.yaml | 15 + .../recipes/finance/workflows/grpo/base.yaml | 37 +- .../finance/workflows/grpo/qwen3_4b.yaml | 4 + .../workflows/grpo/qwen3_4b_finsec.yaml | 1 + .../sdg/template-based-sdg-demo.yaml | 1 - .../recipes/finance/workflows/sft/base.yaml | 4 + .../workflows/sft/nemotron_nano_9b.yaml | 140 ------- scripts/convert_checkpoint_to_hf.sh | 15 +- scripts/serve_vllm_patched.py | 13 +- tests/test_sbatch.py | 208 ++++++++++ 38 files changed, 1776 insertions(+), 410 deletions(-) create mode 100644 dockerfiles/Dockerfile.nemo-rl create mode 100644 dockerfiles/Dockerfile.nemo-skills create mode 100644 dockerfiles/Dockerfile.vllm create mode 100644 dockerfiles/Dockerfile.vllm-grpo create mode 100644 dockerfiles/README.md create mode 100644 dockerfiles/docker_instructions.md create mode 100644 nvflow/lib/runtime.py create mode 100644 nvflow/lib/sbatch.py delete mode 100644 nvflow/recipes/finance/workflows/sft/nemotron_nano_9b.yaml create mode 100644 tests/test_sbatch.py diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index 79be005..d3603af 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -41,4 +41,7 @@ jobs: - name: Test run: | - .venv/bin/pytest tests/ -v --tb=short + # --no-sync mirrors the GitLab CI invocation and prevents `uv run` + # from implicitly auto-installing heavy core deps that we + # deliberately skipped above. + uv run --no-sync pytest tests/ -v --tb=short diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 5c0c12a..947952c 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -59,7 +59,10 @@ test: - uv venv --python 3.12 - uv pip install -e ".[dev]" --no-deps - uv pip install pytest pytest-cov pytest-xdist pytest-timeout PyYAML omegaconf rich - - uv run pytest tests/ -v --tb=short + # --no-sync prevents `uv run` from implicitly auto-installing heavy core + # deps (nemo-skills etc) that we deliberately skipped above. Without it, + # tests would pass here for the wrong reason and fail on GitHub CI. + - uv run --no-sync pytest tests/ -v --tb=short rules: - if: $CI_PIPELINE_SOURCE == "merge_request_event" - if: $CI_COMMIT_BRANCH == "main" diff --git a/INSTALL.md b/INSTALL.md index 34840c8..4b7f59f 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -1,6 +1,6 @@ # Installation & Setup Guide -Quick setup guide for NVFlow - a lightweight orchestration tool for Slurm clusters. +Quick setup guide for NVFlow - a lightweight orchestration tool for Slurm clusters. NVFlow's containers are self-sufficient โ€” all dependencies are pre-installed, so no runtime downloads are needed. Once images and models are staged, the pipeline runs fully offline. ## ๐Ÿ“‹ Steps @@ -17,17 +17,27 @@ Quick setup guide for NVFlow - a lightweight orchestration tool for Slurm cluste > **Note:** This guide assumes you've already completed the [README.md](README.md) setup (installed `uv`, cloned the repo, ran `uv sync`). +### Build Host Requirements (for building container images) + +The `docker build` step needs **internet access** to pull base layers, source from GitHub, and packages from PyPI / NGC / Docker Hub. The resulting `.sqsh` files then run fully offline on the cluster. + +- **Docker Engine** or **Docker Desktop** (any OS - Linux, macOS, Windows/WSL2) +- **`docker login nvcr.io`** - required once for the NeMo-RL base image +- **`docker buildx`** - only needed for multi-arch / cross-arch builds (ships with Docker Desktop; on Linux: `docker buildx version`) + +> **Note:** If your destination cluster is `linux/amd64` (the common case) and your build host is amd64 Linux / Intel macOS / Windows, the default `docker build` works without `buildx`. + ### Cluster Setup Requirements **Required:** -- **Slurm cluster** access with SSH keys +- **Slurm cluster** access with SSH keys (or run directly from a login node) +- **enroot** - on cluster nodes (`enroot version`) -**Only needed for converting Docker images to .sqsh format:** +**Only needed for the parallel container conversion script:** - **yq** - YAML parser ([install guide](https://github.com/mikefarah/yq)) - **curl** - for downloading configs -- **enroot** - on cluster nodes (for container conversion) -> **Note:** If you already have `.sqsh` container images, skip to [Configure Your Cluster](#configure-your-cluster). +> **Note:** If you already have `.sqsh` container images staged on the cluster, skip to [Configure Your Cluster](#configure-your-cluster). **yq (YAML parser):** ```bash @@ -59,23 +69,26 @@ enroot version # Run on cluster node **Get your cluster info:** - Slurm account: `sacctmgr show associations user=$USER` (look for the Account column) - Available partitions: `sinfo` +- Slurm version: `scontrol show config | grep SLURM_VERSION` (25.x needs the enroot Ray template fix - see [Troubleshooting](#troubleshooting)) - Storage paths for data/models/containers --- ## Setup Containers -NeMo-Skills requires Docker containers converted to `.sqsh` format for running on Slurm clusters. +NVFlow uses five containers converted to `.sqsh` format for running on Slurm clusters. **Four are built locally** from self-contained Dockerfiles in [`dockerfiles/`](dockerfiles/); the fifth (`sglang`) is pulled as-is. -**Required containers (4):** +**Required containers (5):** | Container | Source | Tested Version | Action | |-----------|--------|----------------|--------| -| `nemo-skills` | NeMo-Skills Dockerfiles | NeMo-Skills @ `0229040` | **Build** (see Step 1a) | -| `vllm` | Docker Hub | `vllm/vllm-openai:v0.18.1` | **Pull** (standalone SDG/eval) | -| `vllm-grpo` | Docker Hub | `vllm/vllm-openai:v0.17.1` | **Pull** (standalone GRPO rollouts/judge) | -| `sglang` | Docker Hub | `lmsysorg/sglang:v0.5.10.post1` | **Pull** (no build needed) | -| `nemo-rl` | NGC | `nvcr.io/nvidia/nemo-rl:v0.6.0` | **Pull** from NGC (no build needed) | +| `nvflow-nemo-rl` | [`dockerfiles/Dockerfile.nemo-rl`](dockerfiles/Dockerfile.nemo-rl) | base `nvcr.io/nvidia/nemo-rl:v0.6.0` | **Build** (see Step 1) | +| `nvflow-nemo-skills` | [`dockerfiles/Dockerfile.nemo-skills`](dockerfiles/Dockerfile.nemo-skills) | NeMo-Skills @ `0229040` | **Build** (see Step 1) | +| `nvflow-vllm` | [`dockerfiles/Dockerfile.vllm`](dockerfiles/Dockerfile.vllm) | base `vllm/vllm-openai:v0.18.1` | **Build** (SDG/eval) | +| `nvflow-vllm-grpo` | [`dockerfiles/Dockerfile.vllm-grpo`](dockerfiles/Dockerfile.vllm-grpo) | base `vllm/vllm-openai:v0.17.1` | **Build** (GRPO rollouts/judge) | +| `sglang` | Docker Hub | `lmsysorg/sglang:v0.5.10.post1` | **Pull** (no custom Dockerfile) | + +> **Note:** All four custom images are **built**, not pulled. The four Dockerfiles bake in NeMo-Skills source, NeMo-Gym source, pre-built virtual environments, `tiktoken` / `openai_harmony` encoding caches, and `/root/.local โ†’ /opt/uv-python` path relocation so the images run cleanly under `enroot`/`pyxis` on Slurm with no outbound network access. **Optional containers** (not currently used by any NVFlow recipes): @@ -86,109 +99,100 @@ NeMo-Skills requires Docker containers converted to `.sqsh` format for running o | `verl` | NeMo-Skills Dockerfiles | Build | | `trtllm` | `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc8` | Pull from NGC | -### Step 1a: Build NeMo-Skills Containers +### Step 1: Build Docker Images -Clone the NeMo-Skills repo at the **exact commit pinned by NVFlow** to ensure compatibility. The pinned commit is defined in [`pyproject.toml`](pyproject.toml): +NVFlow ships self-contained Dockerfiles in [`dockerfiles/`](dockerfiles/) that pre-install all Python packages, pre-cache tokenizer encodings, and pre-build virtual environments. Run `docker build` on a connected build host: ```bash -# Clone NeMo-Skills and check out the pinned commit -git clone https://github.com/NVIDIA/NeMo-Skills.git -cd NeMo-Skills -git checkout 022904023ad7a83a87662a313cf72e7df5891d55 -``` +cd /path/to/nvflow -> **Tip:** Always use the commit hash from `pyproject.toml` (search for `nemo-skills @`). Building from a different version may cause incompatibilities. +# Build all four custom images (amd64, the common case) +docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:v0.6.0 . +docker build -f dockerfiles/Dockerfile.nemo-skills -t nvflow-nemo-skills:0229040 . +docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:v0.18.1 . +docker build -f dockerfiles/Dockerfile.vllm-grpo -t nvflow-vllm-grpo:v0.17.1 . -Build the `nemo-skills` container using the [NeMo-Skills Dockerfiles](https://github.com/NVIDIA/NeMo-Skills/tree/022904023ad7a83a87662a313cf72e7df5891d55/dockerfiles): +# sglang is pulled as-is, no custom Dockerfile +docker pull lmsysorg/sglang:v0.5.10.post1 +``` -```bash -# Build with the helper script -./dockerfiles/build.sh dockerfiles/Dockerfile.nemo-skills +> **Tip:** The Dockerfiles expose `ARG`s for version pins (`NEMO_SKILLS_COMMIT`, `NEMO_GYM_BRANCH`, `VLLM_VERSION`, `BASE_IMAGE`). Defaults are listed in [`dockerfiles/README.md`](dockerfiles/README.md#version-pins). Keep `NEMO_SKILLS_COMMIT` consistent across `Dockerfile.nemo-skills`, `Dockerfile.nemo-rl`, and `pyproject.toml`. -# Or build directly with docker -docker build -t nemo-skills:latest -f dockerfiles/Dockerfile.nemo-skills . -``` +For **cross-arch builds** (e.g. building an `amd64` image on Apple Silicon, or a multi-arch manifest list pushed directly to a registry), see [`dockerfiles/docker_instructions.md`](dockerfiles/docker_instructions.md#1-build). Cross-arch builds use `docker buildx` with QEMU emulation and are significantly slower than native. -For `vllm`, `vllm-grpo`, and `sglang`, pull pre-built images directly from Docker Hub (no build needed): +For optional containers (`megatron`, `sandbox`, `verl`), build them from the upstream [NeMo-Skills Dockerfiles](https://github.com/NVIDIA-NeMo/Skills/tree/022904023ad7a83a87662a313cf72e7df5891d55/dockerfiles). -```bash -docker pull vllm/vllm-openai:v0.18.1 # standalone for SDG/eval -docker pull vllm/vllm-openai:v0.17.1 # standalone for GRPO rollouts/judge -docker pull lmsysorg/sglang:v0.5.10.post1 -``` +### Step 1b: Sanity-Check Images Before Conversion -For optional containers (`megatron`, `sandbox`, `verl`), build them the same way using their respective Dockerfiles. For arm64 builds, see the [multi-platform instructions](https://github.com/NVIDIA/NeMo-Skills/tree/022904023ad7a83a87662a313cf72e7df5891d55/dockerfiles#building-for-arm64aarch64). +Before the time-consuming `enroot import` step, run the smoke checks in [`dockerfiles/docker_instructions.md` ยง2](dockerfiles/docker_instructions.md#2-sanity-checks-blockers). Each check is a **hard blocker** - if it fails locally, the image will not work in production. They verify the offline-critical pieces: `uv` works offline, the 6 Gym `.venv` symlinks are intact, `tiktoken` / `openai_harmony` caches load with `--network=none`, and `tzdata` is populated. -### Step 1b: Pull NeMo-RL Container (for SFT and GRPO) +### Step 2: Push Images to a Registry (Option A) or Save as Tarball (Option B) -The `nemo-rl` container is available as a pre-built image on NGC: +`enroot` runs on Slurm compute/login nodes (Linux only). There are two paths from a Docker image to a `.sqsh` file - pick whichever fits your offline workflow. -```bash -docker pull nvcr.io/nvidia/nemo-rl:v0.6.0 -``` +#### Option A: Via a private container registry (recommended) -Alternatively, build from source using the [NeMo-RL repository](https://github.com/NVIDIA-NeMo/RL): +Push the built images to a registry accessible from your cluster (Docker Hub, NGC, or a private registry): ```bash -git clone https://github.com/NVIDIA-NeMo/RL.git -cd RL -git checkout v0.6.0 -git submodule update --init --recursive -``` +REGISTRY= -Follow the [NeMo-RL Docker build instructions](https://github.com/NVIDIA-NeMo/RL/blob/main/docs/docker.md#building-the-release-image) to build the release image, then tag and push it to your registry alongside the NeMo-Skills containers. +docker tag nvflow-nemo-rl:v0.6.0 $REGISTRY/nvflow-nemo-rl:v0.6.0 +docker tag nvflow-nemo-skills:0229040 $REGISTRY/nvflow-nemo-skills:0229040 +docker tag nvflow-vllm:v0.18.1 $REGISTRY/nvflow-vllm:v0.18.1 +docker tag nvflow-vllm-grpo:v0.17.1 $REGISTRY/nvflow-vllm-grpo:v0.17.1 -### Step 2: Push Images to a Registry +docker push $REGISTRY/nvflow-nemo-rl:v0.6.0 +docker push $REGISTRY/nvflow-nemo-skills:0229040 +docker push $REGISTRY/nvflow-vllm:v0.18.1 +docker push $REGISTRY/nvflow-vllm-grpo:v0.17.1 -After building, push the images to a container registry accessible from your cluster (Docker Hub, NGC, or a private registry): +# sglang can be pulled directly by enroot (no push needed unless your +# cluster cannot reach Docker Hub). +``` -```bash -# Tag and push the NeMo-Skills container -docker tag nemo-skills:latest your-registry/nemo-skills:latest -docker push your-registry/nemo-skills:latest +> **Why push?** Slurm cluster nodes typically don't have Docker installed, so `enroot` needs to pull images from a registry (or load them from a Docker daemon - see Option B). -# Tag and push vllm (pulled from Docker Hub) -docker tag vllm/vllm-openai:v0.18.1 your-registry/nemo-skills-vllm:latest -docker push your-registry/nemo-skills-vllm:latest +#### Option B: Via a saved tarball (no registry required) -# Tag and push NeMo-RL (pulled from NGC) -docker tag nvcr.io/nvidia/nemo-rl:v0.6.0 your-registry/nemo-skills-nemo-rl:latest -docker push your-registry/nemo-skills-nemo-rl:latest +For offline sites without a private registry, save the Docker image to a tarball, transfer it to a Linux host that has both Docker and `enroot`, load the tarball into the local Docker daemon, then import via `dockerd://`: -# sglang can be pulled directly by enroot (no push needed unless your -# cluster cannot reach Docker Hub) +```bash +# On the build host +docker save nvflow-nemo-rl:v0.6.0 | gzip > nvflow-nemo-rl-v0.6.0.tar.gz +docker save nvflow-nemo-skills:0229040 | gzip > nvflow-nemo-skills-0229040.tar.gz +docker save nvflow-vllm:v0.18.1 | gzip > nvflow-vllm-v0.18.1.tar.gz +docker save nvflow-vllm-grpo:v0.17.1 | gzip > nvflow-vllm-grpo-v0.17.1.tar.gz -# Repeat for any optional images you built (e.g., megatron, sandbox, verl) +# Transfer the .tar.gz files to the cluster (scp / rsync / sneakernet) ``` -> **Why push?** Slurm cluster nodes typically don't have Docker installed, so `enroot` needs to pull images from a registry. Pushing to a registry also lets the automated setup script work. +`enroot import` natively supports only `docker://` (remote registry), `dockerd://` (local Docker daemon), and `podman://` URIs. If the cluster has neither a private registry nor a Docker daemon, run a transient local registry container, push to it, and import via `docker://localhost:5000/...`. -### Step 3: Create Your Container Config +### Step 3: Update Container Config -`containers.yaml` is a **template** with placeholder values -- do not edit it directly. Instead, copy it to a personal file and fill in your registry paths: +Copy the template to a personal file that records the registry / tag references the cluster should pull from: ```bash cp cluster_configs/containers.yaml cluster_configs/my_containers.yaml ``` -Edit `my_containers.yaml` with your actual registry paths: +Edit `cluster_configs/my_containers.yaml` with your registry paths. The YAML **keys** (`nemo-skills`, `nemo-rl`, `vllm`, `vllm-grpo`, `sglang`) match what the workflow code references and must not be renamed; only the registry / tag values change: ```yaml containers: - nemo-skills: your-registry/nemo-skills:latest - vllm: your-registry/nemo-skills-vllm:latest # v0.18.1 for SDG/eval - vllm-grpo: vllm/vllm-openai:v0.17.1 # v0.17.1 for GRPO rollouts/judge - nemo-rl: nvcr.io/nvidia/nemo-rl:v0.6.0 # or your-registry/nemo-skills-nemo-rl:latest - sglang: lmsysorg/sglang:v0.5.10.post1 + nemo-rl: your-registry/nvflow-nemo-rl:v0.6.0 + nemo-skills: your-registry/nvflow-nemo-skills:0229040 + vllm: your-registry/nvflow-vllm:v0.18.1 # v0.18.1 for SDG/eval + vllm-grpo: your-registry/nvflow-vllm-grpo:v0.17.1 # v0.17.1 for GRPO rollouts/judge + sglang: lmsysorg/sglang:v0.5.10.post1 ``` > **Note:** `my_containers.yaml` is gitignored (`cluster_configs/*.yaml` pattern), so your registry paths stay local and won't be committed. ### Step 4: Convert to .sqsh Format -Choose one of the following methods to convert your container images to `.sqsh` format for Slurm. - -#### Option A: Automated Setup (Recommended) +#### Option A: Automated Setup (Recommended, for Option A registries) Use the setup script to download from your registry and convert all containers in parallel. Pass your personal config with `--config`: @@ -197,7 +201,7 @@ Use the setup script to download from your registry and convert all containers i sbatch --account=YOUR_ACCOUNT scripts/setup_containers.sh --config cluster_configs/my_containers.yaml ./containers ``` -The `--config` flag is required -- the script reads image references from the specified YAML file, pulls them via `enroot`, and converts to `.sqsh` format. See [the script](scripts/setup_containers.sh) for additional options (`--platform`, `--force`). +The `--config` flag is required - the script reads image references from the specified YAML file, pulls them via `enroot`, and converts to `.sqsh` format. See [the script](scripts/setup_containers.sh) for additional options (`--platform`, `--force`). **Check progress:** ```bash @@ -206,18 +210,37 @@ tail -f outputs/logs/slurm-containers-.out #### Option B: Manual Conversion -Convert images one at a time using `enroot` on a cluster node: +Convert images one at a time using `enroot` on a cluster node. From a registry, use `docker://$REGISTRY/...`; from a loaded tarball, use `dockerd://...` after `docker load`: ```bash -# Import from your registry -enroot import docker://your-registry/nemo-skills:latest -enroot import docker://your-registry/nemo-skills-vllm:latest -enroot import docker://nvcr.io/nvidia/nemo-rl:v0.6.0 - -# Import from official registries -enroot import docker://vllm/vllm-openai:v0.18.1 -enroot import docker://vllm/vllm-openai:v0.17.1 -enroot import docker://lmsysorg/sglang:v0.5.10.post1 +CONTAINER_DIR= + +enroot import --output $CONTAINER_DIR/nvflow-nemo-rl-v0.6.0.sqsh \ + "docker://$REGISTRY/nvflow-nemo-rl:v0.6.0" # from a registry +# -- or -- +gunzip -c nvflow-nemo-rl-v0.6.0.tar.gz | docker load +enroot import --output $CONTAINER_DIR/nvflow-nemo-rl-v0.6.0.sqsh \ + dockerd://nvflow-nemo-rl:v0.6.0 # from a tarball +``` + +Repeat for `nemo-skills`, `vllm`, `vllm-grpo`, and (if needed) `sglang`. + +**Two things to watch for:** + +- **Registries with a path component need `#` instead of `/`.** `enroot` parses `docker:///` such that everything after the first `/` is image path, which breaks for registries where the host itself contains a path (e.g. `nvcr.io/`). Use `#` to separate host from image path: + ```bash + enroot import --output nvflow-vllm-v0.18.1.sqsh \ + "docker://nvcr.io#/nvflow-vllm:v0.18.1" + ``` +- **Filename colon.** `enroot` writes the Docker tag separator (`:`) literally into the output filename. Either pass `--output` with a shell-safe name (as above) or rename after import: + ```bash + mv "nvflow-nemo-rl:v0.6.0.sqsh" nvflow-nemo-rl-v0.6.0.sqsh + ``` + +If the cluster authenticates to your registry, drop credentials into `~/.config/enroot/.credentials`: + +``` +machine login password ``` Move the resulting `.sqsh` files to your cluster's container storage path. @@ -226,7 +249,7 @@ Move the resulting `.sqsh` files to your cluster's container storage path. ## Download Models -> โš ๏ธ **Important:** Pre-download models to your cluster storage before running workflows. +> โš ๏ธ **Important:** Pre-download models to your cluster storage before running workflows. The runtime sets `HF_HUB_OFFLINE=1`, so any model not already on disk will fail at job time. > > **Why this matters:** > - Avoids wasting expensive GPU time on downloads @@ -303,54 +326,63 @@ stage_kwargs: **Tip:** Download commonly used models once and reuse across all workflows. +### One-Time Connected-Node Stages (Datasets) + +A handful of stages legitimately need internet on **first** run to pull benchmark / seed datasets from HuggingFace or SEC EDGAR. Run them on a connected node (or off-cluster) and ship the resulting artifacts to the cluster - they're reused by every subsequent run. + +| Stage | Pulls from | Why | +|---|---|---| +| `workflow-2 download_sec_filings` | SEC EDGAR | Filings aren't on HF | +| `workflow-3 step-0 create_seed_data` | HF `nogabenyoash/SecQue` | Seed dataset | +| `workflow-1 step-0 prepare_data` (eval) | HF `secque`, `financebench` | Benchmark data | +| `workflow-5 step-4 prepare_data` (GRPO) | HF | Only if `should_download: true` | + +For these stages, **temporarily clear** the three HF offline flags (`HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE`, `TRANSFORMERS_OFFLINE`) in your cluster config. Keep `UV_OFFLINE=true` set - `uv` should never need to resolve packages at runtime. + +> **Note:** `huggingface_hub` interprets `TRANSFORMERS_OFFLINE=1` as `HF_HUB_OFFLINE=1`, so all three need to be off (or unset) for HF dataset pulls to succeed. + --- ## Setup NeMo-RL & NeMo-Gym Sources (for GRPO) -> **Skip this section** if you're only running SDG/eval workflows. This setup is needed for GRPO RL training and recommended for multi-node SFT. +> **Skip this section** if you're using the self-sufficient `nvflow-nemo-rl` image as-is (the recommended path). The image already contains the NeMo-RL source, a pinned NeMo-Gym branch, and a pre-built `.venv` symlinked across all 6 Gym components. No host clones or overlay mounts are required for SDG, SFT, GRPO, or eval workflows. -Both NeMo-RL and NeMo-Gym source trees are overlay-mounted into the NeMo-RL container via cluster config mounts. This ensures the container uses the exact tested source code. +This section is **dev mode only** - read it only if you're actively iterating on NeMo-RL or NeMo-Gym source against the self-sufficient image. -### NeMo-RL Source Clone +### What the self-sufficient image already contains -Mount the NeMo-RL source into the container at `/opt/NeMo-RL`. If you already cloned it in [Step 1b](#step-1b-build-nemo-rl-container-for-sft-and-grpo), reuse that clone: +`nvflow-nemo-rl` is built from [`dockerfiles/Dockerfile.nemo-rl`](dockerfiles/Dockerfile.nemo-rl) on top of `nvcr.io/nvidia/nemo-rl:v0.6.0` and bakes in: -```bash -# Reuse the clone from Step 1b, or: -git clone https://github.com/NVIDIA-NeMo/RL.git -cd RL -git checkout v0.6.0 -git submodule update --init --recursive -``` +- NeMo-Skills @ `0229040` installed into the frozen `/opt/nemo_rl_venv` +- NeMo-Gym source at `/opt/NeMo-RL/3rdparty/Gym-workspace/Gym`, checked out at the `ude/finance-sec-search-v2` branch (override via `NEMO_GYM_BRANCH` build arg) +- A pre-built Gym `.venv` symlinked across all 6 components (`equivalence_llm_judge`, `finance_sec_search`, `simple_agent`, `finance_agent`, `openai_model`, `vllm_model`) +- `/root/.local/share/uv/python` relocated to `/opt/uv-python` and `/root/.local/bin` to `/opt/uv-bin` so the venvs survive enroot/pyxis mounting `$HOME` over `/root` +- `tiktoken` / `openai_harmony` encoding caches at `/opt/tiktoken_cache` -### NeMo-Gym Clone +GRPO stages call `installation_command: source /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate` and find everything they need inside the image. -Clone NeMo-Gym and mount it inside the NeMo-RL source tree. The Gym overlay is independent from the Gym submodule inside RL -- this lets RL and Gym evolve on separate branches. Check your workflow config (e.g., `grpo/base.yaml`) for the tested Gym branch or commit: +### Do NOT overlay-mount source over the image paths -```bash -git clone https://github.com/NVIDIA-NeMo/Gym.git -cd Gym -git checkout ude/finance-sec-search # finance_agent environment (until merged to main) -``` +Bind-mounting a host clone at `/opt/NeMo-RL` or `/opt/NeMo-RL/3rdparty/Gym-workspace/Gym` **shadows the baked `.venv`**, and `installation_command` fails with `No such file or directory` - breaking `prepare_data`, `collect_rollouts`, `compute_rewards`, and `training` for GRPO. -### Cluster Config Mounts - -Add both overlay mounts to your cluster config (see [Configure Your Cluster](#configure-your-cluster)): +The older dev-mode overlay snippets in `template-slurm.yaml` are commented out for exactly this reason: ```yaml mounts: - - /path/to/RL:/opt/NeMo-RL - - /path/to/Gym:/opt/NeMo-RL/3rdparty/Gym-workspace/Gym + # DO NOT use these with the self-sufficient image โ€” they shadow the baked .venv + # - :/opt/NeMo-RL + # - :/opt/NeMo-RL/3rdparty/Gym-workspace/Gym ``` -No local `uv sync` is needed for either -- the container's `installation_command` handles dependency installation at runtime. +### Dev mode: iterating on NeMo-RL / NeMo-Gym source + +If you really need to iterate on NeMo-RL or NeMo-Gym source against this image, clone the source trees (NeMo-RL at `v0.6.0` with submodules, NeMo-Gym at `ude/finance-sec-search-v2`), uncomment the two overlay mounts in `cluster_configs/my_cluster.yaml`, and set `NRL_FORCE_REBUILD_VENVS=true` in `env_vars` so Ray workers rebuild their venvs against your source. Your host clone must contain a `.venv` ABI-compatible with the image, and `NRL_FORCE_REBUILD_VENVS=true` **requires internet** at job time โ€” only use it on a connected node, never in production. ### Prefetch SEC Filings Cache (for `finance_sec_search`) If using the `finance_sec_search` NeMo-Gym environment, you must prefetch the SEC filings cache to a shared mounted path. The default `~/.cache` does **not** work inside Slurm containers. -1. Follow the prefetch instructions in `Gym/resources_servers/finance_sec_search/README.md` -2. Set `cache_dir` in the environment config overlay to point to the shared mounted path +The GRPO workflow includes a dedicated `prefetch_cache` stage that runs on a connected node and populates the cache under your `workflow-5-grpo/` output directory. See [`docs/recipes/finance/workflows/06-grpo.md`](docs/recipes/finance/workflows/06-grpo.md) for the full prefetch flow. --- @@ -370,9 +402,33 @@ Edit `cluster_configs/my_cluster.yaml` and replace all `` values: 1. **SSH settings** - Your cluster login node, username, SSH key path (ONLY for remote job submission from local machine) 2. **Slurm account/partition** - Run `sacctmgr show associations user=$USER` and `sinfo` 3. **Container paths** - Copy from `outputs/logs/slurm-containers-.out` after running setup_containers.sh -4. **Mount points** - Map your cluster paths to container paths +4. **Mount points** - Map your cluster paths to container paths (at minimum `:/hf_models` and `:/workspace`) 5. **Environment variables** - Set `HF_HOME` to a path visible inside the container (see [env_vars docs](docs/cluster-configuration.md#environment-variables)) and any API keys +### Step 3: Keep the Air-Gap Enforcement Block + +`template-slurm.yaml` ships with the offline flags pre-populated - leave them on: + +```yaml +env_vars: + # --- AIR-GAPPED ENFORCEMENT (recommended) --- + - HF_HUB_OFFLINE=1 + - HF_DATASETS_OFFLINE=1 + - TRANSFORMERS_OFFLINE=1 + - UV_OFFLINE=true + + # Pre-baked tiktoken / openai_harmony cache (set as ENV in vllm/vllm-grpo + # already; setting here applies them uniformly to nemo-skills and nemo-rl) + - TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache + + # Do NOT set in self-sufficient mode โ€” forces Ray workers to re-resolve via uv + # - NRL_FORCE_REBUILD_VENVS=true +``` + +For the one-time connected-node stages listed in [Download Models](#one-time-connected-node-stages-datasets) above, comment out the three `HF_*_OFFLINE` flags just for that submission, then re-enable. + > **Note:** The template includes detailed comments for each section. Your personal config (`my_cluster.yaml`) is gitignored to protect secrets. > > ๐Ÿ“– **For detailed documentation of all configuration fields, see the [Cluster Configuration Guide](docs/cluster-configuration.md)**. @@ -388,7 +444,7 @@ uv run python -c "from nemo_skills.pipeline.cli import generate; print('โœ… OK') # 2. Check containers exist ls -lh /*.sqsh -# 3. Test SSH to cluster +# 3. Test SSH to cluster (only if submitting from a local machine) ssh -i @ "echo 'โœ… SSH OK'" # 4. Test cluster config loads @@ -437,25 +493,26 @@ module load enroot # if available # Or contact your cluster admin ``` +### Docker build fails on `docker login` for NeMo-RL base image +The `nvflow-nemo-rl` build pulls from `nvcr.io/nvidia/nemo-rl:v0.6.0` (NGC). Run `docker login nvcr.io` once (username `$oauthtoken`, password = your [NGC API key](https://ngc.nvidia.com/setup/api-key)). + +### `enroot import` quirks (filename colon, `#` separator for `nvcr.io`) +See [Two things to watch for](#step-4-convert-to-sqsh-format) in Step 4. + ### SSH connection failed ```bash -# Check key permissions chmod 600 - -# Test manual connection ssh -i @ ``` ### Container paths wrong -- Use **absolute paths** in cluster config -- Check file exists: `ls -l /.sqsh` -- Re-run container setup if needed +Use **absolute paths** in cluster config and verify each `.sqsh` exists (`ls -l /*.sqsh`). Re-run container setup if needed. ### HF_HOME / cache "No such file or directory" -- `HF_HOME` (and other path-valued env vars) must resolve to a path **visible inside the container** -- Use a mount destination (e.g., `/workspace/cache/huggingface`) or a host path that is transparently mounted (e.g., `/shared/.../cache` when `- /shared:/shared` is in `mounts`) -- Paths that exist only on the host and have no corresponding mount will fail with `No such file or directory` -- Common mistake: using `$HOME` or `~/.cache` -- these do not resolve inside containers unless explicitly mounted +`HF_HOME` (and every other path-valued env var) must resolve **inside the container** -- use a mount destination (e.g. `/workspace/cache/huggingface`) or a transparently-mounted host path (e.g. `/shared/...` when `- /shared:/shared` is in `mounts`). `$HOME` and `~/.cache` will not resolve. + +### Air-gapped runtime errors at job time +For symptoms specific to the self-sufficient runtime - GRPO `installation_command` failing with "No such file or directory", `OfflineModeIsEnabled`, `uv` trying to reach PyPI, `tiktoken` / `openai_harmony` failing offline - see the [Offline Runtime](docs/recipes/finance/troubleshooting.md#offline-runtime) section in the finance troubleshooting guide. ### Slurm jobs won't submit - Verify account: `sacctmgr show associations user=$USER` @@ -506,11 +563,14 @@ Then head back to the [README.md](README.md#-quick-start) Quick Start section to ## Reference -- **NeMo-Skills**: https://github.com/NVIDIA/NeMo-Skills -- **NeMo-Skills Dockerfiles**: https://github.com/NVIDIA/NeMo-Skills/tree/022904023ad7a83a87662a313cf72e7df5891d55/dockerfiles +- **NVFlow Dockerfiles**: [`dockerfiles/README.md`](dockerfiles/README.md) +- **NVFlow Self-Sufficient Build / Deploy Guide**: [`dockerfiles/docker_instructions.md`](dockerfiles/docker_instructions.md) +- **Cluster Configuration Guide**: [`docs/cluster-configuration.md`](docs/cluster-configuration.md) +- **NeMo-Skills**: https://github.com/NVIDIA-NeMo/Skills +- **NeMo-Skills Dockerfiles (upstream reference)**: https://github.com/NVIDIA-NeMo/Skills/tree/022904023ad7a83a87662a313cf72e7df5891d55/dockerfiles - **NeMo-RL**: https://github.com/NVIDIA-NeMo/RL -- **NeMo-RL Docker Build**: https://github.com/NVIDIA-NeMo/RL/blob/main/docs/docker.md#building-the-release-image -- **Official Container Config**: https://github.com/NVIDIA/NeMo-Skills/blob/main/cluster_configs/example-slurm.yaml +- **NeMo-Gym**: https://github.com/NVIDIA-NeMo/Gym +- **Official Container Config (NeMo-Skills)**: https://github.com/NVIDIA-NeMo/Skills/blob/main/cluster_configs/example-slurm.yaml - **Slurm Docs**: https://slurm.schedmd.com/ - **Enroot**: https://github.com/NVIDIA/enroot diff --git a/README.md b/README.md index c995b75..2640664 100644 --- a/README.md +++ b/README.md @@ -123,9 +123,18 @@ pytest ## ๐Ÿ”ง Cluster Setup -To run workflows on a Slurm cluster, you need to configure containers and cluster settings. - -> **See [INSTALL.md](INSTALL.md)** for complete cluster setup (containers, cluster configuration, verification). +To run workflows on a Slurm cluster you need to: (1) build the four NVFlow +container images from the Dockerfiles in [`dockerfiles/`](dockerfiles/), +(2) convert them to `.sqsh` for Slurm, and (3) write a cluster config +(`cluster_configs/my_cluster.yaml`). The containers are self-sufficient โ€” +all dependencies are pre-installed, so no runtime downloads are needed. + +> **See [INSTALL.md](INSTALL.md)** for the complete setup guide +> (build, sanity-check, `.sqsh` conversion, model staging, cluster +> configuration, and verification). +> +> **See [`dockerfiles/docker_instructions.md`](dockerfiles/docker_instructions.md)** +> for the build / multi-arch / sanity-check reference. Once cluster setup is complete, set the config directory: diff --git a/cluster_configs/containers.yaml b/cluster_configs/containers.yaml index b2780fd..14281b9 100644 --- a/cluster_configs/containers.yaml +++ b/cluster_configs/containers.yaml @@ -10,33 +10,47 @@ # # Edit my_containers.yaml with your registry paths # sbatch --account= scripts/setup_containers.sh --config cluster_configs/my_containers.yaml ./containers # +# Naming: +# YAML KEYS (nemo-skills, nemo-rl, vllm, vllm-grpo, sglang) match the +# short names the workflow code uses to look up containers -- do NOT rename +# them. Only the values (registry/tag refs) change between deployments. +# # Format: -# - Simple string: image reference (e.g., your-registry/nemo-skills:latest) +# - Simple string: image reference (e.g., your-registry/nvflow-nemo-skills:0229040) # - Nested object: Platform-specific tags (amd64/arm64 keys) +# +# Self-sufficient containers (all deps pre-installed, no runtime downloads): +# nemo-rl -> dockerfiles/Dockerfile.nemo-rl (base: nvcr.io/nvidia/nemo-rl:v0.6.0) +# nemo-skills -> dockerfiles/Dockerfile.nemo-skills (base: ubuntu:22.04, NeMo-Skills @ 0229040) +# vllm -> dockerfiles/Dockerfile.vllm (base: vllm/vllm-openai:v0.18.1) +# vllm-grpo -> dockerfiles/Dockerfile.vllm-grpo (base: vllm/vllm-openai:v0.17.1) +# sglang -> pulled as-is from Docker Hub (no custom Dockerfile) containers: # --------------------------------------------------------------------------- - # Required: Build from NeMo-Skills Dockerfiles (see INSTALL.md Step 1a) + # Required: Built locally from dockerfiles/, then pushed to your registry + # (see INSTALL.md Step 1 and Step 2). # --------------------------------------------------------------------------- - # Tested: NeMo-Skills @ 0229040 - nemo-skills: /nemo-skills: + # Tested: nvflow-nemo-rl:v0.6.0 (extends nvcr.io/nvidia/nemo-rl:v0.6.0) + nemo-rl: /nvflow-nemo-rl:v0.6.0 + + # Tested: nvflow-nemo-skills:0229040 (NeMo-Skills @ commit 0229040) + nemo-skills: /nvflow-nemo-skills:0229040 + + # Tested: nvflow-vllm:v0.18.1 (extends vllm/vllm-openai:v0.18.1; pre-cached + # tiktoken + openai_harmony; multi-arch amd64 + arm64) + vllm: /nvflow-vllm:v0.18.1 + + # Tested: nvflow-vllm-grpo:v0.17.1 (extends vllm/vllm-openai:v0.17.1; pinned + # to match NeMo-RL v0.6.0 colocated vLLM) + vllm-grpo: /nvflow-vllm-grpo:v0.17.1 # --------------------------------------------------------------------------- - # Required: Pull pre-built from Docker Hub / NGC (no build needed) + # Required: Pulled as-is from Docker Hub (no custom Dockerfile) # --------------------------------------------------------------------------- - # Tested: vllm/vllm-openai:v0.18.1 (standalone SDG/eval) - vllm: /nemo-skills-vllm: - # Tested: vllm/vllm-openai:v0.17.1 (standalone GRPO rollouts/judge) - vllm-grpo: vllm/vllm-openai:v0.17.1 # Tested: lmsysorg/sglang:v0.5.10.post1 sglang: lmsysorg/sglang:v0.5.10.post1 - # --------------------------------------------------------------------------- - # Required: Pull from NGC (see INSTALL.md Step 1b) - # --------------------------------------------------------------------------- - # Tested: nvcr.io/nvidia/nemo-rl:v0.6.0 - nemo-rl: nvcr.io/nvidia/nemo-rl:v0.6.0 - # --------------------------------------------------------------------------- # Optional: Not currently used by NVFlow recipes # Uncomment and update if needed for your workflows. diff --git a/cluster_configs/template-slurm.yaml b/cluster_configs/template-slurm.yaml index 077633a..86e14a3 100644 --- a/cluster_configs/template-slurm.yaml +++ b/cluster_configs/template-slurm.yaml @@ -5,12 +5,12 @@ # # Then update all values with your settings. # -# Container versions tested with this release: -# nemo-skills: NeMo-Skills @ 0229040 -# vllm: vllm/vllm-openai v0.18.1 (standalone SDG/eval) -# vllm-grpo: vllm/vllm-openai v0.17.1 (standalone GRPO rollouts/judge) +# Container versions tested with this release (self-sufficient, no runtime downloads): +# nemo-skills: nvflow-nemo-skills (NeMo-Skills @ 0229040) +# vllm: nvflow-vllm (base vllm/vllm-openai v0.18.1, standalone SDG/eval) +# vllm-grpo: nvflow-vllm-grpo (base vllm/vllm-openai v0.17.1, GRPO rollouts/judge) # sglang: lmsysorg/sglang v0.5.10.post1 -# nemo-rl: nvcr.io/nvidia/nemo-rl:v0.6.0 (includes vLLM 0.17.1 colocated) +# nemo-rl: nvflow-nemo-rl (base nvcr.io/nvidia/nemo-rl:v0.6.0, pre-built venvs + Gym) # # Reference: https://github.com/NVIDIA/NeMo-Skills @@ -79,12 +79,12 @@ extra_sandbox_args: # ============================================================================= # After converting containers to .sqsh format (see INSTALL.md), paste paths here. containers: - # Required containers - nemo-skills: /nemo-skills.sqsh # Orchestration client (eval, SDG, data prep) - vllm: /vllm.sqsh # vLLM v0.18.1 standalone (SDG, eval) - vllm-grpo: /vllm-grpo.sqsh # vLLM v0.17.1 standalone (GRPO rollouts, judge) - sglang: /sglang.sqsh # sglang inference server (SDG stages 3-5) - nemo-rl: /nemo-rl.sqsh # NeMo-RL v0.6.0 for SFT and GRPO training + # Required containers (self-sufficient โ€” all deps pre-installed, no runtime downloads) + nemo-skills: /nvflow-nemo-skills.sqsh # Orchestration client (eval, SDG, data prep) + vllm: /nvflow-vllm.sqsh # vLLM v0.18.1 standalone (SDG, eval) + vllm-grpo: /nvflow-vllm-grpo.sqsh # vLLM v0.17.1 standalone (GRPO rollouts, judge) + sglang: /sglang.sqsh # sglang inference server (SDG stages 3-5) + nemo-rl: /nvflow-nemo-rl.sqsh # NeMo-RL v0.6.0 for SFT and GRPO training # Optional containers (not currently used by NVFlow recipes) # trtllm: /trtllm.sqsh # megatron: /megatron.sqsh @@ -98,14 +98,16 @@ containers: mounts: - :/hf_models # HuggingFace models - :/workspace # Your workspace - # --- GRPO / RL Training (required for collect_rollouts and training stages) --- - # NeMo-RL source mount: overlays the container's built-in /opt/NeMo-RL. - - :/opt/NeMo-RL - # Gym overlay: mount your NeMo-Gym clone inside the NeMo-RL source tree. - # Harmless for SFT/SDG/eval -- only accessed by GRPO stages. - - :/opt/NeMo-RL/3rdparty/Gym-workspace/Gym # Add more mounts as needed: # - /lustre/data:/data + # + # --- NeMo-RL / NeMo-Gym source overlays (dev mode only) --- + # The nvflow-nemo-rl container includes NeMo-RL source and pre-built Gym + # venvs. Mounting host clones here shadows the container's venvs and + # breaks GRPO stages. Only uncomment for local source iteration with + # NRL_FORCE_REBUILD_VENVS=true enabled below. + # - :/opt/NeMo-RL + # - :/opt/NeMo-RL/3rdparty/Gym-workspace/Gym # ============================================================================= # Timeouts (per partition) @@ -131,10 +133,30 @@ env_vars: - TOKENIZERS_PARALLELISM=false # Disable HF tokenizer Rayon threads (prevents vLLM RefCell race) - VIRTUAL_ENV= # Unset to prevent host venv from interfering with container - VIRTUAL_ENV_PROMPT= # Unset venv prompt - # Rebuild Ray venvs when NeMo-RL source mount changes (e.g., new branch/commit). - # Without this, workers reuse stale cached venvs and may fail with import errors. - # Safe to leave enabled โ€” only rebuilds when the source tree actually changes. + + # --- Offline enforcement (recommended) --- + # Prevents accidental network calls at runtime. Containers are self-sufficient. + # Clear HF_*_OFFLINE temporarily for one-time stages that download external + # data (download_sec_filings, create_seed_data). UV_OFFLINE should stay set. + - HF_HUB_OFFLINE=1 + - HF_DATASETS_OFFLINE=1 + - TRANSFORMERS_OFFLINE=1 + - UV_OFFLINE=true + + # Pre-cached tiktoken encodings (baked into vllm/vllm-grpo containers; + # set here for uniform coverage across all container types). + - TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache + + # Ray worker port range โ€” avoids collisions with OS ephemeral ports. + - MIN_WORKER_PORT=7000 + - MAX_WORKER_PORT=8000 + + # Only enable in dev mode when iterating on NeMo-RL/Gym source overlays. + # Forces Ray workers to rebuild venvs from mounted source (requires internet). # - NRL_FORCE_REBUILD_VENVS=true + # API keys (keep these secret, don't commit to git!) # - HF_TOKEN= # - WANDB_API_KEY= diff --git a/dockerfiles/Dockerfile.nemo-rl b/dockerfiles/Dockerfile.nemo-rl new file mode 100644 index 0000000..56c21ad --- /dev/null +++ b/dockerfiles/Dockerfile.nemo-rl @@ -0,0 +1,131 @@ +# ============================================================================= +# NVFlow NeMo-RL Container +# ============================================================================= +# Extends the NeMo-RL nightly container with NeMo-Skills and the NeMo-Gym +# finance agent for NVFlow workflows. The base image ships with frozen +# environments and pre-built Ray venvs; this Dockerfile adds the NeMo-Skills +# package, replaces the Gym submodule with a feature branch that includes the +# finance-SEC-search resource server and finance agent, pre-builds all Gym +# component venvs, and relocates paths for Slurm/enroot compatibility. +# +# Build: +# docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:latest . +# ============================================================================= + +ARG BASE_IMAGE=nvcr.io/nvidia/nemo-rl:v0.6.0 +FROM ${BASE_IMAGE} + +# --- Symlink for NeMo-Skills code that references /opt/NeMo-RL (wrong case) -- +RUN ln -sf /opt/nemo-rl /opt/NeMo-RL + +# --- Upgrade uv (nemo-gym may require newer features than what the nightly ships) +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +# --- Pre-cache Python interpreter for uv (air-gapped safety net) ------------- +# Gym's cli_setup_command runs `uv venv --python `. If +# skip_venv_if_present ever misses, uv still needs a local interpreter. +RUN /root/.local/bin/uv python install 3.12 + +# --- Install NeMo-Skills into the frozen venv -------------------------------- +ARG NEMO_SKILLS_COMMIT=022904023ad7a83a87662a313cf72e7df5891d55 +RUN git clone https://github.com/NVIDIA-NeMo/Skills.git /opt/NeMo-Skills && \ + cd /opt/NeMo-Skills && git checkout ${NEMO_SKILLS_COMMIT} && \ + /root/.local/bin/uv pip install --python /opt/nemo_rl_venv/bin/python . + +# --- Replace NeMo-Gym submodule with feature branch ------------------------- +# The feature branch includes the finance-SEC-search resource server and +# finance agent that are not yet on main. +ARG NEMO_GYM_BRANCH=ude/finance-sec-search-v2 +RUN rm -rf /opt/nemo-rl/3rdparty/Gym-workspace/Gym && \ + git clone --branch ${NEMO_GYM_BRANCH} \ + https://github.com/NVIDIA-NeMo/Gym.git \ + /opt/nemo-rl/3rdparty/Gym-workspace/Gym + +# --- Pre-build Gym venv ------------------------------------------------------ +WORKDIR /opt/nemo-rl/3rdparty/Gym-workspace/Gym +RUN /root/.local/bin/uv venv .venv --python 3.12 && \ + . .venv/bin/activate && \ + /root/.local/bin/uv sync --active --extra dev + +# Install finance-specific dependencies into Gym venv +# uvicorn>=0.37.0 is required for timeout_worker_healthcheck support; +# uv sync resolves from the parent nemo-rl workspace lock (0.35.0) instead +# of the Gym lock, so we force the correct version here. +RUN . .venv/bin/activate && \ + /root/.local/bin/uv pip install aiohttp beautifulsoup4 "tavily==1.1.0" tenacity "uvicorn>=0.37.0" + +# --- Symlink component venvs to the main Gym venv --------------------------- +# Each NeMo-Gym component expects its own .venv/; symlinking avoids multi-GB +# duplication and guarantees every component runs with the same packages. +RUN for component in \ + resources_servers/equivalence_llm_judge \ + resources_servers/finance_sec_search \ + responses_api_agents/simple_agent \ + responses_api_agents/finance_agent \ + responses_api_models/openai_model \ + responses_api_models/vllm_model; do \ + dir="/opt/nemo-rl/3rdparty/Gym-workspace/Gym/$component"; \ + [ -d "$dir" ] && ln -sf /opt/nemo-rl/3rdparty/Gym-workspace/Gym/.venv "$dir/.venv"; \ + done + +WORKDIR / + +# --- Install Gym into the NemoGym Ray venv ------------------------------------ +# The pre-built Ray venv from the base image is stale (built from the old Gym +# submodule). Install the new Gym branch editable + all deps so the Ray actor +# can import nemo_gym without missing modules (e.g. gprof2dot, pydot). +RUN /root/.local/bin/uv pip install \ + --python /opt/ray_venvs/nemo_rl.environments.nemo_gym.NemoGym/bin/python \ + -e /opt/nemo-rl/3rdparty/Gym-workspace/Gym + +# --- Align numpy across all Ray venvs to match the main venv ---------------- +# NeMo-Skills may upgrade numpy; mismatched versions cause pickle failures +# when Ray serializes data between the main process and worker processes. +RUN MAIN_NP=$(/opt/nemo_rl_venv/bin/python -c "import numpy; print(numpy.__version__)") && \ + for venv in /opt/ray_venvs/*/; do \ + "$venv/bin/pip" install --no-cache-dir "numpy==$MAIN_NP" 2>/dev/null || true; \ + done + +# --- Relocate /root/.local/ โ†’ /opt/ ----------------------------------------- +# enroot/pyxis on Slurm mounts the user's home directory over /root at runtime, +# which shadows everything uv installed there during the Docker build. +# NOTE: Do NOT move /root/.cache/uv โ€” base-image venvs symlink into it. +RUN REAL_PYTHON=$(readlink /opt/nemo_rl_venv/bin/python) && \ + mv /root/.local/share/uv/python /opt/uv-python && \ + find /opt/uv-python -maxdepth 1 -type l | while read link; do \ + target=$(readlink "$link") && \ + new_target=$(echo "$target" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ + ln -sf "$new_target" "$link"; \ + done && \ + NEW_PYTHON=$(echo "$REAL_PYTHON" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ + ln -sf "$NEW_PYTHON" /opt/nemo_rl_venv/bin/python && \ + sed -i "s|/root/.local/share/uv/python|/opt/uv-python|g" /opt/nemo_rl_venv/pyvenv.cfg && \ + mv /root/.local/bin /opt/uv-bin + +# --- Fix pre-built Ray venvs (same /root/ relocation) ----------------------- +RUN for cfg in /opt/ray_venvs/*/pyvenv.cfg; do \ + sed -i "s|/root/.local/share/uv/python|/opt/uv-python|g" "$cfg"; \ + done && \ + find /opt/ray_venvs/ -type l | while read link; do \ + target=$(readlink "$link") && \ + case "$target" in */root/.local/share/uv/python*) \ + new_target=$(echo "$target" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ + ln -sf "$new_target" "$link" ;; \ + esac; \ + done + +# --- Fix Gym venv (same /root/ relocation) ---------------------------------- +RUN GYM_VENV=/opt/nemo-rl/3rdparty/Gym-workspace/Gym/.venv && \ + sed -i "s|/root/.local/share/uv/python|/opt/uv-python|g" "$GYM_VENV/pyvenv.cfg" && \ + find "$GYM_VENV" -type l | while read link; do \ + target=$(readlink "$link") && \ + case "$target" in */root/.local/share/uv/python*) \ + new_target=$(echo "$target" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ + ln -sf "$new_target" "$link" ;; \ + esac; \ + done + +# --- Runtime environment ----------------------------------------------------- +ENV VIRTUAL_ENV=/opt/nemo_rl_venv +ENV PATH=/opt/uv-bin:/opt/nemo_rl_venv/bin:$PATH +ENV UV_PYTHON_INSTALL_DIR=/opt/uv-python diff --git a/dockerfiles/Dockerfile.nemo-skills b/dockerfiles/Dockerfile.nemo-skills new file mode 100644 index 0000000..a67c6eb --- /dev/null +++ b/dockerfiles/Dockerfile.nemo-skills @@ -0,0 +1,138 @@ +# ============================================================================= +# NVFlow NeMo-Skills Container +# ============================================================================= +# Self-contained Dockerfile that builds the NeMo-Skills evaluation container +# with all required packages pre-installed. +# +# Build: +# docker build -f dockerfiles/Dockerfile.nemo-skills -t nvflow-nemo-skills:latest . +# +# Upstream source: +# https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.nemo-skills +# ============================================================================= + + +# Clone NeMo-Skills at a pinned commit (replaces build-context COPY commands) +FROM scratch AS nemo-skills-src +ARG NEMO_SKILLS_COMMIT=022904023ad7a83a87662a313cf72e7df5891d55 +ADD --keep-git-dir=true https://github.com/NVIDIA-NeMo/Skills.git#${NEMO_SKILLS_COMMIT} / + + +# =========================================================================== +# BEGIN UPSTREAM (adapted from NeMo-Skills Dockerfile.nemo-skills) +# Source: https://github.com/NVIDIA-NeMo/Skills/blob/0229040/dockerfiles/Dockerfile.nemo-skills +# Modifications: +# - COPY commands changed to COPY --from=nemo-skills-src +# - Added `tzdata` to apt packages (required by pyarrow/pandas; populates +# /usr/share/zoneinfo so libc tz lookups resolve, e.g. "UTC") +# =========================================================================== + +FROM ubuntu:22.04 + +RUN apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y \ + python3.10 \ + python3-pip \ + curl \ + wget \ + git \ + git-lfs \ + ffmpeg \ + tzdata && \ + ln -s /usr/bin/python3 /usr/bin/python && \ + rm -rf /var/cache/apt/archives /var/lib/apt/lists/* + +RUN pip install --upgrade pip setuptools uv + +RUN apt update && \ + apt install -y software-properties-common && \ + add-apt-repository -y ppa:apptainer/ppa && \ + apt update && apt -y install apptainer && \ + add-apt-repository -y ppa:apptainer/ppa && \ + apt update && apt install -y apptainer-suid && \ + rm -rf /var/cache/apt/archives /var/lib/apt/lists/* + +RUN mkdir /opt/benchmarks +RUN git clone https://github.com/google-research/google-research.git /opt/benchmarks/google-research --depth=1 + +RUN git clone https://github.com/ShishirPatil/gorilla.git /opt/gorilla +RUN cd /opt/gorilla && git checkout 86d0374d0db52623c5092a73f82c22b87b7e9a25 +RUN cd /opt/gorilla/berkeley-function-call-leaderboard && pip install --no-cache-dir -e . --extra-index-url https://download.pytorch.org/whl/cpu + +RUN apt remove -y python3-blinker + +ARG IFBENCH_COMMIT=c6767a19bd82ac0536cab950f2f8f6bcc6fabe7c +ARG IFBENCH_REPO=https://github.com/allenai/IFBench.git +ARG IFBENCH_DIR=/opt/benchmarks/IFBench +RUN git init "$IFBENCH_DIR" && cd "$IFBENCH_DIR" && git remote add origin "$IFBENCH_REPO" && \ + git fetch --depth 1 origin "${IFBENCH_COMMIT}" && git reset --hard FETCH_HEAD +RUN cd ${IFBENCH_DIR} && pip install -r requirements.txt + +COPY --from=nemo-skills-src /dockerfiles/ifbench.patch /opt/benchmarks/IFBench/ifbench.patch +RUN cd /opt/benchmarks/IFBench && git apply ifbench.patch + +RUN pip install langdetect absl-py immutabledict nltk ipython && \ + python -c "import nltk; from spacy.cli import download; nltk.download('punkt'); nltk.download('punkt_tab'); \ + nltk.download('stopwords'); nltk.download('averaged_perceptron_tagger_eng'); download('en_core_web_sm')" + +RUN mkdir -p /opt/NeMo-Skills/requirements /opt/NeMo-Skills/core +COPY --from=nemo-skills-src /pyproject.toml /opt/NeMo-Skills/pyproject.toml +COPY --from=nemo-skills-src /README.md /opt/NeMo-Skills/README.md +COPY --from=nemo-skills-src /requirements/ /opt/NeMo-Skills/requirements/ +COPY --from=nemo-skills-src /core/requirements.txt /opt/NeMo-Skills/core/requirements.txt +RUN pip install git+https://github.com/NVIDIA/NeMo-speech-data-processor@29b9b1ec0ceaf3ffa441c1d01297371b3f8e11d2 +ARG CACHEBUST=4 +RUN echo "httpx>=0.28.1" > /tmp/overrides.txt && \ + uv pip install --system --no-cache --override /tmp/overrides.txt \ + -r /opt/NeMo-Skills/core/requirements.txt \ + -r /opt/NeMo-Skills/requirements/pipeline.txt +RUN pip install ddgs + +# =========================================================================== +# END UPSTREAM +# =========================================================================== + + +# =========================================================================== +# NVFlow Additional Packages +# =========================================================================== +# Pre-install packages used by NVFlow workflow steps so they are available +# at runtime without needing to download anything. +# =========================================================================== + +RUN pip install --no-cache-dir --ignore-requires-python \ + jsonlines \ + tiktoken \ + markdownify \ + "backoff>=2.2.1" \ + func-timeout \ + "tavily==1.1.0" \ + "model-library==0.1.8" \ + "compute-eval @ git+https://github.com/NVIDIA/compute-eval.git@2d14770" + +# Pre-cache tiktoken encodings so no downloads are needed at runtime. +# cl100k_base is used by question_context_utils.py for token counting. +ENV TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache +RUN mkdir -p /opt/tiktoken_cache && \ + python3 -c "import tiktoken; tiktoken.get_encoding('cl100k_base')" + +# SEC data-prep dependencies for workflow-2 (download_sec_filings) and +# workflow-3 step-0 (create_seed_data). Baked in so the recipe-level +# installation_command can stay empty -- no PyPI fetch at job-launch time. +RUN pip install --no-cache-dir --ignore-requires-python \ + edgartools==5.20.2 \ + sec-parser \ + pandas \ + pyarrow \ + tqdm \ + httpx \ + tzdata \ + requests \ + beautifulsoup4 + +# Fail-fast smoke check so the image won't ship missing a needed dep. +RUN python -c "from edgar import set_identity; \ + import sec_parser, pandas, pyarrow, httpx, tzdata, requests; \ + from bs4 import BeautifulSoup; \ + from datasets import load_dataset; \ + print('SEC-prep + create_seed_data deps OK')" diff --git a/dockerfiles/Dockerfile.vllm b/dockerfiles/Dockerfile.vllm new file mode 100644 index 0000000..46bb563 --- /dev/null +++ b/dockerfiles/Dockerfile.vllm @@ -0,0 +1,56 @@ +# ============================================================================= +# NVFlow vLLM Container +# ============================================================================= +# Self-contained Dockerfile that builds the vLLM inference container with +# pre-cached tokenizer encodings. +# +# Build: +# docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:latest . +# +# Upstream source: +# https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.vllm +# ============================================================================= + + +# =========================================================================== +# BEGIN UPSTREAM (NeMo-Skills Dockerfile.vllm) +# =========================================================================== + +ARG VLLM_VERSION=v0.18.1 +FROM vllm/vllm-openai:${VLLM_VERSION} + +RUN pip install "vllm[audio]" +RUN pip install qwen-vl-utils +RUN pip install ray + +# =========================================================================== +# END UPSTREAM +# =========================================================================== + + +# =========================================================================== +# NVFlow Additional Layers +# =========================================================================== +# Pre-cache tokenizer encodings so no downloads are needed at runtime. +# =========================================================================== + +ENV TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache +ENV TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache +ENV TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache +RUN mkdir -p /opt/tiktoken_cache + +# Download tiktoken encoding files explicitly with curl. +# The Rust tiktoken-rs client inside openai_harmony fails to download under +# QEMU arm64 emulation (docker buildx), so we fetch them reliably here and +# point TIKTOKEN_ENCODINGS_BASE at the directory. This also makes the image +# fully air-gapped on both amd64 and arm64. +RUN curl -fSL -o /opt/tiktoken_cache/o200k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken && \ + curl -fSL -o /opt/tiktoken_cache/cl100k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken + +# Verify the harmony encoding loads from the pre-downloaded files +RUN python3 -c "\ +from openai_harmony import load_harmony_encoding, HarmonyEncodingName; \ +load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS); \ +print('openai_harmony encoding loaded OK')" diff --git a/dockerfiles/Dockerfile.vllm-grpo b/dockerfiles/Dockerfile.vllm-grpo new file mode 100644 index 0000000..3ade4f7 --- /dev/null +++ b/dockerfiles/Dockerfile.vllm-grpo @@ -0,0 +1,55 @@ +# ============================================================================= +# NVFlow vLLM-GRPO Container +# ============================================================================= +# vLLM container pinned to v0.17.1 for GRPO rollouts and judge inference. +# The main vLLM container (Dockerfile.vllm) uses v0.18.1 for SDG/eval. +# +# Build: +# docker build -f dockerfiles/Dockerfile.vllm-grpo -t nvflow-vllm-grpo:latest . +# +# Upstream source: +# https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.vllm +# ============================================================================= + + +# =========================================================================== +# BEGIN UPSTREAM (NeMo-Skills Dockerfile.vllm) +# =========================================================================== + +ARG VLLM_VERSION=v0.17.1 +FROM vllm/vllm-openai:${VLLM_VERSION} + +RUN pip install "vllm[audio]" +RUN pip install qwen-vl-utils + +# =========================================================================== +# END UPSTREAM +# =========================================================================== + + +# =========================================================================== +# NVFlow Additional Layers +# =========================================================================== +# Pre-cache tokenizer encodings so no downloads are needed at runtime. +# =========================================================================== + +ENV TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache +ENV TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache +ENV TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache +RUN mkdir -p /opt/tiktoken_cache + +# Download tiktoken encoding files explicitly with curl. +# The Rust tiktoken-rs client inside openai_harmony fails to download under +# QEMU arm64 emulation (docker buildx), so we fetch them reliably here and +# point TIKTOKEN_ENCODINGS_BASE at the directory. This also makes the image +# fully air-gapped on both amd64 and arm64. +RUN curl -fSL -o /opt/tiktoken_cache/o200k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken && \ + curl -fSL -o /opt/tiktoken_cache/cl100k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken + +# Verify the harmony encoding loads from the pre-downloaded files +RUN python3 -c "\ +from openai_harmony import load_harmony_encoding, HarmonyEncodingName; \ +load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS); \ +print('openai_harmony encoding loaded OK')" diff --git a/dockerfiles/README.md b/dockerfiles/README.md new file mode 100644 index 0000000..88abbfd --- /dev/null +++ b/dockerfiles/README.md @@ -0,0 +1,52 @@ +# NVFlow Container Images + +NVFlow uses five container images, all designed to run fully offline on +air-gapped Slurm clusters. Four are **built locally** from the +self-contained Dockerfiles in this directory; the fifth (`sglang`) is pulled +as-is from Docker Hub. The Dockerfiles are build recipes โ€” running +`docker build` against each one on a connected host produces the actual +images. + +For complete documentation โ€” build instructions, sanity checks, deployment +steps, air-gapped design rationale, and rebuild guidance โ€” see +**[docker_instructions.md](docker_instructions.md)**. + +## Quick Start + +```bash +# Requires `docker login nvcr.io` for the NGC registry (nemo-rl base image) +docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:v0.6.0 . +docker build -f dockerfiles/Dockerfile.nemo-skills -t nvflow-nemo-skills:0229040 . +docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:v0.18.1 . +docker build -f dockerfiles/Dockerfile.vllm-grpo -t nvflow-vllm-grpo:v0.17.1 . + +# Multi-arch builds (amd64 + arm64) โ€” push directly to a registry +REGISTRY= +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.vllm -t $REGISTRY/nvflow-vllm:v0.18.1 --push . +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.vllm-grpo -t $REGISTRY/nvflow-vllm-grpo:v0.17.1 --push . + +# sglang โ€” pull directly, no custom Dockerfile needed +docker pull lmsysorg/sglang:v0.5.10.post1 +``` + +## Images + +| Image | Base | Purpose | +|-------|------|---------| +| `nvflow-nemo-rl` | `nvcr.io/nvidia/nemo-rl:v0.6.0` | SFT, GRPO training, collect_rollouts, compute_rewards | +| `nvflow-nemo-skills` | `ubuntu:22.04` | SDG pipeline, evaluation, data preparation | +| `nvflow-vllm` | `vllm/vllm-openai:v0.18.1` | Standalone vLLM inference (SDG, eval) โ€” multi-arch (amd64 + arm64) | +| `nvflow-vllm-grpo` | `vllm/vllm-openai:v0.17.1` | Standalone vLLM inference (GRPO rollouts, judge) โ€” multi-arch | +| `sglang` | `lmsysorg/sglang:v0.5.10.post1` | SGLang inference server (pulled as-is) | + +## Version Pins + +| Build Arg | Default | Where to find the right value | +|-----------|---------|-------------------------------| +| `BASE_IMAGE` (nemo-rl) | `nvcr.io/nvidia/nemo-rl:v0.6.0` | [NGC NeMo-RL tags](https://catalog.ngc.nvidia.com) | +| `NEMO_SKILLS_COMMIT` | `022904023ad7a83a87662a313cf72e7df5891d55` (`0229040`) | Should match across `Dockerfile.nemo-skills` and `Dockerfile.nemo-rl` | +| `NEMO_GYM_BRANCH` | `ude/finance-sec-search-v2` | NeMo-Gym branch with finance agent | +| `VLLM_VERSION` (vllm) | `v0.18.1` | [vLLM releases](https://github.com/vllm-project/vllm/releases) | +| `VLLM_VERSION` (vllm-grpo) | `v0.17.1` | Pinned to match NeMo-RL v0.6.0 colocated vLLM | diff --git a/dockerfiles/docker_instructions.md b/dockerfiles/docker_instructions.md new file mode 100644 index 0000000..fc2a989 --- /dev/null +++ b/dockerfiles/docker_instructions.md @@ -0,0 +1,392 @@ +# NVFlow Air-Gapped Docker Images + +Build, validate, and deploy the five NVFlow container images for use on +air-gapped Slurm clusters. Four of them are produced by running +`docker build` against the self-contained Dockerfiles in this directory; +the fifth (`sglang`) is pulled as-is from Docker Hub. All images are built +on a connected host (the only step that needs internet) and then run fully +offline on the cluster. + +## Images + +| Image | Base | Purpose | +|---|---|---| +| `nvflow-nemo-rl` | `nvcr.io/nvidia/nemo-rl:v0.6.0` | SFT, GRPO training, collect_rollouts, compute_rewards | +| `nvflow-nemo-skills` | `ubuntu:22.04` | SDG pipeline, evaluation, data preparation, SEC data prep | +| `nvflow-vllm` | `vllm/vllm-openai:v0.18.1` | Standalone vLLM (SDG, eval) โ€” multi-arch (amd64 + arm64) | +| `nvflow-vllm-grpo` | `vllm/vllm-openai:v0.17.1` | Standalone vLLM (GRPO rollouts, judge) โ€” multi-arch | +| `sglang` | `lmsysorg/sglang:v0.5.10.post1` | SGLang inference server (pulled as-is, no custom Dockerfile) | + +## Version pins + +| Build arg | Default | Where to find the right value | +|---|---|---| +| `BASE_IMAGE` (nemo-rl) | `nvcr.io/nvidia/nemo-rl:v0.6.0` | [NGC NeMo-RL tags](https://catalog.ngc.nvidia.com) | +| `NEMO_SKILLS_COMMIT` | `022904023ad7a83a87662a313cf72e7df5891d55` (`0229040`) | Must match across `Dockerfile.nemo-skills` and `Dockerfile.nemo-rl` | +| `NEMO_GYM_BRANCH` | `ude/finance-sec-search-v2` | NeMo-Gym branch with finance agent | +| `VLLM_VERSION` (vllm) | `v0.18.1` | [vLLM releases](https://github.com/vllm-project/vllm/releases) | +| `VLLM_VERSION` (vllm-grpo) | `v0.17.1` | Pinned to match NeMo-RL v0.6.0 colocated vLLM | + +## 1. Build + +**Build host requirements:** any OS with Docker Engine or Docker Desktop and +internet access (Linux, macOS, Windows/WSL2 all work). The destination Slurm +cluster is the constraint โ€” it's almost always `linux/amd64`, so all examples +below produce amd64 images. + +> **Default platform = host architecture.** `docker build` produces an image +> for the build host's arch. On amd64 Linux / Intel macOS / Windows that's +> `linux/amd64`. On Apple Silicon, Graviton, or other arm64 hosts it's +> `linux/arm64` โ€” to get amd64 from those hosts, add `--platform linux/amd64` +> via `buildx` (see arm64 section below). + +### amd64 (default on amd64 build hosts) + +```bash +cd /path/to/nvflow + +docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:v0.6.0 . +docker build -f dockerfiles/Dockerfile.nemo-skills -t nvflow-nemo-skills:0229040 . +docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:v0.18.1 . +docker build -f dockerfiles/Dockerfile.vllm-grpo -t nvflow-vllm-grpo:v0.17.1 . + +# sglang โ€” pulled directly, no custom Dockerfile +docker pull lmsysorg/sglang:v0.5.10.post1 +``` + +### Cross-arch builds (e.g. amd64 image on Apple Silicon) + +Cross-architecture builds need `buildx` plus QEMU emulation registered on +the host. Register QEMU once per build host (Linux only โ€” Docker Desktop +ships QEMU pre-registered): + +```bash +docker run --privileged --rm tonistiigi/binfmt --install all +``` + +Then build with an explicit `--platform`: + +```bash +# amd64 image from an arm64 host (most common cross-arch case for Slurm) +docker buildx build --platform linux/amd64 \ + -f dockerfiles/Dockerfile.vllm \ + -t nvflow-vllm:v0.18.1 \ + --load . +``` + +QEMU-emulated builds are significantly slower than native and can take +several hours, especially for `Dockerfile.nemo-rl`. Prefer a native build +host when possible. + +### linux/arm64 single-arch build + +Only the two vLLM images are arm64-friendly today. From an arm64 host the +plain `docker build` works; from an amd64 host, use `buildx` with QEMU: + +```bash +docker buildx build --platform linux/arm64 \ + -f dockerfiles/Dockerfile.vllm \ + -t nvflow-vllm:v0.18.1-arm64 \ + --load . +``` + +`--load` only supports a single platform at a time; for multi-arch see below. + +### Multi-arch build (amd64 + arm64) โ€” push to registry + +For `vllm` and `vllm-grpo`, build for both architectures and push the manifest +list in one shot. Multi-arch builds **must** push to a registry โ€” the local +Docker image store can't hold a manifest list, so `--load` is not an option: + +```bash +REGISTRY= + +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.vllm \ + -t $REGISTRY/nvflow-vllm:v0.18.1 \ + --provenance=false --sbom=false --push . + +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.vllm-grpo \ + -t $REGISTRY/nvflow-vllm-grpo:v0.17.1 \ + --provenance=false --sbom=false --push . + +# Verify both architectures are in the manifest list +docker buildx imagetools inspect $REGISTRY/nvflow-vllm:v0.18.1 +``` + +`--provenance=false --sbom=false` keeps the manifest list compatible with +older registries / consumers that don't understand attestation manifests. + +## 2. Sanity checks (blockers) + +Run these against the locally-built Docker images before the time-consuming +enroot import step. Each check below is a hard blocker โ€” if it fails, the +image will not work in production. + +> If your build host arch differs from the image arch (e.g. running checks +> against an amd64 image on Apple Silicon), the checks will run under QEMU +> emulation as long as QEMU is registered (Docker Desktop ships it; on Linux +> see the `tonistiigi/binfmt` step above). Without QEMU, `docker run` will +> fail with `exec format error` โ€” defer the checks to the cluster after +> enroot import in that case. + +### nemo-rl + +```bash +IMAGE=nvflow-nemo-rl:v0.6.0 + +# A. uv works offline (paths relocated out of /root) +docker run --rm -e UV_OFFLINE=true $IMAGE bash -c \ + "uv python list --only-installed | grep 3.12" +# Expect: cpython-3.12.x at /opt/uv-python/... + +# B. main venv has no stale /root/.local references +docker run --rm $IMAGE bash -c ' + grep -rl "/root/.local" \ + /opt/nemo_rl_venv/pyvenv.cfg \ + /opt/ray_venvs/*/pyvenv.cfg \ + /opt/nemo-rl/3rdparty/Gym-workspace/Gym/.venv/pyvenv.cfg \ + 2>/dev/null || echo "All clean"' +# Expect: All clean + +# C. all 6 Gym component venvs are symlinked +docker run --rm $IMAGE bash -c ' + GYM=/opt/nemo-rl/3rdparty/Gym-workspace/Gym + for c in \ + resources_servers/equivalence_llm_judge \ + resources_servers/finance_sec_search \ + responses_api_agents/simple_agent \ + responses_api_agents/finance_agent \ + responses_api_models/openai_model \ + responses_api_models/vllm_model; do + [ -L "$GYM/$c/.venv" ] && echo "OK: $c" || echo "MISSING: $c" + done' +# Expect: 6x "OK: ..." + +# D. uvicorn pin (timeout_worker_healthcheck kwarg required by Gym servers) +docker run --rm $IMAGE bash -c ' + /opt/nemo-rl/3rdparty/Gym-workspace/Gym/.venv/bin/python -c " +import uvicorn, inspect +assert \"timeout_worker_healthcheck\" in inspect.signature(uvicorn.run).parameters, uvicorn.__version__ +print(\"uvicorn\", uvicorn.__version__, \"OK\")"' +# Expect: uvicorn 0.37.x OK +``` + +### nemo-skills + +```bash +IMAGE=nvflow-nemo-skills:0229040 + +# A. tiktoken pre-cache loads offline +docker run --rm --network=none -e HF_HUB_OFFLINE=1 $IMAGE bash -c ' + python3 -c "import tiktoken; tiktoken.get_encoding(\"cl100k_base\"); print(\"OK\")"' +# Expect: OK + +# B. tzdata populated (otherwise pyarrow timestamps blow up at job time: +# ArrowInvalid: timezone "UTC" not found in timezone database) +docker run --rm $IMAGE bash -c ' + python3 -c "import pyarrow as pa; pa.array([], type=pa.timestamp(\"ns\", tz=\"UTC\")); print(\"OK\")"' +# Expect: OK + +# C. SDG + SEC data-prep deps importable +# (used by workflow-2 download_sec and workflow-3 step-0 create_seed_data; +# all baked into Dockerfile.nemo-skills, no runtime pip install) +docker run --rm $IMAGE bash -c ' + python3 -c "import jsonlines, tiktoken, markdownify, backoff, func_timeout, tavily, edgartools, sec_parser, pandas, pyarrow, requests, bs4; print(\"OK\")"' +# Expect: OK +``` + +### vllm and vllm-grpo + +Run the same set against both images: + +```bash +for IMAGE in nvflow-vllm:v0.18.1 nvflow-vllm-grpo:v0.17.1; do + echo "=== $IMAGE ===" + + # A. tiktoken encoding files present + docker run --rm $IMAGE bash -c "ls /opt/tiktoken_cache/*.tiktoken" + # Expect: o200k_base.tiktoken and cl100k_base.tiktoken + + # B. openai_harmony loads offline (no network) + docker run --rm --network=none $IMAGE bash -c ' + python3 -c "from openai_harmony import load_harmony_encoding, HarmonyEncodingName; \ + load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS); print(\"OK\")"' + # Expect: OK +done +``` + +## 3. Convert to `.sqsh` for the Slurm cluster + +`enroot` runs on the Slurm compute/login nodes (Linux only). There are two +paths from a Docker image to a `.sqsh` file โ€” pick whichever fits your air-gap +workflow. + +### Option A: via a private container registry (recommended) + +Push each image, then `enroot import` from the registry on the cluster: + +```bash +REGISTRY= + +docker tag nvflow-nemo-rl:v0.6.0 $REGISTRY/nvflow-nemo-rl:v0.6.0 +docker tag nvflow-nemo-skills:0229040 $REGISTRY/nvflow-nemo-skills:0229040 +docker tag nvflow-vllm:v0.18.1 $REGISTRY/nvflow-vllm:v0.18.1 +docker tag nvflow-vllm-grpo:v0.17.1 $REGISTRY/nvflow-vllm-grpo:v0.17.1 + +docker push $REGISTRY/nvflow-nemo-rl:v0.6.0 +docker push $REGISTRY/nvflow-nemo-skills:0229040 +docker push $REGISTRY/nvflow-vllm:v0.18.1 +docker push $REGISTRY/nvflow-vllm-grpo:v0.17.1 +``` + +Then on the cluster (typically a CPU partition): + +```bash +CONTAINER_DIR= +REGISTRY= + +enroot import \ + --output $CONTAINER_DIR/nvflow-nemo-rl-v0.6.0.sqsh \ + "docker://$REGISTRY/nvflow-nemo-rl:v0.6.0" + +# Repeat for nemo-skills, vllm, vllm-grpo, and (optionally) sglang. +``` + +If the cluster authenticates to your registry, drop credentials into +`~/.config/enroot/.credentials`: + +``` +machine login password +``` + +### Option B: via a saved tarball (no registry required) + +For fully air-gapped sites without a private registry, save the Docker image +to a tarball, transfer it to a Linux host that has both Docker and `enroot`, +load the tarball into the local Docker daemon, then import via `dockerd://`: + +```bash +# On the build host +docker save nvflow-nemo-rl:v0.6.0 | gzip > nvflow-nemo-rl-v0.6.0.tar.gz +# Transfer the .tar.gz to the cluster (scp / rsync / sneakernet) + +# On the cluster (requires a Docker daemon accessible to your user) +gunzip -c nvflow-nemo-rl-v0.6.0.tar.gz | docker load +enroot import \ + --output $CONTAINER_DIR/nvflow-nemo-rl-v0.6.0.sqsh \ + dockerd://nvflow-nemo-rl:v0.6.0 +``` + +> `enroot import` natively supports only `docker://` (remote registry), +> `dockerd://` (local Docker daemon), and `podman://` URIs. If the cluster +> has neither a private registry nor a Docker daemon, run a transient local +> registry container, push to it, and import via `docker://localhost:5000/...`. + +### Two things to watch for in either option + +- **Registries with a path component need `#` instead of `/`.** `enroot` parses + `docker:///` such that everything after the first `/` is image + path, which breaks for registries where the host itself contains a path + (e.g. `nvcr.io/`). Use `#` to separate host from image path: + ```bash + enroot import --output nvflow-vllm-v0.18.1.sqsh \ + "docker://nvcr.io#/nvflow-vllm:v0.18.1" + ``` +- **Filename colon.** `enroot` writes the Docker tag separator (`:`) literally + into the output filename. Either pass `--output` with a shell-safe name (as + above) or rename after import: + ```bash + mv "nvflow-nemo-rl:v0.6.0.sqsh" nvflow-nemo-rl-v0.6.0.sqsh + ``` + +## 4. Cluster config (`my_cluster.yaml`) + +`cluster_configs/my_cluster.yaml` is **not tracked in git** โ€” create it from +the template below, replacing the `<...>` placeholders with site-specific +values. The two blocks below are the minimum required for air-gapped +operation. + +### Container paths (point at the `.sqsh` files from step 3) + +```yaml +containers: + nemo-rl: /nvflow-nemo-rl-v0.6.0.sqsh + nemo-skills: /nvflow-nemo-skills-0229040.sqsh + vllm: /nvflow-vllm-v0.18.1.sqsh + vllm-grpo: /nvflow-vllm-grpo-v0.17.1.sqsh + # sglang: /sglang-v0.5.10.post1.sqsh +``` + +### Air-gap enforcement + +These environment variables turn off all outbound package and model fetches: + +```yaml +env_vars: + # Disable HuggingFace network access (Hub, datasets, transformers). + - HF_HUB_OFFLINE=1 + - HF_DATASETS_OFFLINE=1 + - TRANSFORMERS_OFFLINE=1 + + # Disable uv package and Python interpreter downloads. + - UV_OFFLINE=true + + # Point tiktoken / openai_harmony at the cache baked into the images. + # Required for nemo-skills and nemo-rl (vllm/vllm-grpo set them as ENV). + - TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache +``` + +### Don't bind-mount NeMo-RL or NeMo-Gym source over the image paths + +The air-gapped `nvflow-nemo-rl` image already contains NeMo-Gym venv +at `/opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv` (sanity check **C** in +section 2 verifies this). GRPO stages source that venv via +`installation_command: source .../Gym/.venv/bin/activate` before running. + +Older dev-mode `my_cluster.yaml` templates often include host source overlays +like: + +```yaml +mounts: + # DO NOT use these with the air-gapped image โ€” they shadow the baked .venv + # - /RL:/opt/NeMo-RL + # - /Gym:/opt/NeMo-RL/3rdparty/Gym-workspace/Gym +``` + +These bind-mounts hide the baked `.venv` symlink and the `installation_command` +fails with `No such file or directory` โ€” breaking `prepare_data`, +`collect_rollouts`, `compute_rewards`, and `training` for GRPO. Only add +these mounts if you are deliberately iterating on NeMo-RL/Gym source against a +host `.venv` you've built to be ABI-compatible with the image. + +### Don't enable this in offline mode + +```yaml +# - NRL_FORCE_REBUILD_VENVS=true # forces Ray workers to re-resolve via uv + # (requires internet; will fail under air-gap) +``` + +## Notes for one-time / connected-node operations + +A few stages legitimately need internet on first run. Run them on a +connected node (or off-cluster) and ship the resulting artifacts onto the +air-gapped cluster: + +| Stage | Why it needs internet | +|---|---| +| `workflow-2 download_sec_filings` | Downloads filings from SEC EDGAR (not HF, but still external). | +| `workflow-3 step-0 create_seed_data` | Pulls `nogabenyoash/SecQue` from HuggingFace. | +| `workflow-1 step-0 prepare_data` (eval) | Pulls finance benchmark datasets (`secque`, `financebench`) from HuggingFace. | +| `workflow-5 step-4 prepare_data` (GRPO) | Only if `should_download: true`; default `should_download: false` requires no internet. | + +For these stages, temporarily clear the three HF flags +(`HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE`, `TRANSFORMERS_OFFLINE`). Keep +`UV_OFFLINE=true` set โ€” `uv` should never need to resolve packages at runtime. + +Note: `huggingface_hub` interprets `TRANSFORMERS_OFFLINE=1` as +`HF_HUB_OFFLINE=1`, so all three need to be off (or unset) for HF dataset +pulls to succeed. diff --git a/docs/architecture/ARCHITECTURE.md b/docs/architecture/ARCHITECTURE.md index 1d6c5e0..1d58a20 100644 --- a/docs/architecture/ARCHITECTURE.md +++ b/docs/architecture/ARCHITECTURE.md @@ -668,25 +668,29 @@ graph TB ### Container Architecture ``` -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ Container Images (.sqsh format) โ”‚ -โ”œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ค -โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ nemo-skills โ”‚ โ”‚ vLLM โ”‚ โ”‚ -โ”‚ โ”‚ (0229040) โ”‚ โ”‚ (v0.18.1) โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ SGLang โ”‚ โ”‚ NeMo-RL โ”‚ โ”‚ -โ”‚ โ”‚ (v0.5.10) โ”‚ โ”‚ (v0.6.0) โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ NeMo FW โ”‚ โ”‚ PyTorch โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” +โ”‚ Air-gapped container images (.sqsh format) โ”‚ +โ”‚ (built locally from dockerfiles/Dockerfile.* โ€” see INSTALL.md)โ”‚ +โ”œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ค +โ”‚ โ”‚ +โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ +โ”‚ โ”‚ nvflow-nemo-skills โ”‚ โ”‚ nvflow-vllm โ”‚ โ”‚ +โ”‚ โ”‚ (0229040) โ”‚ โ”‚ (v0.18.1) โ”‚ โ”‚ +โ”‚ โ”‚ SDG/eval/data prep โ”‚ โ”‚ SDG/eval inference โ”‚ โ”‚ +โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ +โ”‚ โ”‚ +โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ +โ”‚ โ”‚ sglang (pulled) โ”‚ โ”‚ nvflow-nemo-rl โ”‚ โ”‚ +โ”‚ โ”‚ (v0.5.10.post1) โ”‚ โ”‚ (v0.6.0) โ”‚ โ”‚ +โ”‚ โ”‚ SDG inference โ”‚ โ”‚ SFT/GRPO + Gym venv โ”‚ โ”‚ +โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ +โ”‚ โ”‚ +โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ +โ”‚ โ”‚ nvflow-vllm-grpo โ”‚ โ”‚ +โ”‚ โ”‚ (v0.17.1) โ”‚ โ”‚ +โ”‚ โ”‚ GRPO rollout/judge โ”‚ โ”‚ +โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ +โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ†“ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ Shared Filesystem Mounts โ”‚ diff --git a/docs/cluster-configuration.md b/docs/cluster-configuration.md index 4228950..5ca331b 100644 --- a/docs/cluster-configuration.md +++ b/docs/cluster-configuration.md @@ -253,22 +253,33 @@ mounts: | Root Lustre | `/lustre` | Access entire shared filesystem | | Dataset directory | `/data` | Training/evaluation datasets | +### Do NOT bind-mount NeMo-RL / NeMo-Gym source over the image paths + +The self-sufficient `nvflow-nemo-rl` image (built from [`dockerfiles/Dockerfile.nemo-rl`](../dockerfiles/Dockerfile.nemo-rl)) already contains: + +- NeMo-RL source at `/opt/NeMo-RL` (and `/opt/nemo-rl` lowercase alias) +- NeMo-Gym at `/opt/NeMo-RL/3rdparty/Gym-workspace/Gym` (branch `ude/finance-sec-search-v2`) +- A pre-built `.venv` symlinked across all 6 Gym components + +GRPO stages call `installation_command: source /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate`. Bind-mounting a host source tree at `/opt/NeMo-RL` or `/opt/NeMo-RL/3rdparty/Gym-workspace/Gym` **shadows the baked `.venv`** and breaks `prepare_data`, `collect_rollouts`, `compute_rewards`, and `training` with `No such file or directory`. + +The overlay mounts in `template-slurm.yaml` are commented out for exactly this reason. Only uncomment them if you're deliberately iterating on NeMo-RL / Gym source against a host `.venv` you've built to be ABI-compatible with the image. In that dev-mode case you must also set `NRL_FORCE_REBUILD_VENVS=true` (see [Environment Variables](#environment-variables) below) -- which requires internet, so it can only be used on a connected node. + ### Model-Specific Cluster Configs -Some models require additional mounts not needed by others. Rather than cluttering a single config with conditional mounts, use separate cluster config files: +Some models require additional cluster-level differences (e.g. different timeouts, partitions, or env vars). Rather than cluttering a single config with conditional logic, use separate cluster config files: -| Cluster Config | Used By | Extra Mounts | Notes | -|----------------|---------|-------------|-------| -| `my_cluster.yaml` | Qwen3, Gemma3 (dense models) | None | Default for all standard models | -| `my_cluster_nemotron.yaml` | Nemotron-3-Nano (MoE) | `/path/to/RL:/opt/NeMo-RL` | NeMo-RL overlay for MoE support | +| Cluster Config | Used By | Notes | +|----------------|---------|-------| +| `my_cluster.yaml` | Qwen3, Gemma3 (dense models) | Default for all standard models | +| `my_cluster_nemotron.yaml` | Nemotron-3-Nano (MoE) | Use only if Nemotron needs different mounts/env -- the self-sufficient `nvflow-nemo-rl` image now handles MoE without a host overlay | **How it works:** - `base.yaml` (SFT workflow) sets `cluster: my_cluster` as the default -- `nemotron-3-nano.yaml` overrides with `cluster: my_cluster_nemotron` -- Both configs are identical except for the NeMo-RL overlay mount +- A model config can override with `cluster: my_cluster_nemotron` - Keep both configs in sync when making infrastructure changes -> **Tip:** The NeMo-RL overlay mount is temporary. Once MoE support is merged into the main NeMo-RL branch, `my_cluster_nemotron.yaml` can be retired. See the [SFT Workflow Guide](recipes/finance/workflows/04-sft.md#nemotron-3-nano-special-requirements) for details. +> **Note:** Previous versions of this guide recommended a NeMo-RL host overlay (`/path/to/RL:/opt/NeMo-RL`) for Nemotron-3-Nano MoE support. With the self-sufficient `nvflow-nemo-rl` image that overlay is no longer required and would shadow the baked `.venv`. See the [SFT Workflow Guide](recipes/finance/workflows/04-sft.md) for the current setup. --- @@ -330,7 +341,15 @@ env_vars: - TOKENIZERS_PARALLELISM=false - VIRTUAL_ENV= - VIRTUAL_ENV_PROMPT= - # NeMo-RL / GRPO (uncomment as needed) + # --- Air-gap enforcement (recommended; on by default in template-slurm.yaml) --- + - HF_HUB_OFFLINE=1 + - HF_DATASETS_OFFLINE=1 + - TRANSFORMERS_OFFLINE=1 + - UV_OFFLINE=true + - TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache + # NeMo-RL / GRPO dev-mode only (do NOT enable in self-sufficient mode) # - NRL_FORCE_REBUILD_VENVS=true # API keys (keep secret, don't commit to git!) - HF_TOKEN=hf_... @@ -349,11 +368,27 @@ env_vars: | `VIRTUAL_ENV` | *(empty)* | Unset to prevent host virtualenv from leaking into containers | | `VIRTUAL_ENV_PROMPT` | *(empty)* | Unset to prevent host venv prompt from leaking into containers | -#### NeMo-RL / GRPO Variables +#### Air-Gap Enforcement Variables + +These variables prevent the runtime from making outbound network calls and from missing baked-in tokenizer encodings. `template-slurm.yaml` ships them pre-populated; leave them set in normal production. + +| Variable | Value | Purpose | +|----------|-------|---------| +| `HF_HUB_OFFLINE` | `1` | Disables HuggingFace Hub network access (model + tokenizer downloads) | +| `HF_DATASETS_OFFLINE` | `1` | Disables `datasets` network access | +| `TRANSFORMERS_OFFLINE` | `1` | Disables `transformers` network access. `huggingface_hub` treats this as equivalent to `HF_HUB_OFFLINE=1` | +| `UV_OFFLINE` | `true` | Prevents `uv` from resolving / downloading packages or Python interpreters at runtime. Keep this set **always** -- containers ship with frozen venvs | +| `TIKTOKEN_CACHE_DIR` | `/opt/tiktoken_cache` | Points `tiktoken` at the cache baked into the images | +| `TIKTOKEN_RS_CACHE_DIR` | `/opt/tiktoken_cache` | Points the Rust `tiktoken-rs` client at the cache (used by `openai_harmony`) | +| `TIKTOKEN_ENCODINGS_BASE` | `/opt/tiktoken_cache` | Required for `openai_harmony` to load `HARMONY_GPT_OSS` offline | + +> **One-time connected-node stages:** A few stages (`download_sec_filings`, `create_seed_data`, eval `prepare_data`, GRPO `prepare_data` with `should_download: true`) need internet on first run to pull benchmark/seed datasets. For those submissions, **temporarily comment out** `HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE`, and `TRANSFORMERS_OFFLINE`. Keep `UV_OFFLINE=true` set in all cases. See [INSTALL.md โ†’ One-Time Connected-Node Stages](../INSTALL.md#one-time-connected-node-stages-datasets). + +#### NeMo-RL / GRPO Variables (Dev Mode Only) | Variable | Value | Purpose | |----------|-------|---------| -| `NRL_FORCE_REBUILD_VENVS` | `true` | Forces Ray workers to rebuild their virtual environments from the mounted NeMo-RL source tree instead of reusing cached venvs. **Enable this** when you update the NeMo-RL or Gym overlay mount (new branch, new commit). Without it, workers may use stale cached venvs with outdated code, causing import errors or silent behavior differences. Safe to leave enabled โ€” only triggers a rebuild when the source tree actually changes | +| `NRL_FORCE_REBUILD_VENVS` | `true` | **Dev mode only.** Forces Ray workers to rebuild their virtual environments from the mounted NeMo-RL source tree instead of reusing cached venvs. Requires internet (uses `uv` to resolve packages) -- **do not enable in self-sufficient production**. Only relevant when you've bind-mounted a host NeMo-RL / Gym source clone over `/opt/NeMo-RL` and want Ray workers to pick up the new source | #### API Keys (Secrets) diff --git a/docs/recipes/finance/troubleshooting.md b/docs/recipes/finance/troubleshooting.md index 5b66396..3e53e9e 100644 --- a/docs/recipes/finance/troubleshooting.md +++ b/docs/recipes/finance/troubleshooting.md @@ -5,6 +5,7 @@ Comprehensive troubleshooting guide for common issues across all finance recipe ## Quick Navigation - [Cluster & Infrastructure](#cluster--infrastructure) +- [Offline Runtime](#self-sufficient-runtime) - [Resource Issues](#resource-issues) - [Data Issues](#data-issues) - [Training Issues](#training-issues) @@ -98,6 +99,88 @@ scontrol show config | grep SLURM_VERSION --- +## Offline Runtime + +The default NVFlow images (`nvflow-nemo-rl`, `nvflow-nemo-skills`, `nvflow-vllm`, `nvflow-vllm-grpo`) are built to run with **no outbound network access** at job time. Most "weird" runtime errors on a freshly-deployed cluster trace back to a missing offline asset, a stale overlay mount, or an env var that was cleared. + +For the full build / deploy / verify flow, see [INSTALL.md](../../../INSTALL.md) and [`dockerfiles/docker_instructions.md`](../../../dockerfiles/docker_instructions.md). + +### GRPO `installation_command` fails with `No such file or directory` + +**Problem:** A GRPO stage (`prepare_data`, `collect_rollouts`, `compute_rewards`, or `training`) fails immediately after `source /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate` with: + +``` +bash: /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate: No such file or directory +``` + +**Cause:** You bind-mounted a host clone of NeMo-RL or NeMo-Gym at `/opt/NeMo-RL` (or `/opt/NeMo-RL/3rdparty/Gym-workspace/Gym`), which shadows the baked `.venv` inside the `nvflow-nemo-rl` image. + +**Solution:** Remove the overlay mounts from `cluster_configs/my_cluster.yaml`. The self-sufficient image already contains everything GRPO needs: + +```yaml +mounts: + # COMMENT THESE OUT (or delete) for normal production runs: + # - :/opt/NeMo-RL + # - :/opt/NeMo-RL/3rdparty/Gym-workspace/Gym +``` + +See [INSTALL.md โ†’ Setup NeMo-RL & NeMo-Gym Sources](../../../INSTALL.md#setup-nemo-rl--nemo-gym-sources-for-grpo) for when (rarely) the overlay is correct. + +### `huggingface_hub.errors.OfflineModeIsEnabled` / `LocalEntryNotFoundError` + +**Problem:** A stage fails trying to pull a model or dataset from HuggingFace Hub. + +**Cause:** Air-gap mode is on (`HF_HUB_OFFLINE=1`, etc.) but the asset isn't pre-staged on disk. + +**Solution:** +- **Models:** Pre-download to your mounted `hf_models` directory with `hf download` -- see [INSTALL.md โ†’ Download Models](../../../INSTALL.md#download-models). +- **Datasets / SEC filings:** Some stages (`download_sec_filings`, `create_seed_data`, eval `prepare_data`, GRPO `prepare_data` with `should_download: true`) need internet on first run. Run them on a connected node with the three `HF_*_OFFLINE` flags **temporarily commented out** in `my_cluster.yaml`; keep `UV_OFFLINE=true` set. The artifacts persist under `/workspace` and are reused by every subsequent run. + +### `uv` errors with "package not installed" or tries to resolve from PyPI + +**Problem:** A Ray worker or stage script fails because `uv` is trying to download a package. + +**Cause (usual):** Someone enabled `NRL_FORCE_REBUILD_VENVS=true` in offline mode. That flag forces Ray workers to re-resolve packages via `uv`, which requires internet. + +**Solution:** Comment out `NRL_FORCE_REBUILD_VENVS` in `my_cluster.yaml`. It's only safe to enable on a connected node when you've bind-mounted a host NeMo-RL source overlay and changed the source tree -- see [`docs/cluster-configuration.md`](../../cluster-configuration.md#nemo-rl--grpo-variables-dev-mode-only). + +**Cause (rare):** A baked venv is genuinely missing a dependency. Rebuild the image with the missing package added to the Dockerfile and re-run the sanity checks from [`dockerfiles/docker_instructions.md` ยง2](../../../dockerfiles/docker_instructions.md#2-sanity-checks-blockers). + +### `tiktoken` / `openai_harmony` fails to load offline + +**Problem:** An SDG, eval, or GRPO stage crashes with `tiktoken` / `openai_harmony` trying to fetch encoding files from `openaipublic.blob.core.windows.net`. + +**Cause:** The pre-cached encodings env vars are not set inside the container. `vllm` / `vllm-grpo` set them via image `ENV`, but `nemo-skills` and `nemo-rl` rely on the cluster config. + +**Solution:** Confirm `cluster_configs/my_cluster.yaml` has all three set: + +```yaml +env_vars: + - TIKTOKEN_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_RS_CACHE_DIR=/opt/tiktoken_cache + - TIKTOKEN_ENCODINGS_BASE=/opt/tiktoken_cache +``` + +Verify the cache exists inside the image: +```bash +docker run --rm nvflow-nemo-skills:0229040 ls /opt/tiktoken_cache +# Expect: cl100k_base.tiktoken (and o200k_base.tiktoken in vllm images) +``` + +### `pyarrow` / `pandas` blow up with `ArrowInvalid: timezone "UTC" not found` + +**Problem:** SEC data prep fails on first record with a timezone-related error. + +**Cause:** `tzdata` is not populated in the image. + +**Solution:** Already fixed in `Dockerfile.nemo-skills` (apt `tzdata`). If you see this in a custom-built image, confirm `tzdata` is installed: +```bash +docker run --rm nvflow-nemo-skills:0229040 bash -c \ + 'python3 -c "import pyarrow as pa; pa.array([], type=pa.timestamp(\"ns\", tz=\"UTC\")); print(\"OK\")"' +``` + +--- + ## Resource Issues ### Out of Memory (OOM) Errors diff --git a/nvflow/core/workflow_runner.py b/nvflow/core/workflow_runner.py index d88a547..cc6843a 100644 --- a/nvflow/core/workflow_runner.py +++ b/nvflow/core/workflow_runner.py @@ -244,6 +244,13 @@ def run( # Validate that requested stages exist in config self._validate_stages(stages_to_run, all_stages) + # Ensure cluster-level extra_sbatch_args reach every Slurm submission. + # Installed lazily here (not at CLI startup) because importing + # nemo_skills.pipeline pulls in torch/transformers (~15s cold cache). + from nvflow.lib.sbatch import apply_sbatch_args_autopatch + + apply_sbatch_args_autopatch() + # Warn about sibling stages that are declared in pipeline_stages # but not currently registered (e.g., their import failed). self._preflight_pipeline_health(all_stages, stages_to_run) diff --git a/nvflow/lib/rl/rollout.py b/nvflow/lib/rl/rollout.py index ce84a7f..4a3e21e 100644 --- a/nvflow/lib/rl/rollout.py +++ b/nvflow/lib/rl/rollout.py @@ -592,6 +592,7 @@ def _build_client_cmd( ' "+head_server.port=$HEAD_SERVER_PORT" \\\n' ' "+port_range_low=1024" \\\n' ' "+port_range_high=8999" \\\n' + ' "+skip_venv_if_present=true" \\\n' f"{judge_ng_run_overrides}" ' > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 &\n' "NG_RUN_PID=$!\n" @@ -1060,14 +1061,9 @@ def _build_collection_jobs( chunk_job_specs: dict[int, list[dict]] = {} job_log_dir = f"{p.output_dir}/logs" - sbatch_kwargs = None - raw_args = cluster_config.get("extra_sbatch_args") or [] - if raw_args: - sbatch_kwargs = {} - for arg in raw_args: - arg = arg.lstrip("-") - key, _, val = arg.partition("=") - sbatch_kwargs[key] = val if val else True + from nvflow.lib.sbatch import parse_extra_sbatch_args + + sbatch_kwargs = parse_extra_sbatch_args(cluster_config) or None for seed, chunk_id in remaining: job_lbl = f"rs{seed}_chunk{chunk_id}" diff --git a/nvflow/lib/rl/verify.py b/nvflow/lib/rl/verify.py index 0f3fc71..5d60d7c 100644 --- a/nvflow/lib/rl/verify.py +++ b/nvflow/lib/rl/verify.py @@ -144,6 +144,7 @@ def _build_verify_cmd( ' "+policy_model.responses_api_models.vllm_model.model=unused" \\\n' ' "+head_server.host=127.0.0.1" \\\n' ' "+head_server.port=$HEAD_SERVER_PORT" \\\n' + ' "+skip_venv_if_present=true" \\\n' f"{judge_ng_run_overrides}" ' > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 &\n' "NG_RUN_PID=$!\n" diff --git a/nvflow/lib/runtime.py b/nvflow/lib/runtime.py new file mode 100644 index 0000000..7dcb771 --- /dev/null +++ b/nvflow/lib/runtime.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Dual-mode Python interpreter resolution for Slurm jobs. + +NVFlow containers ship with pre-built venvs (airgap mode). When those +venvs are absent (dev mode with host-mounted NeMo-RL source), the code +falls back to ``uv run`` which resolves dependencies on the fly. + +All resolution happens **at bash runtime on the cluster node**, not at +Python submission time, because the venv paths exist inside the container. +""" + +from __future__ import annotations + +__all__ = [ + "NRL_PYTHON_PREAMBLE", + "ray_venv_python_preamble", +] + +# Bash preamble for SFT and GRPO training stages. +# Sets $NRL_PYTHON to the pre-built venv interpreter if available, +# otherwise falls back to uv with the NeMo-RL project context. +NRL_PYTHON_PREAMBLE = ( + "if [ -x /opt/nemo_rl_venv/bin/python ]; then " + " NRL_PYTHON=/opt/nemo_rl_venv/bin/python; " + "else " + " export UV_PROJECT=/opt/NeMo-RL; " + ' NRL_PYTHON="uv run --active python"; ' + "fi" +) + + +def ray_venv_python_preamble(venv_path: str, uv_extra: str) -> str: + """Bash preamble resolving a Ray worker venv Python. + + Used by checkpoint conversion scripts which need specific Ray venvs + (megatron, dtensor v1/v2) that are only present in the airgap container. + + Sets ``$CONVERT_PYTHON`` for use in the generated bash script. + """ + return ( + f"if [ -x {venv_path} ]; then " + f" CONVERT_PYTHON={venv_path}; " + f"else " + f" export UV_PROJECT=/opt/NeMo-RL; " + f' CONVERT_PYTHON="uv run --extra {uv_extra} python"; ' + f"fi" + ) diff --git a/nvflow/lib/sbatch.py b/nvflow/lib/sbatch.py new file mode 100644 index 0000000..61982c8 --- /dev/null +++ b/nvflow/lib/sbatch.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Plumb ``cluster_config['extra_sbatch_args']`` into every Slurm submission. + +Cluster configs may declare global Slurm flags via ``extra_sbatch_args`` +(e.g. ``--qos=xlarge_qos``, ``--exclusive``). These must reach the +generated sbatch script for every stage โ€” SDG, SFT, eval, and GRPO. + +This module patches ``nemo_skills.pipeline.utils.exp.get_executor``, the +single funnel through which all ``SlurmExecutor`` instances are created. +One wrapper, two rebindings (``exp`` module + ``utils`` package re-export), +complete coverage. + +:func:`parse_extra_sbatch_args` is also used by ``nvflow/lib/rl/rollout.py`` +which previously inlined the same parsing loop. +""" + +from __future__ import annotations + +import functools +from typing import Any + +__all__ = [ + "apply_sbatch_args_autopatch", + "parse_extra_sbatch_args", +] + + +def parse_extra_sbatch_args(cluster_config: dict | None) -> dict[str, Any]: + """Parse ``extra_sbatch_args`` list into a kwargs dict. + + Accepts entries like ``--qos=xlarge_qos`` (value) or ``--exclusive`` + (boolean flag). Leading dashes are stripped. Empty / non-string + entries are skipped so malformed YAML doesn't crash submission. + """ + if not cluster_config: + return {} + out: dict[str, Any] = {} + for arg in cluster_config.get("extra_sbatch_args") or []: + if not isinstance(arg, str): + continue + key, sep, val = arg.lstrip("-").partition("=") + if not key: + continue + out[key] = val if sep else True + return out + + +_PATCHED = False + + +def apply_sbatch_args_autopatch() -> None: + """Wrap ``get_executor`` so every Slurm submission honours ``extra_sbatch_args``. + + Patches both the home module (``exp``) and the ``utils`` package + re-export so callers that import via either path see the wrapper. + + Safe to call repeatedly; subsequent calls are no-ops. Silently does + nothing if ``nemo-skills`` isn't importable. + """ + global _PATCHED + if _PATCHED: + return + + try: + import nemo_skills.pipeline.utils as utils_mod + import nemo_skills.pipeline.utils.exp as exp_mod + except ImportError: + return + + _PATCHED = True + + _original_get_executor = exp_mod.get_executor + + @functools.wraps(_original_get_executor) + def _patched_get_executor(cluster_config, *args, **kwargs): + extras = parse_extra_sbatch_args(cluster_config) + if extras: + existing = kwargs.get("sbatch_kwargs") or {} + merged = {**extras, **existing} + kwargs["sbatch_kwargs"] = merged + return _original_get_executor(cluster_config, *args, **kwargs) + + exp_mod.get_executor = _patched_get_executor + utils_mod.get_executor = _patched_get_executor + + +def _reset_for_tests() -> None: + """Test-only hook: clear the idempotency flag and undo monkey-patches.""" + global _PATCHED + if not _PATCHED: + return + try: + import nemo_skills.pipeline.utils as utils_mod + import nemo_skills.pipeline.utils.exp as exp_mod + + for mod in (exp_mod, utils_mod): + wrapped = getattr(mod, "get_executor", None) + original = getattr(wrapped, "__wrapped__", None) + if original is not None: + mod.get_executor = original + except ImportError: + pass + _PATCHED = False diff --git a/nvflow/recipes/finance/stages/rl/training.py b/nvflow/recipes/finance/stages/rl/training.py index 2f75001..ae33ec3 100644 --- a/nvflow/recipes/finance/stages/rl/training.py +++ b/nvflow/recipes/finance/stages/rl/training.py @@ -151,7 +151,9 @@ def _resolve_nemo_rl_config(self, config: dict) -> dict: config_paths = [VLLM_MODEL_FOR_TRAINING] for env_cfg in environments.values(): config_paths.extend(env_cfg.get("config_paths", [])) - merged.setdefault("env", {}).setdefault("nemo_gym", {})["config_paths"] = config_paths + nemo_gym = merged.setdefault("env", {}).setdefault("nemo_gym", {}) + nemo_gym["config_paths"] = config_paths + nemo_gym["skip_venv_if_present"] = True if config.get("training_datasets"): merged["data"]["train"] = config["training_datasets"] @@ -582,12 +584,14 @@ def _build_train_cmd( ) or get_timeout_str(cluster_config, partition) hf_model = config.get("hf_checkpoint_path", config["model_name"]) + from nvflow.lib.runtime import NRL_PYTHON_PREAMBLE + cmd = ( + f"{NRL_PYTHON_PREAMBLE} && " f"{config_snippet} && " f"export PYTHONPATH=$PYTHONPATH:/nemo_run/code:/opt/NeMo-RL && " - f"export UV_PROJECT=/opt/NeMo-RL && " f"echo 'Starting training' && " - f"uv run --active python /opt/NeMo-RL/examples/nemo_gym/run_grpo_nemo_gym.py " + f"$NRL_PYTHON /opt/NeMo-RL/examples/nemo_gym/run_grpo_nemo_gym.py " f" --config {config_path}" f" ++policy.model_name={hf_model}" f" ++cluster.gpus_per_node={prepared.num_gpus}" diff --git a/nvflow/recipes/finance/stages/sdg/create_seed_data.py b/nvflow/recipes/finance/stages/sdg/create_seed_data.py index 2758fcf..4c2f2ed 100644 --- a/nvflow/recipes/finance/stages/sdg/create_seed_data.py +++ b/nvflow/recipes/finance/stages/sdg/create_seed_data.py @@ -53,9 +53,6 @@ def execute( console.detail("Num seed questions", num_seed_questions) console.blank() - install_cmd = "pip install -q --root-user-action=ignore jsonlines datasets edgartools requests beautifulsoup4" - - # Build command cmd = ( f"python -m nvflow.recipes.finance.utils.sdg.create_seed_data " f"--output_file {output_file} " @@ -79,7 +76,6 @@ def execute( log_dir=log_dir, expname=expname, run_after=run_after, - installation_command=install_cmd, **config.get("stage_kwargs", {}), ) diff --git a/nvflow/recipes/finance/stages/sdg/generate_questions.py b/nvflow/recipes/finance/stages/sdg/generate_questions.py index bdacf7b..362cf67 100644 --- a/nvflow/recipes/finance/stages/sdg/generate_questions.py +++ b/nvflow/recipes/finance/stages/sdg/generate_questions.py @@ -57,17 +57,16 @@ def execute( console.detail("Inline args", str(inline_args)) console.blank() - prepped_file = input_file.replace(".jsonl", "_prepped.jsonl") + output_dir = Path(output_file).parent + prepped_file = str(output_dir / (Path(input_file).stem + "_prepped.jsonl")) console.status(f"Prepped file: {prepped_file}") - output_dir = Path(output_file).parent generation_folder = output_dir / Path(input_file).stem prep_log_dir = str(generation_folder / "prep-logs") # Run prep command to generate company/year combinations run_cmd( ctx=wrap_arguments( - f"pip install -q --root-user-action=ignore jsonlines && " f"python3 -m nvflow.recipes.finance.utils.sdg.prepare_question_gen_data " f"--input_file {input_file} --company_info_file {company_info_file} " f"--start_year {start_year} --end_year {end_year} --output_file {prepped_file}" diff --git a/nvflow/recipes/finance/stages/sdg/map_questions_to_context.py b/nvflow/recipes/finance/stages/sdg/map_questions_to_context.py index aa6f871..d50eda3 100644 --- a/nvflow/recipes/finance/stages/sdg/map_questions_to_context.py +++ b/nvflow/recipes/finance/stages/sdg/map_questions_to_context.py @@ -58,7 +58,6 @@ def execute( run_cmd( ctx=wrap_arguments( - f"pip install -q --root-user-action=ignore jsonlines tiktoken markdownify && " f"python3 -m nvflow.recipes.finance.utils.shared.question_context_utils --input_file {input_file} --filings_metadata {filings_metadata} --filings_dir {filings_dir} --output_file {output_file} --token_limit {token_limit}" ), cluster=cluster, diff --git a/nvflow/recipes/finance/stages/sft/training.py b/nvflow/recipes/finance/stages/sft/training.py index a172f84..f05760a 100644 --- a/nvflow/recipes/finance/stages/sft/training.py +++ b/nvflow/recipes/finance/stages/sft/training.py @@ -561,12 +561,14 @@ def _build_train_cmd( timeout = get_timeout_str(cluster_config, partition) hf_model = config.get("hf_checkpoint_path", config["model_name"]) + from nvflow.lib.runtime import NRL_PYTHON_PREAMBLE + cmd = ( + f"{NRL_PYTHON_PREAMBLE} && " f"{config_snippet} && " f"export PYTHONPATH=$PYTHONPATH:/nemo_run/code:/opt/NeMo-RL && " - f"export UV_PROJECT=/opt/NeMo-RL && " f"echo 'Starting training' && " - f"uv run --active python /opt/NeMo-RL/examples/run_sft.py " + f"$NRL_PYTHON /opt/NeMo-RL/examples/run_sft.py " f" --config {config_path}" f" ++policy.model_name={hf_model}" f" ++cluster.gpus_per_node={prepared.num_gpus}" diff --git a/nvflow/recipes/finance/utils/evaluation/checkpoint_converter.py b/nvflow/recipes/finance/utils/evaluation/checkpoint_converter.py index 6a961f9..f812dd2 100644 --- a/nvflow/recipes/finance/utils/evaluation/checkpoint_converter.py +++ b/nvflow/recipes/finance/utils/evaluation/checkpoint_converter.py @@ -188,6 +188,27 @@ def get_hf_output_paths(run_path: str | Path, step: int) -> tuple[Path, Path]: return hf_model_path, convert_log_dir +MEGATRON_VENV_PYTHON = ( + "/opt/ray_venvs/" + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" + "/bin/python" +) + +# DTensor v1 Ray venv -- has the fsdp extra (torch.distributed.checkpoint). +DTENSOR_V1_VENV_PYTHON = ( + "/opt/ray_venvs/" + "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker" + "/bin/python" +) + +# DTensor v2 Ray venv -- has the automodel extra (nemo_automodel). +DTENSOR_V2_VENV_PYTHON = ( + "/opt/ray_venvs/" + "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2" + "/bin/python" +) + + def build_conversion_script( megatron_path: str | Path, hf_output_path: str | Path, @@ -197,6 +218,9 @@ def build_conversion_script( Build a bash script that runs the conversion. The script invokes this module as a CLI tool on the cluster. + Uses the Megatron Ray venv Python because the megatron package + (including megatron.bridge) is only installed there, not in + /opt/nemo_rl_venv/. Args: megatron_path: Path to Megatron checkpoint @@ -206,11 +230,14 @@ def build_conversion_script( Returns: Bash script as a string """ + from nvflow.lib.runtime import ray_venv_python_preamble + + megatron_preamble = ray_venv_python_preamble(MEGATRON_VENV_PYTHON, "mcore") script = f""" set -e +{megatron_preamble} -export UV_PROJECT=/opt/NeMo-RL -uv run --extra mcore python -m nvflow.recipes.finance.utils.evaluation.checkpoint_converter \\ +$CONVERT_PYTHON -m nvflow.recipes.finance.utils.evaluation.checkpoint_converter \\ --megatron-path "{megatron_path}" \\ --hf-output-path "{hf_output_path}" \\ --model-name "{model_name}" @@ -226,8 +253,12 @@ def build_dcp_conversion_script( """Build a bash script that converts a DTensor checkpoint to HF format. Supports both checkpoint formats: - - **v1 (DCP)**: ``.metadata`` file โ†’ calls ``convert_dcp_to_hf.py`` - - **v2 (safetensors)**: ``shard-*.safetensors`` files โ†’ calls ``offline_hf_consolidation.py`` + - **v1 (DCP)**: ``.metadata`` โ†’ ``convert_dcp_to_hf.py`` in the DTensor v1 Ray venv. + - **v2 (safetensors)**: ``shard-*.safetensors`` โ†’ ``offline_hf_consolidation.py`` + in the DTensor v2 Ray venv. + + Each branch prefers the pre-built Ray venv if available, falling back + to ``uv run --extra`` in dev mode. The script runs ON THE CLUSTER. It resolves the run subdirectory under checkpoint_path (flat or GRPO layout), auto-detects the format, then @@ -239,6 +270,8 @@ def build_dcp_conversion_script( step: Checkpoint step number hf_output_path: Where to write HF model (known at submit time) """ + from nvflow.lib.runtime import ray_venv_python_preamble + step_name = f"step_{step}" script = f""" set -euo pipefail @@ -289,9 +322,8 @@ def build_dcp_conversion_script( --recreate-hf-metadata "$MODEL_DIR" "$STEP_DIR/config.yaml" fi - cd /opt/NeMo-RL - export UV_PROJECT=/opt/NeMo-RL - uv run --extra automodel python /opt/NeMo-RL/3rdparty/Automodel-workspace/Automodel/tools/offline_hf_consolidation.py \\ + {ray_venv_python_preamble(DTENSOR_V2_VENV_PYTHON, "automodel")} + $CONVERT_PYTHON /opt/NeMo-RL/3rdparty/Automodel-workspace/Automodel/tools/offline_hf_consolidation.py \\ --model-name unused \\ --input-dir "$MODEL_DIR" \\ --output-dir "$HF_OUTPUT" @@ -302,8 +334,9 @@ def build_dcp_conversion_script( echo "Detected DTensor v1 (DCP) checkpoint" echo "Converting: $STEP_DIR -> $HF_OUTPUT" + {ray_venv_python_preamble(DTENSOR_V1_VENV_PYTHON, "fsdp")} cd /opt/NeMo-RL - uv run examples/converters/convert_dcp_to_hf.py \\ + $CONVERT_PYTHON examples/converters/convert_dcp_to_hf.py \\ --config="$STEP_DIR/config.yaml" \\ --dcp-ckpt-path="$WEIGHTS_DIR" \\ --hf-ckpt-path="$HF_OUTPUT" diff --git a/nvflow/recipes/finance/workflows/download_sec_filings.yaml b/nvflow/recipes/finance/workflows/download_sec_filings.yaml index 42a31f8..fbb5eec 100644 --- a/nvflow/recipes/finance/workflows/download_sec_filings.yaml +++ b/nvflow/recipes/finance/workflows/download_sec_filings.yaml @@ -26,12 +26,6 @@ stages: sec_identity_email: your.email@email.com # UPDATE: Your email (required by SEC) sec_identity_company: YourCompany # UPDATE: Your company name dependencies: [] - stage_kwargs: - installation_command: >- - pip install -q --root-user-action=ignore - edgartools==5.20.2 sec-parser pandas pyarrow tqdm httpx tzdata && - ln -sfn $(python -c "import tzdata; print(tzdata.__path__[0])")/zoneinfo - /usr/share/zoneinfo 2>/dev/null || true # Demo stage: 7 companies for quick testing demo: @@ -40,12 +34,6 @@ stages: sec_identity_email: your.email@email.com # UPDATE: Your email (required by SEC) sec_identity_company: YourCompany # UPDATE: Your company name dependencies: [] - stage_kwargs: - installation_command: >- - pip install -q --root-user-action=ignore - edgartools==5.20.2 sec-parser pandas pyarrow tqdm httpx tzdata && - rm -rf /usr/share/zoneinfo && - ln -s $(python -c "import tzdata; print(tzdata.__path__[0])")/zoneinfo /usr/share/zoneinfo # Production stage: Full S&P 500 sap-500: @@ -54,9 +42,3 @@ stages: sec_identity_email: your.email@email.com # UPDATE: Your email (required by SEC) sec_identity_company: YourCompany # UPDATE: Your company name dependencies: [] - stage_kwargs: - installation_command: >- - pip install -q --root-user-action=ignore - edgartools==5.20.2 sec-parser pandas pyarrow tqdm httpx tzdata && - rm -rf /usr/share/zoneinfo && - ln -s $(python -c "import tzdata; print(tzdata.__path__[0])")/zoneinfo /usr/share/zoneinfo diff --git a/nvflow/recipes/finance/workflows/eval/base.yaml b/nvflow/recipes/finance/workflows/eval/base.yaml index 44a6c17..ad30c7c 100644 --- a/nvflow/recipes/finance/workflows/eval/base.yaml +++ b/nvflow/recipes/finance/workflows/eval/base.yaml @@ -76,8 +76,12 @@ benchmarks: datasets_dir: /workspace/nvflow/recipes/finance/datasets +# All eval deps are in the nemo-skills container +installation_command: "true" + # num_gpus defaults to cluster's gpus_per_node at runtime (portable across clusters) -conversion: {} +conversion: + installation_command: "true" # ============================================================================ # Prepare Data Stage (used by baselines.yaml and shared reference) @@ -89,3 +93,5 @@ stages: # Re-add finance_agent to this list when the benchmark is re-enabled. dataset_names: [secque, financebench] output_dir: /workspace/nvflow/recipes/finance/datasets + stage_kwargs: + installation_command: "true" diff --git a/nvflow/recipes/finance/workflows/eval/demo.yaml b/nvflow/recipes/finance/workflows/eval/demo.yaml index 4c3b885..d95e070 100644 --- a/nvflow/recipes/finance/workflows/eval/demo.yaml +++ b/nvflow/recipes/finance/workflows/eval/demo.yaml @@ -19,6 +19,21 @@ workflow: cluster: my_cluster base_output_dir: /workspace/outputs/finance/demo/workflow-1-baseline-eval +# Local vLLM judge (overrides OpenAI default in base.yaml) +judge: + model: /hf_models/openai/gpt-oss-120b + server_type: vllm + server_address: null + server_gpus: 2 + server_nodes: 1 + server_args: "--max-model-len 32768 --async-scheduling" + extra_args: >- + ++prompt_config=/workspace/nvflow/recipes/finance/prompts/sec_judge.yaml + ++inference.temperature=0.0 + ++inference.tokens_to_generate=8192 + ++generation_key=judgement + ++skip_extraction=True + pipeline_stages: - prepare_data - qwen3-4b diff --git a/nvflow/recipes/finance/workflows/grpo/base.yaml b/nvflow/recipes/finance/workflows/grpo/base.yaml index 9667f5d..7fb51cf 100644 --- a/nvflow/recipes/finance/workflows/grpo/base.yaml +++ b/nvflow/recipes/finance/workflows/grpo/base.yaml @@ -21,10 +21,16 @@ # Each environment has its own data pipeline and NeMo-Gym config. # Use --environment to run a single environment. # -# NeMo-RL / NeMo-Gym versions (pinned for reproducibility): -# NeMo-RL main: e5a729cc438ea71bafa7138204f861196598a9b2 -# NeMo-Gym feat/finance-sec-search-improvements: f24573bb157a1b4799591df1087d6db0b25e4e4c -# Mounted as overlay via my_cluster.yaml at /opt/NeMo-RL. +# NeMo-RL / NeMo-Gym versions: +# Baked into the air-gapped `nvflow-nemo-rl` image (see dockerfiles/Dockerfile.nemo-rl): +# - Base: nvcr.io/nvidia/nemo-rl:v0.6.0 +# - NeMo-Skills: 022904023ad7a83a87662a313cf72e7df5891d55 (`0229040`) +# - NeMo-Gym branch: ude/finance-sec-search-v2 (override via NEMO_GYM_BRANCH build arg) +# - Pre-built Gym .venv symlinked across all 6 components +# Do NOT bind-mount host clones at /opt/NeMo-RL or +# /opt/NeMo-RL/3rdparty/Gym-workspace/Gym in production -- the overlays +# shadow the baked .venv and break GRPO stages. See INSTALL.md +# "Setup NeMo-RL & NeMo-Gym Sources (for GRPO)" for the dev-mode exception. # ============================================================================ # Recipe identifier (defines which recipe this workflow belongs to) @@ -36,8 +42,9 @@ workflow: type: "training" # Workflow category description: "GRPO RL training for financial reasoning with NeMo-Gym" -# Cluster to run jobs on (references cluster_configs/.yaml) -# Uses my_cluster which mounts NeMo-RL + Gym overlays at /opt/NeMo-RL. +# Cluster to run jobs on (references cluster_configs/.yaml). +# my_cluster.yaml selects the air-gapped `nvflow-nemo-rl` image, which already +# contains NeMo-RL source + a pre-built Gym .venv -- no host overlays needed. cluster: my_cluster # Root directory for all finance pipeline outputs (shared across workflows). @@ -379,15 +386,14 @@ stages: num_gpus: 1 # nemo-rl container needs GPU node for nvidia-container-cli gym_path: "/opt/NeMo-RL/3rdparty/Gym-workspace/Gym" - # Sets up a NeMo-Gym venv in /tmp so writes stay container-local (not on - # the bind-mounted /opt/NeMo-RL). SLURM_JOB_ID suffix guarantees - # uniqueness even if two jobs share a node. + # Activate NeMo-Gym venv: use pre-built if available, otherwise build via uv. # Anchored as &gym_install -- reused by prefetch_cache, collect_rollouts, and compute_rewards. installation_command: &gym_install >- - cd /opt/NeMo-RL/3rdparty/Gym-workspace/Gym - && uv venv /tmp/gym-venv-\${SLURM_JOB_ID} --python 3.12 - && source /tmp/gym-venv-\${SLURM_JOB_ID}/bin/activate - && uv sync --active --extra dev + if [ -f /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate ]; then + source /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate; + else + cd /opt/NeMo-RL && uv sync --extra nemo_gym; + fi # -------------------------------------------------------------------------- # Prefetch Cache (optional, CPU-only; no step-N -- writes to cache-finance-sec-search) @@ -624,9 +630,8 @@ stages: wandb_project: finance-grpo wandb_mode: disabled # online | offline | disabled - # Installs NeMo-Gym dependencies in the NeMo-RL container. - installation_command: >- - cd /opt/NeMo-RL && uv sync --extra nemo_gym + # NeMo-Gym dependencies are pre-installed in the container. + installation_command: "true" # Runtime environment customization # stage_kwargs: diff --git a/nvflow/recipes/finance/workflows/grpo/qwen3_4b.yaml b/nvflow/recipes/finance/workflows/grpo/qwen3_4b.yaml index 8bd2806..264824f 100644 --- a/nvflow/recipes/finance/workflows/grpo/qwen3_4b.yaml +++ b/nvflow/recipes/finance/workflows/grpo/qwen3_4b.yaml @@ -57,8 +57,11 @@ _training_policy_32k: &training_policy_32k vllm_cfg: max_model_len: ${_env_seq_len_32k} sequence_packing: + enabled: true train_mb_tokens: ${_env_seq_len_32k} logprob_mb_tokens: ${_env_seq_len_32k} + logprob_chunk_size: 2048 + make_sequence_length_divisible_by: 16 environments: equivalence_llm_judge: @@ -181,6 +184,7 @@ stages: eval: eval_steps: [20] checkpoint_path: ${directories.step-8-training}/equivalence_llm_judge/grpo-qwen3-4b-16g-tp2-cp1-seq32k + eval_output_dir: ${directories.step-9-eval}/equivalence_llm_judge base_output_dir: ${model_output_dir}/equivalence_llm_judge format: fsdp # Triggers DCP/safetensors โ†’ HF conversion (auto-detects v1/v2) base_model: ${hf_model_path} diff --git a/nvflow/recipes/finance/workflows/grpo/qwen3_4b_finsec.yaml b/nvflow/recipes/finance/workflows/grpo/qwen3_4b_finsec.yaml index 26816d2..f0e3d5e 100644 --- a/nvflow/recipes/finance/workflows/grpo/qwen3_4b_finsec.yaml +++ b/nvflow/recipes/finance/workflows/grpo/qwen3_4b_finsec.yaml @@ -124,6 +124,7 @@ stages: eval: eval_steps: [20] checkpoint_path: ${directories.step-8-training}/finance_sec_search/grpo-qwen3-4b-64g-tp2-cp8-seq128k + eval_output_dir: ${directories.step-9-eval}/finance_sec_search base_output_dir: ${model_output_dir}/finance_sec_search format: megatron base_model: ${hf_model_path} diff --git a/nvflow/recipes/finance/workflows/sdg/template-based-sdg-demo.yaml b/nvflow/recipes/finance/workflows/sdg/template-based-sdg-demo.yaml index 93232da..c000f2e 100644 --- a/nvflow/recipes/finance/workflows/sdg/template-based-sdg-demo.yaml +++ b/nvflow/recipes/finance/workflows/sdg/template-based-sdg-demo.yaml @@ -21,7 +21,6 @@ filings_dir: /workspace/outputs/finance/demo/workflow-2-download-sec/step-0-down # Demo Stage Overrides - Only override what's different from prod # ============================================================================ stages: - # -------------------------------------------------------------------------- # Stage 0: Create Seed Data - Filter to 7 companies and 30 questions for demo # -------------------------------------------------------------------------- diff --git a/nvflow/recipes/finance/workflows/sft/base.yaml b/nvflow/recipes/finance/workflows/sft/base.yaml index 51587a8..c9aac66 100644 --- a/nvflow/recipes/finance/workflows/sft/base.yaml +++ b/nvflow/recipes/finance/workflows/sft/base.yaml @@ -194,6 +194,10 @@ stages: # This provides pyproject.toml on all nodes for the Ray _env_builder, # and resolves the UV_PROJECT=/opt/NeMo-RL path used by training.py. + # All training deps are in the nemo-rl container + stage_kwargs: + installation_command: "true" + # -------------------------------------------------------------------------- # Stage 5: Evaluate Checkpoints on Finance Benchmarks # -------------------------------------------------------------------------- diff --git a/nvflow/recipes/finance/workflows/sft/nemotron_nano_9b.yaml b/nvflow/recipes/finance/workflows/sft/nemotron_nano_9b.yaml deleted file mode 100644 index d348b1f..0000000 --- a/nvflow/recipes/finance/workflows/sft/nemotron_nano_9b.yaml +++ /dev/null @@ -1,140 +0,0 @@ -# ============================================================================ -# NVIDIA-Nemotron-Nano-9B-v2 SFT Training Configuration -# ============================================================================ -# Inherits shared pipeline logic from base.yaml and provides all -# Nemotron-Nano-9B-v2 specific settings (tokenizer, model paths, parallelism). -# -# Model: https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2 -# Architecture: Mamba2-Transformer Hybrid (4 attention layers) -# Context Length: 128K -# Reasoning: Thinking model - uses /think system prompt to enable reasoning -# -# Usage: -# uv run nflow run-all --config workflows/sft/nemotron_nano_9b.yaml -# ============================================================================ - -# Inherit base SFT workflow -_base_: base.yaml - -# Use cluster config with Nemotron-Nano-9B-v2 container support -cluster: my_cluster_nemotron - -# Model-specific output directory (use /workspace for cluster mounts) -base_output_dir: /workspace/outputs/finance/sap-500/workflow-4-sft/nemotron_nano_9b - -# ============================================================================ -# Model-Specific Stage Overrides -# ============================================================================ - -stages: - - # -------------------------------------------------------------------------- - # Stage 1: Prepare for SFT - Nemotron-Nano-9B tokenizer - # -------------------------------------------------------------------------- - # Nemotron uses /think system prompt to enable reasoning (not enable_thinking) - prepare_for_sft: - prepare_data_kwargs: - ctx_args: >- - ++tokenizer=/hf_models/nvidia/NVIDIA-Nemotron-Nano-9B-v2 - ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml - ++system_message="/think" - - # -------------------------------------------------------------------------- - # Stage 3: Sequence Length Grouping - Nemotron-Nano-9B tokenizer - # -------------------------------------------------------------------------- - sequence_length_grouping: - tokenizer_path: /hf_models/nvidia/NVIDIA-Nemotron-Nano-9B-v2 - - # -------------------------------------------------------------------------- - # Stage 4: Training - Nemotron-Nano-9B-v2 configuration (NeMo-RL format) - # -------------------------------------------------------------------------- - training: - model_name: nvidia/NVIDIA-Nemotron-Nano-9B-v2 - hf_checkpoint_path: /hf_models/nvidia/NVIDIA-Nemotron-Nano-9B-v2 - - # Use Megatron backend (standard for Nemotron-Nano-9B) - backend: megatron - - # Cluster configuration for 9B model - total_gpus: 16 - dependent_jobs: 1 # Test resume with 1 dependent job - - # Use base preset with model-specific overrides - preset: "sft-base" - - # Nemotron-Nano-9B specific configuration - # All hyperparameters preserved from legacy nemotron-nano-9b preset - overrides: - sft: - max_num_epochs: 1 # QUICK TEST CONFIG - max_num_steps: 100 # Stop after 100 steps - val_period: 50 - - # === FULL TRAINING CONFIG (uncomment when ready) === - # max_num_epochs: 3 - # val_period: 100 - - checkpointing: - save_period: 50 # Quick test: save every 50 steps - # save_period: 250 # Full training: save every 250 steps - keep_top_k: 50 - checkpoint_must_save_by: "00:03:40:00" - - policy: - train_global_batch_size: 16 # QUICK TEST: Smaller batch for faster iteration - # train_global_batch_size: 128 # FULL TRAINING - train_micro_batch_size: 1 - max_total_sequence_length: 49152 # 48K (matches max_token_length filter) - - # Megatron parallelism: TP ร— PP ร— CP = 4ร—1ร—4 = 16 (must equal total_gpus) - megatron_cfg: - tensor_model_parallel_size: 4 - pipeline_model_parallel_size: 1 - context_parallel_size: 4 - sequence_parallel: true - activation_checkpointing: true - bias_activation_fusion: true - apply_rope_fusion: true # Applies to 4 attention layers - empty_unused_memory_level: 1 # Memory optimization (moderate) - - # Optimizer settings (from legacy nemotron-nano-9b preset) - optimizer: - optimizer: "adam" - lr: 5e-6 - min_lr: 5e-7 - weight_decay: 0.01 - bf16: true - clip_grad: 1.0 - use_distributed_optimizer: true - - # Scheduler settings - scheduler: - lr_decay_style: "cosine" - lr_warmup_iters: 10 # QUICK TEST: Must be < max_num_steps - # lr_warmup_iters: 2000 # FULL TRAINING - lr_warmup_init: 1e-7 - - data: - force_reprocess: false - num_workers: 10 - - # Debug: understand the Python environment - stage_kwargs: - installation_command: >- - echo "=== ENVIRONMENT DEBUG ===" && - echo "Which python: $(which python)" && - echo "Python version: $(python --version)" && - echo "VIRTUAL_ENV=$VIRTUAL_ENV" && - echo "PYTHONPATH=$PYTHONPATH" && - echo "NEMO_RL_VENV_DIR=$NEMO_RL_VENV_DIR" && - echo "=== Check if datasets is installed ===" && - python -c "import datasets; print(f'datasets found: {datasets.__file__}')" || echo "datasets NOT found" && - echo "=== pip list | grep datasets ===" && - pip list | grep -i datasets || echo "No datasets in pip list" && - echo "=== Site packages location ===" && - python -c "import site; print(site.getsitepackages())" && - echo "=== Installing datasets explicitly ===" && - pip install datasets && - echo "=== Verify datasets after install ===" && - python -c "import datasets; print(f'datasets installed: {datasets.__file__}')" && - sed -i 's/hf_overrides=hf_overrides,//' /nemo_run/code/nemo_skills/training/nemo_rl/convert_megatron_to_hf.py diff --git a/scripts/convert_checkpoint_to_hf.sh b/scripts/convert_checkpoint_to_hf.sh index f53437a..951a662 100755 --- a/scripts/convert_checkpoint_to_hf.sh +++ b/scripts/convert_checkpoint_to_hf.sh @@ -115,15 +115,16 @@ if [[ "$BACKEND" == "megatron" ]]; then CONVERT_CALL+="hf_output_path=\\\"${OUTPUT_DIR}\\\", " CONVERT_CALL+="model_name=\\\"${MODEL_NAME}\\\")" - FULL_CMD="export UV_PROJECT=/opt/NeMo-RL \ - && export PYTHONPATH=\$PYTHONPATH:/nemo_run/code \ - && uv run --extra mcore python -c \ + # MegatronPolicyWorker Ray venv has the mcore extra; PYTHONPATH lets it find nvflow. + MCORE_VENV_PYTHON="/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/bin/python" + FULL_CMD="export PYTHONPATH=\$PYTHONPATH:/nemo_run/code \ + && ${MCORE_VENV_PYTHON} -c \ \"from nvflow.recipes.finance.utils.evaluation.checkpoint_converter import convert_checkpoint; ${CONVERT_CALL}\"" else - # DTensor/FSDP: NeMo-RL's converter + rsync tokenizer from checkpoint - FULL_CMD="export UV_PROJECT=/opt/NeMo-RL \ - && cd /opt/NeMo-RL \ - && uv run examples/converters/convert_dcp_to_hf.py \ + # DTensor/FSDP: DTensorPolicyWorker Ray venv has the fsdp extra. + DTENSOR_V1_VENV_PYTHON="/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker/bin/python" + FULL_CMD="cd /opt/NeMo-RL \ + && ${DTENSOR_V1_VENV_PYTHON} examples/converters/convert_dcp_to_hf.py \ --config=\"${STEP_PATH}/config.yaml\" \ --dcp-ckpt-path=\"${STEP_PATH}/policy/weights\" \ --hf-ckpt-path=\"${OUTPUT_DIR}\" \ diff --git a/scripts/serve_vllm_patched.py b/scripts/serve_vllm_patched.py index f97df77..b680e69 100644 --- a/scripts/serve_vllm_patched.py +++ b/scripts/serve_vllm_patched.py @@ -191,10 +191,21 @@ def _patch_hermes_tool_parser() -> None: # WORKAROUND(vllm-0.17-hermes) def _ensure_tiktoken_cache() -> None: # WORKAROUND(harmony-aarch64) - """Download tiktoken vocab files if on aarch64 and set env vars.""" + """Download tiktoken vocab files if on aarch64 and set env vars. + + Skips the download when TIKTOKEN_CACHE_DIR or TIKTOKEN_RS_CACHE_DIR is + already set (e.g. pointing at files baked into the container), which is + required for airgap / offline environments. + """ if platform.machine() not in ("aarch64", "arm64"): return + existing = os.environ.get("TIKTOKEN_CACHE_DIR") or os.environ.get("TIKTOKEN_RS_CACHE_DIR") + if existing: + print(f"{_TAG} Tiktoken cache already configured ({existing}), skipping download.") + os.environ.setdefault("TIKTOKEN_ENCODINGS_BASE", existing) + return + cache_dir = Path("/tmp/tiktoken-encodings") cache_dir.mkdir(parents=True, exist_ok=True) diff --git a/tests/test_sbatch.py b/tests/test_sbatch.py new file mode 100644 index 0000000..1b1c63b --- /dev/null +++ b/tests/test_sbatch.py @@ -0,0 +1,208 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for ``nvflow.lib.sbatch``. + +Covers the ``parse_extra_sbatch_args`` parser and the ``get_executor`` +autopatch that plumbs cluster-level ``extra_sbatch_args`` into every +Slurm submission. +""" + +from __future__ import annotations + +import sys +import types +from typing import Any + +import pytest + +from nvflow.lib import sbatch +from nvflow.lib.sbatch import parse_extra_sbatch_args + +# --------------------------------------------------------------------------- +# parse_extra_sbatch_args +# --------------------------------------------------------------------------- + + +def test_parse_value_form(): + cfg = {"extra_sbatch_args": ["--qos=xlarge_qos", "--partition=gb200"]} + assert parse_extra_sbatch_args(cfg) == {"qos": "xlarge_qos", "partition": "gb200"} + + +def test_parse_boolean_form(): + cfg = {"extra_sbatch_args": ["--exclusive", "--requeue"]} + assert parse_extra_sbatch_args(cfg) == {"exclusive": True, "requeue": True} + + +def test_parse_strips_dashes(): + cfg = {"extra_sbatch_args": ["---qos=xlarge_qos"]} + assert parse_extra_sbatch_args(cfg) == {"qos": "xlarge_qos"} + + +def test_parse_skips_garbage_entries(): + cfg = {"extra_sbatch_args": ["", "--qos=xlarge_qos", None, 42, "---"]} # type: ignore[list-item] + assert parse_extra_sbatch_args(cfg) == {"qos": "xlarge_qos"} + + +@pytest.mark.parametrize( + "cluster_config", + [None, {}, {"extra_sbatch_args": None}, {"extra_sbatch_args": []}], # type: ignore[dict-item] +) +def test_parse_empty(cluster_config): + assert parse_extra_sbatch_args(cluster_config) == {} + + +_REPO_ROOT = __import__("pathlib").Path(__file__).resolve().parent.parent + + +def test_rollout_uses_shared_parser(): + """rollout.py must use parse_extra_sbatch_args, not an inline loop. + + Reads the file directly so the test does not require ``nemo_skills`` + (which ``rollout`` imports at module load) to be installed. + """ + src = (_REPO_ROOT / "nvflow" / "lib" / "rl" / "rollout.py").read_text() + assert "parse_extra_sbatch_args" in src + + +def test_autopatch_is_installed_by_workflow_runner_not_cli(): + """Autopatch must be in WorkflowRunner.run(), not CLI startup. + + Reads the files directly so the test does not require ``typer`` + (a CLI-only dep) to be installed. + """ + cli_src = (_REPO_ROOT / "nvflow" / "cli" / "main.py").read_text() + runner_src = (_REPO_ROOT / "nvflow" / "core" / "workflow_runner.py").read_text() + + assert "apply_sbatch_args_autopatch" not in cli_src + assert "apply_sbatch_args_autopatch" in runner_src + + +# --------------------------------------------------------------------------- +# apply_sbatch_args_autopatch +# --------------------------------------------------------------------------- + + +@pytest.fixture +def fake_nemo_skills(monkeypatch): + """Minimal fake nemo_skills so the autopatch can target get_executor.""" + calls: list[dict[str, Any]] = [] + + def _make(name: str, **attrs): + mod = types.ModuleType(name) + for k, v in attrs.items(): + setattr(mod, k, v) + sys.modules[name] = mod + return mod + + def _original_get_executor(cluster_config, *args, **kwargs): + calls.append({"cluster_config": cluster_config, "args": args, "kwargs": kwargs}) + + _make("nemo_skills") + _make("nemo_skills.pipeline") + utils_mod = _make("nemo_skills.pipeline.utils", get_executor=_original_get_executor) + exp_mod = _make("nemo_skills.pipeline.utils.exp", get_executor=_original_get_executor) + + sbatch._reset_for_tests() + try: + yield exp_mod, utils_mod, calls + finally: + sbatch._reset_for_tests() + for mod_name in [ + "nemo_skills.pipeline.utils.exp", + "nemo_skills.pipeline.utils", + "nemo_skills.pipeline", + "nemo_skills", + ]: + sys.modules.pop(mod_name, None) + + +def test_autopatch_injects_qos_into_get_executor(fake_nemo_skills): + exp_mod, _, calls = fake_nemo_skills + + sbatch.apply_sbatch_args_autopatch() + + exp_mod.get_executor( + {"extra_sbatch_args": ["--qos=xlarge_qos"]}, + "container", + num_nodes=1, + ) + + assert len(calls) == 1 + assert calls[0]["kwargs"]["sbatch_kwargs"] == {"qos": "xlarge_qos"} + + +def test_autopatch_existing_kwargs_win(fake_nemo_skills): + exp_mod, _, calls = fake_nemo_skills + + sbatch.apply_sbatch_args_autopatch() + + exp_mod.get_executor( + {"extra_sbatch_args": ["--qos=xlarge_qos", "--time=02:00:00"]}, + "container", + sbatch_kwargs={"qos": "explicit_override"}, + ) + + assert calls[0]["kwargs"]["sbatch_kwargs"] == { + "qos": "explicit_override", + "time": "02:00:00", + } + + +def test_autopatch_noop_when_no_extras(fake_nemo_skills): + exp_mod, _, calls = fake_nemo_skills + + sbatch.apply_sbatch_args_autopatch() + + exp_mod.get_executor({"extra_sbatch_args": []}, "container") + + assert len(calls) == 1 + assert calls[0]["kwargs"].get("sbatch_kwargs") is None + + +def test_autopatch_patches_utils_reexport(fake_nemo_skills): + _, utils_mod, calls = fake_nemo_skills + + sbatch.apply_sbatch_args_autopatch() + + utils_mod.get_executor( + {"extra_sbatch_args": ["--qos=test"]}, + "container", + ) + + assert len(calls) == 1 + assert calls[0]["kwargs"]["sbatch_kwargs"] == {"qos": "test"} + + +def test_autopatch_is_idempotent(fake_nemo_skills): + exp_mod, utils_mod, _ = fake_nemo_skills + + sbatch.apply_sbatch_args_autopatch() + first = (exp_mod.get_executor, utils_mod.get_executor) + + sbatch.apply_sbatch_args_autopatch() + second = (exp_mod.get_executor, utils_mod.get_executor) + + assert first == second + + +def test_autopatch_skips_when_nemo_skills_absent(): + """No crash when nemo-skills isn't importable.""" + sbatch._reset_for_tests() + + for mod_name in list(sys.modules): + if mod_name == "nemo_skills" or mod_name.startswith("nemo_skills."): + sys.modules[mod_name] = None # type: ignore[assignment] + + sbatch.apply_sbatch_args_autopatch() # must not raise