diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..d67219f --- /dev/null +++ b/.dockerignore @@ -0,0 +1,41 @@ +# Build context excludes for dockerfiles/Dockerfile.nvflow (COPY . -> /opt/nvflow). +# The client image bakes the source + a fresh uv venv. We KEEP .git in the context +# (hatch-vcs versioning needs it at build); the image then squashes it to a single +# history-free snapshot commit (nemo-run only needs `git archive HEAD`). We drop +# everything heavy, generated, or secret. Other NVFlow Dockerfiles clone their +# sources (they do not COPY this context), so these excludes are safe for them. + +# Prebuilt/host venv (the image builds its own) and uv cache +.venv/ +**/.venv/ +.uv-cache/ +uv-cache/ + +# Huge untracked local caches / outputs (cache/ is tens of GB) +cache/ +outputs/ +htmlcov/ + +# Python/pytest/tooling artifacts +**/__pycache__/ +*.pyc +*.pyo +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ + +# Editor/local backups. Scope to cluster-config backups only -- a bare *.bak* +# would also drop TRACKED dataset files (e.g. datasets/finance_agent/*.bak), +# which then show up as phantom deletions and mark the baked tree dirty. +cluster_configs/*.bak* +*~ + +# Secrets / personal cluster config (gitignored; must never be baked into an image) +cluster_configs/my_cluster.yaml +**/*.env +.env + +# Container image artifacts that should never enter the context +*.sqsh +*.tar +*.tar.gz diff --git a/.github/workflows/secrets-detector.yml b/.github/workflows/secrets-detector.yml index dcb2ca9..2794a59 100644 --- a/.github/workflows/secrets-detector.yml +++ b/.github/workflows/secrets-detector.yml @@ -30,6 +30,11 @@ jobs: run: | curl -sSfL https://raw.githubusercontent.com/trufflesecurity/trufflehog/main/scripts/install.sh | sh -s -- -b /usr/local/bin + # --results=verified,unknown fails on confirmed secrets and on anything + # TruffleHog could not check, but drops findings it actively verified as + # not a secret. Without it, any 40-hex string near a keyword trips a + # detector -- pinned upstream git SHAs read as Weights & Biases keys. + # Keep the flags in sync with the secrets-detector job in .gitlab-ci.yml. - name: Scan for secrets run: | - trufflehog git file://. --since-commit ${{ github.event.pull_request.base.sha }} --fail --no-update + trufflehog git file://. --since-commit ${{ github.event.pull_request.base.sha }} --fail --no-update --results=verified,unknown diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index d3603af..bb77b63 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -35,9 +35,10 @@ jobs: run: | curl -LsSf https://astral.sh/uv/install.sh | sh uv venv --python 3.12 - uv pip install pytest pytest-cov pytest-timeout - uv pip install PyYAML omegaconf rich + # Dependency list is shared with .gitlab-ci.yml so the two pipelines + # cannot drift. Add test-only deps to tests/requirements-ci.txt. uv pip install -e . --no-deps + uv pip install -r tests/requirements-ci.txt - name: Test run: | diff --git a/.gitignore b/.gitignore index 9d0f4f5..2bc4922 100644 --- a/.gitignore +++ b/.gitignore @@ -25,14 +25,13 @@ wheels/ *.egg # Virtual environments -.venv/ +.venv venv/ ENV/ env/ # UV .uv/ - # IDE .vscode/ .idea/ @@ -46,6 +45,12 @@ env/ htmlcov/ .tox/ +# HF / tokenizer / debug caches. Created by ad-hoc local Python sessions +# (e.g. AutoConfig.from_pretrained(..., trust_remote_code=True)) when +# HF_HOME defaults to $cwd/cache. Pipeline runs use the Lustre HF cache, +# never this directory -- keep it out of git to avoid accidental commits. +/cache/ + # Data and outputs (customize based on your needs) /data /data/ @@ -82,6 +87,7 @@ nvflow/recipes/finance/datasets/finance_agent/*.json # Backup files *.bak +*.bak-* # OS .DS_Store diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 947952c..b5ac76d 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -43,7 +43,24 @@ dco-check: stage: lint image: python:3.12-slim script: - - python scripts/check_dco.py + - python3 scripts/check_dco.py + rules: + - if: $CI_PIPELINE_SOURCE == "merge_request_event" + +# Mirrors .github/workflows/secrets-detector.yml. The GitHub copy only runs on +# pull requests, and this repo opens one per release, so without this job a +# finding stays invisible until release day. Keep the flags in sync with the +# GitHub workflow. +secrets-detector: + stage: lint + image: python:3.12-slim + variables: + GIT_DEPTH: 0 + before_script: + - apt-get update -qq && apt-get install -y -qq git curl > /dev/null + - curl -sSfL https://raw.githubusercontent.com/trufflesecurity/trufflehog/main/scripts/install.sh | sh -s -- -b /usr/local/bin + script: + - trufflehog git file://. --since-commit "$CI_MERGE_REQUEST_DIFF_BASE_SHA" --fail --no-update --results=verified,unknown rules: - if: $CI_PIPELINE_SOURCE == "merge_request_event" @@ -54,11 +71,14 @@ test: stage: test image: python:3.12-slim script: - # Lightweight install: skip heavy core deps (nemo-skills ~200+ packages) - # that unit tests don't need. Only install the project + test deps. + # Lightweight install: the project plus tests/requirements-ci.txt, but not + # the heavy nemo-skills stack (~200 packages incl. torch). Tests that + # genuinely need nemo-skills use pytest.importorskip and skip here; they + # run in the full-deps CI. The dependency list is shared with + # .github/workflows/unit-tests.yml β€” edit it there, in one place. - 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 pip install -e . --no-deps + - uv pip install -r tests/requirements-ci.txt # --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. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 483ff73..df3c38c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -43,5 +43,5 @@ repos: rev: v1.7.1 hooks: - id: mypy - additional_dependencies: [types-PyYAML, types-tqdm] + additional_dependencies: [types-PyYAML, types-tqdm, types-requests] args: [--ignore-missing-imports] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b51dcf6..0d5bde7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -63,12 +63,12 @@ By making a contribution to this project, I certify that: this project or the open source license(s) involved. ``` -## Merge Requests +## Pull Requests 1. Fork the repository and create a branch from `main`. 2. Make your changes with signed-off commits (`git commit -s`). -3. Push your branch and open a Merge Request into `main`. -4. Ensure the MR description references any related issues and describes the change clearly. +3. Push your branch and open a Pull Request into `main`. +4. Ensure the PR description references any related issues and describes the change clearly. ## Code and Documentation @@ -80,7 +80,6 @@ By making a contribution to this project, I certify that: By contributing, you agree that your contributions will be licensed under the same license as the project: the Apache License, Version 2.0. See [LICENSE](LICENSES/LICENSE) for the full text. -## IP Review and Open Source Compliance +## Third-Party Code -- **Ongoing modifications**: For changes to project code (including contributions by third parties), follow NVIDIA's IP review process: [https://nv/ip_review_process](https://nv/ip_review_process). -- **Open Source compliance**: This project follows NVIDIA OSRB recommendations for Apache 2.0 release. +If your contribution adds or updates a third-party dependency, make sure its license is compatible with Apache 2.0 and record it in [`LICENSES/THIRD_PARTY_SW_LICENSE_INFO.md`](LICENSES/THIRD_PARTY_SW_LICENSE_INFO.md). diff --git a/INSTALL.md b/INSTALL.md index 4b7f59f..2e041a2 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -1,31 +1,54 @@ # Installation & Setup Guide -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. +Operator guide for getting NVFlow running on a Slurm cluster: install the client, stage containers and models, configure your cluster, and verify. NVFlow's containers are self-sufficient β€” all dependencies are pre-installed, so once images and models are staged, the pipeline runs fully offline. -## πŸ“‹ Steps +> This guide sets up the **cluster side**. How you run the `nflow` **client** β€” +> local install or the airgapped `nvflow-client` container, and whether it +> submits directly or over an SSH tunnel β€” is summarized in +> [Choose your client setup](#choose-your-client-setup) just below. -1. [Prerequisites](#prerequisites) -2. [Setup Containers](#setup-containers) -3. [Download Models](#download-models) -4. [Setup NeMo-RL & NeMo-Gym Sources (for GRPO)](#setup-nemo-rl--nemo-gym-sources-for-grpo) -5. [Configure Your Cluster](#configure-your-cluster) -6. [Verify Installation](#verify-installation) +> **Building artifacts?** Producing the container images is a maintainer task and +> lives under [`docs/maintainers/`](docs/maintainers/), not on this page. ---- +## Choose your client setup -## Prerequisites +`nflow` only submits Slurm jobs β€” the heavy work runs on the cluster. Two +independent choices decide how you run it: **how you provision the client**, and +**how it reaches Slurm**. + +| You run `nflow` on… | Provision the client | Reach Slurm | Guide | +|---|---|---|---| +| Cluster login/dev node (internet) | `uv sync` | direct | [README](README.md#-installation) | +| Laptop / dev box (internet) | `uv sync` | SSH tunnel | [remote-launch](docs/remote-launch.md) | +| Anywhere, airgapped / no install | `nvflow-client` container | direct or SSH tunnel | [remote-launch](docs/remote-launch.md) | + +Steps 1–6 below are the **cluster side** (stage images/models, write +`my_cluster.yaml`, verify) and apply to every row above. -> **Note:** This guide assumes you've already completed the [README.md](README.md) setup (installed `uv`, cloned the repo, ran `uv sync`). +## πŸ“‹ What you'll do -### Build Host Requirements (for building container images) +Six steps, top to bottom. Each step below opens with a **Goal** and ends with a **βœ… Done when** check so you always know where you are. -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. +| Step | What it does | Who needs it | +|------|--------------|--------------| +| 1. [Prerequisites](#prerequisites) | Confirm cluster access + required tools | Everyone | +| 2. [Setup Containers](#setup-containers) | Stage the five `.sqsh` images on the cluster | Everyone | +| 3. [Download Models](#download-models) | Pre-stage the HF models your workflows use | Everyone | +| 4. [GRPO Prerequisites](#grpo-prerequisites) | SEC cache prefetch | **GRPO only β€” else skip** | +| 5. [Configure Your Cluster](#configure-your-cluster) | Write `cluster_configs/my_cluster.yaml` | Everyone | +| 6. [Verify Installation](#verify-installation) | Sanity-check the whole setup | Everyone | -- **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`) +> **Shortcut:** if a maintainer already staged the `.sqsh` images and models for you, you only need Steps 1, 5, and 6. + +--- + +## Prerequisites -> **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`. +**Step 1 of 6 Β· Goal:** confirm you can reach the cluster and have the tools the setup needs. + +> **Note:** This guide assumes you've already installed the client (see the [README](README.md#-installation): install `uv`, clone the repo, run `uv sync`). +> +> πŸ”Œ **Airgapped / no internet on the install host?** Skip the local install and drive `nflow` from the prebuilt `nvflow-client` container (CLI + venv baked in, no `uv sync`, no client internet) β€” see **[docs/remote-launch.md](docs/remote-launch.md)**. You still stage the worker images and models on the cluster (Steps 2–3 below); only the install differs. ### Cluster Setup Requirements @@ -39,7 +62,9 @@ The `docker build` step needs **internet access** to pull base layers, source fr > **Note:** If you already have `.sqsh` container images staged on the cluster, skip to [Configure Your Cluster](#configure-your-cluster). -**yq (YAML parser):** +
+Install yq (only needed for the parallel conversion script) + ```bash # Check if installed yq --version @@ -59,6 +84,7 @@ chmod +x $HOME/bin/yq echo 'export PATH="$HOME/bin:$PATH"' >> $HOME/.bashrc source $HOME/.bashrc ``` +
**curl & enroot:** ```bash @@ -72,183 +98,31 @@ enroot version # Run on cluster node - 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 +**βœ… Done when:** `enroot version` works on a cluster node and you know your Slurm account, a partition, and your storage paths. + --- ## Setup Containers -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 (5):** - -| Container | Source | Tested Version | Action | -|-----------|--------|----------------|--------| -| `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): - -| Container | Source | Action | -|-----------|--------|--------| -| `megatron` | NeMo-Skills Dockerfiles | Build | -| `sandbox` | NeMo-Skills Dockerfiles | Build | -| `verl` | NeMo-Skills Dockerfiles | Build | -| `trtllm` | `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc8` | Pull from NGC | - -### Step 1: Build Docker Images - -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 -cd /path/to/nvflow - -# 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 . - -# sglang is pulled as-is, no custom Dockerfile -docker pull lmsysorg/sglang:v0.5.10.post1 -``` - -> **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`. - -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 optional containers (`megatron`, `sandbox`, `verl`), build them from the upstream [NeMo-Skills Dockerfiles](https://github.com/NVIDIA-NeMo/Skills/tree/022904023ad7a83a87662a313cf72e7df5891d55/dockerfiles). - -### Step 1b: Sanity-Check Images Before Conversion - -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 2: Push Images to a Registry (Option A) or Save as Tarball (Option B) - -`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. - -#### Option A: Via a private container registry (recommended) - -Push the built images to a registry accessible from your cluster (Docker Hub, NGC, or a private registry): - -```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 - -# sglang can be pulled directly by enroot (no push needed unless your -# cluster cannot reach Docker Hub). -``` - -> **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). - -#### Option B: Via a saved tarball (no registry required) - -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://`: - -```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 - -# Transfer the .tar.gz files to the cluster (scp / rsync / sneakernet) -``` - -`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: Update Container Config - -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 `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-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 - -#### Option A: Automated Setup (Recommended, for Option A registries) +**Step 2 of 6 Β· Goal:** have the five `.sqsh` container images staged on your cluster, with their paths in hand. -Use the setup script to download from your registry and convert all containers in parallel. Pass your personal config with `--config`: +NVFlow runs its cluster jobs inside five `.sqsh` container images. As an operator you only need the `.sqsh` files **staged on your cluster** and their paths recorded in your cluster config. -```bash -# Run from a cluster login node (sbatch requires Slurm access) -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`). - -**Check progress:** -```bash -tail -f outputs/logs/slurm-containers-.out -``` +- **Already have `.sqsh` files staged** (by a maintainer or a previous setup)? Note their paths and skip to [Configure Your Cluster](#configure-your-cluster). +- **Need to build / convert them yourself?** See **[docs/maintainers/containers.md](docs/maintainers/containers.md)** β€” build host requirements, `docker build`, push/save, and `enroot import` to `.sqsh`. -#### Option B: Manual Conversion +Your cluster config references the images by fixed keys β€” `nemo-rl`, `nemo-skills`, `vllm`, `vllm-grpo`, and `sglang`. The build guide's "Update Container Config" step explains how to set them; [Configure Your Cluster](#configure-your-cluster) ties them into your run config. -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 -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 -``` +> **`nemo-gym` (needed for GRPO & DG-SDG):** the Gym-only stages β€” GRPO `prepare_data` / `prefetch_cache` and the DG-SDG gym stages β€” run in a dedicated **CPU-only** `nemo-gym` image ([`dockerfiles/Dockerfile.nemo-gym`](docs/maintainers/containers.md#gym-worker-cpu-only)). Stage this sixth image if you run **GRPO or DG-SDG**; **SFT-only** and **eval-only** runs don't need it. -Move the resulting `.sqsh` files to your cluster's container storage path. +**βœ… Done when:** five `.sqsh` files exist on the cluster and you have their absolute paths for the config. --- ## Download Models +**Step 3 of 6 Β· Goal:** pre-download the models your chosen workflows need to the cluster's HF models directory. + > ⚠️ **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:** @@ -314,7 +188,8 @@ stage_kwargs: server_type: sglang ``` -**Models needed per workflow:** +
+Which models does each workflow need? | Model | Demo SDG | Demo SFT | Demo GRPO | Demo Eval | Production GRPO | |-------|:--------:|:--------:|:---------:|:---------:|:---------------:| @@ -325,6 +200,7 @@ stage_kwargs: | `Qwen/Qwen3-30B-A3B` | | | | | βœ“ | **Tip:** Download commonly used models once and reuse across all workflows. +
### One-Time Connected-Node Stages (Datasets) @@ -337,46 +213,21 @@ A handful of stages legitimately need internet on **first** run to pull benchmar | `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. +For these stages, **temporarily clear** the three HF offline flags (`HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE`, `TRANSFORMERS_OFFLINE`) in your cluster config. `UV_OFFLINE` is unrelated β€” leave it at its default (unset); none of these stages invoke `uv`. > **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 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. - -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. - -### What the self-sufficient image already contains - -`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: +**βœ… Done when:** the models for your workflow are on disk under your mounted `hf_models` directory. -- 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` - -GRPO stages call `installation_command: source /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate` and find everything they need inside the image. - -### Do NOT overlay-mount source over the image paths - -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. +--- -The older dev-mode overlay snippets in `template-slurm.yaml` are commented out for exactly this reason: +## GRPO Prerequisites -```yaml -mounts: - # DO NOT use these with the self-sufficient image β€” they shadow the baked .venv - # - :/opt/NeMo-RL - # - :/opt/NeMo-RL/3rdparty/Gym-workspace/Gym -``` +**Step 4 of 6 Β· GRPO only β€” skip this entire step if you're not running GRPO.** -### Dev mode: iterating on NeMo-RL / NeMo-Gym source +**There are no sources to clone.** The `nvflow-nemo-rl` trainer image bakes NeMo-RL together with one prebuilt NeMo-Gym venv per component, so GRPO `training` resolves no packages at job runtime and needs no bind-mount. The Gym-only stages (`collect_rollouts`, `compute_rewards`, `prefetch_cache`, `prepare_data`) run on the equally self-contained `nvflow-nemo-gym` image. -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. +> How that image is built, and its internals, are covered in **[docs/development/nemo-rl-gym.md](docs/development/nemo-rl-gym.md)**. ### Prefetch SEC Filings Cache (for `finance_sec_search`) @@ -384,28 +235,32 @@ If using the `finance_sec_search` NeMo-Gym environment, you must prefetch the SE 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. +**βœ… Done when:** (GRPO users) the SEC filings cache is prefetched if you use `finance_sec_search`. Everyone else: nothing to do β€” move on. + --- ## Configure Your Cluster -### Step 1: Create Your Cluster Config +**Step 5 of 6 Β· Goal:** create and fill in `cluster_configs/my_cluster.yaml`. + +### Create Your Cluster Config ```bash # Copy template cp cluster_configs/template-slurm.yaml cluster_configs/my_cluster.yaml ``` -### Step 2: Edit Your Config +### Edit Your Config 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 (at minimum `:/hf_models` and `:/workspace`) +4. **Mount points** - Map your cluster paths to container paths (at minimum `:/hf_models` and a writable data dir `:/workspace` for outputs + caches). Recipe code and assets ship via the nemo-run packaged snapshot (`/nemo_run/code`), so the repo is not mounted β€” see [Mount Points](docs/cluster-configuration.md#mounts). 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 +### Keep the Air-Gap Enforcement Block `template-slurm.yaml` ships with the offline flags pre-populated - leave them on: @@ -415,16 +270,15 @@ env_vars: - HF_HUB_OFFLINE=1 - HF_DATASETS_OFFLINE=1 - TRANSFORMERS_OFFLINE=1 - - UV_OFFLINE=true + # UV_OFFLINE: left UNSET (global flag). The images bake every venv they need, + # so no stage resolves packages at runtime either way. + # - 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. @@ -433,10 +287,14 @@ For the one-time connected-node stages listed in [Download Models](#one-time-con > > πŸ“– **For detailed documentation of all configuration fields, see the [Cluster Configuration Guide](docs/cluster-configuration.md)**. +**βœ… Done when:** `my_cluster.yaml` has no remaining `` values and keeps the air-gap enforcement block. + --- ## Verify Installation +**Step 6 of 6 Β· Goal:** confirm the whole setup before running a real workflow. + ```bash # 1. Test NeMo-Skills import uv run python -c "from nemo_skills.pipeline.cli import generate; print('βœ… OK')" @@ -454,6 +312,8 @@ uv run python -c "from omegaconf import OmegaConf; OmegaConf.load('cluster_confi uv run nflow list-stages ``` +**βœ… Done when:** all five checks above pass. You're ready to run a workflow β€” see [Next Steps](#next-steps). + --- ## Troubleshooting @@ -471,20 +331,7 @@ uv sync --reinstall ``` ### yq not found -```bash -# macOS -brew install yq - -# Linux (auto-detects architecture) -mkdir -p $HOME/bin -ARCH=$(uname -m); case "$ARCH" in x86_64) ARCH=amd64 ;; aarch64) ARCH=arm64 ;; armv7l) ARCH=arm ;; i686) ARCH=386 ;; esac -wget "https://github.com/mikefarah/yq/releases/latest/download/yq_linux_${ARCH}" -O $HOME/bin/yq -chmod +x $HOME/bin/yq - -# Add to PATH if needed -echo 'export PATH="$HOME/bin:$PATH"' >> $HOME/.bashrc -source $HOME/.bashrc -``` +`yq` is only needed for the parallel `.sqsh` conversion script. Install it per the [Prerequisites β†’ Install yq](#prerequisites) block above. ### enroot not available ```bash @@ -493,11 +340,8 @@ 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. +### Container build / conversion issues +Building images, `docker login nvcr.io`, and `enroot import` quirks (filename colon, `#` separator for `nvcr.io`) are covered in **[docs/maintainers/containers.md](docs/maintainers/containers.md)**. ### SSH connection failed ```bash @@ -521,29 +365,7 @@ For symptoms specific to the self-sufficient runtime - GRPO `installation_comman ### Ray Cluster Initialization Hangs -**Problem:** Ray cluster hangs during initialization, workers fail to connect - -**Symptoms:** -- Training jobs hang after "Starting Ray cluster" -- Error: `execve(): bad interpreter: No such file or directory` -- Ray workers show connection failures in logs - -**Cause:** SLURM container reattachment issues (SLURM 25.x may be affected) - -**Solution:** -Add `ray_template` to your cluster config: - -```yaml -# In cluster_configs/my_cluster.yaml -executor: slurm -ray_template: "ray_enroot.sub.j2" # Fixes Ray cluster initialization -``` - -**Check your SLURM version:** -```bash -scontrol show config | grep SLURM_VERSION -# Confirmed: SLURM 25.11.2 needs this fix, SLURM 24.x works without it -``` +If training jobs hang after "Starting Ray cluster" (or you see `execve(): bad interpreter: No such file or directory`), your Slurm likely needs the enroot Ray template: set `ray_template: "ray_enroot.sub.j2"` in `my_cluster.yaml`. Full symptoms, cause, and SLURM-version notes are in [cluster-configuration.md β†’ Ray Cluster Configuration](docs/cluster-configuration.md#ray-cluster-configuration). --- @@ -563,11 +385,12 @@ Then head back to the [README.md](README.md#-quick-start) Quick Start section to ## Reference +- **Build the containers**: [`docs/maintainers/containers.md`](docs/maintainers/containers.md) +- **NeMo-RL / NeMo-Gym trainer image & Gym venvs**: [`docs/development/nemo-rl-gym.md`](docs/development/nemo-rl-gym.md) - **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-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 diff --git a/LICENSES/THIRD_PARTY_SW_LICENSE_INFO.md b/LICENSES/THIRD_PARTY_SW_LICENSE_INFO.md index d3a6be2..8e13e68 100644 --- a/LICENSES/THIRD_PARTY_SW_LICENSE_INFO.md +++ b/LICENSES/THIRD_PARTY_SW_LICENSE_INFO.md @@ -13,22 +13,22 @@ This document lists third-party open source and other software packages used in | wrapt | 2.1.0 | BSD | [License](https://raw.githubusercontent.com/GrahamDumpleton/wrapt/develop/LICENSE) | | websockets | 15.0.1 | BSD | [License](https://github.com/python-websockets/websockets/blob/main/LICENSE) | | wcwidth | 0.5.3 | MIT | [License](https://raw.githubusercontent.com/jquast/wcwidth/master/LICENSE) | -| wandb | 0.24.1 | MIT | [License](https://github.com/wandb/wandb/blob/main/LICENSE) | +| wandb | 0.28.1 | MIT | [License](https://github.com/wandb/wandb/blob/main/LICENSE) | | virtualenv | 20.36.1 | MIT | [License](https://raw.githubusercontent.com/pypa/virtualenv/main/LICENSE) | -| urllib3 | 1.26.20 | MIT | [License](https://github.com/urllib3/urllib3/blob/main/LICENSE.txt) | +| urllib3 | 2.7.0 | MIT | [License](https://github.com/urllib3/urllib3/blob/main/LICENSE.txt) | | typing-inspection | 0.4.2 | MIT | [License](https://github.com/pydantic/typing-inspection/blob/main/LICENSE) | | typing-inspect | 0.9.0 | MIT | [License](https://raw.githubusercontent.com/ilevkivskyi/typing_inspect/refs/heads/master/LICENSE) | | types-PyYAML | 6.0.12.20250915 | Apache 2.0 | [License](https://pypi.org/project/types-PyYAML/) | | typer-slim | 0.21.1 | MIT | [License](https://raw.githubusercontent.com/fastapi/typer/master/LICENSE) | -| typer | 0.19.2 | MIT | [License](https://raw.githubusercontent.com/fastapi/typer/master/LICENSE) | -| triton | 3.6.0 | MIT | [License](https://github.com/triton-lang/triton/blob/main/LICENSE) | +| typer | 0.21.1 | MIT | [License](https://raw.githubusercontent.com/fastapi/typer/master/LICENSE) | +| triton | 3.7.1 | MIT | [License](https://github.com/triton-lang/triton/blob/main/LICENSE) | | tree-sitter-yaml | 0.7.2 | MIT | [License](https://github.com/ikatyang/tree-sitter-yaml/blob/master/LICENSE) | | tree-sitter-language-pack | 0.13.0 | Apache 2.0 | [License](https://github.com/Goldziher/tree-sitter-language-pack/blob/main/LICENSE) | | tree-sitter-language-pack | 0.13.0 | MIT | [License](https://github.com/Goldziher/tree-sitter-language-pack/blob/main/LICENSE) | | tree-sitter-embedded-template | 0.25.0 | MIT | [License](https://raw.githubusercontent.com/tree-sitter/tree-sitter-embedded-template/master/LICENSE) | | tree-sitter-c-sharp | 0.23.1 | MIT | [License](https://github.com/tree-sitter/tree-sitter-c-sharp/blob/master/LICENSE) | | tree-sitter | 0.25.2 | MIT | [License](https://raw.githubusercontent.com/tree-sitter/py-tree-sitter/master/LICENSE) | -| transformers | 4.57.1 | Apache 2.0 | [License](https://github.com/huggingface/transformers/blob/main/LICENSE) | +| transformers | 5.13.1 | Apache 2.0 | [License](https://github.com/huggingface/transformers/blob/main/LICENSE) | | tqdm | 4.67.2 | MIT and MPL | [License](https://raw.githubusercontent.com/tqdm/tqdm/master/LICENCE) | | torchx | 0.7.0 | BSD | [License](https://github.com/meta-pytorch/torchx/blob/main/LICENSE) | | tomlkit | 0.13.3 | MIT | [License](https://github.com/sdispater/tomlkit/blob/master/LICENSE) | @@ -39,7 +39,7 @@ This document lists third-party open source and other software packages used in | tabulate | 0.9.0 | MIT | [License](https://github.com/astanin/python-tabulate/blob/master/LICENSE) | | stack-data | 0.6.3 | MIT | [License](https://pypi.org/project/stack-data/) | | sse-starlette | 3.2.0 | BSD | [License](https://github.com/sysid/sse-starlette/blob/main/LICENSE) | -| soupsieve | 2.8.3 | MIT | [License](https://github.com/facelessuser/soupsieve/blob/main/LICENSE.md) | +| soupsieve | 2.9.1 | MIT | [License](https://github.com/facelessuser/soupsieve/blob/main/LICENSE.md) | | sniffio | 1.3.1 | Apache 2.0 | [License](https://github.com/python-trio/sniffio/blob/master/LICENSE) | | smmap | 5.0.2 | BSD | [License](https://github.com/gitpython-developers/smmap/blob/main/LICENSE) | | smart-open | 7.5.0 | MIT | [License](https://github.com/getcrest/smart-open/blob/dev/LICENSE) | @@ -49,16 +49,15 @@ This document lists third-party open source and other software packages used in | sentry-sdk | 2.51.0 | MIT | [License](https://github.com/getsentry/sentry-python/blob/master/LICENSE) | | sentence-transformers | 5.2.2 | Apache 2.0 | [License](https://github.com/huggingface/sentence-transformers/blob/main/LICENSE) | | scikit-learn | 1.8.0 | BSD | [License](https://github.com/scikit-learn/scikit-learn/blob/main/COPYING) | -| safetensors | 0.7.0 | Apache 2.0 | [License](https://github.com/huggingface/safetensors/blob/main/LICENSE) | +| safetensors | 0.8.0 | Apache 2.0 | [License](https://github.com/huggingface/safetensors/blob/main/LICENSE) | | safehttpx | 0.1.7 | MIT | [License](https://github.com/gradio-app/safehttpx/blob/main/LICENSE) | | sacrebleu | 2.6.0 | Apache 2.0 | [License](https://github.com/mjpost/sacrebleu/blob/master/LICENSE.txt) | | ruff | 0.14.14 | MIT | [License](https://github.com/astral-sh/ruff/blob/main/LICENSE) | | rpds-py | 0.30.0 | MIT | [License](https://raw.githubusercontent.com/crate-py/rpds/main/LICENSE) | | regex | 2026.1.15 | Apache 2.0 | [License](https://spdx.org/licenses/Apache-2.0.html) | -| regex | 2025.9.18 | Apache 2.0 | [License](https://spdx.org/licenses/Apache-2.0.html) | | referencing | 0.37.0 | MIT | [License](https://github.com/python-jsonschema/referencing/blob/main/COPYING) | | rank-bm25 | 0.2.2 | Apache 2.0 | [License](https://raw.githubusercontent.com/dorianbrown/rank_bm25/master/LICENSE) | -| pytorch | 2.10.0 | BSD | [License](https://github.com/pytorch/pytorch/blob/main/LICENSE) | +| pytorch | 2.13.0 | BSD | [License](https://github.com/pytorch/pytorch/blob/main/LICENSE) | | pythonjedi | 0.19.2 | MIT | [License](https://github.com/davidhalter/jedi/blob/master/LICENSE.txt) | | python3-rich | 14.3.2 | MIT | [License](https://github.com/Textualize/rich/blob/master/LICENSE) | | python3-charset-normalizer | 3.4.4 | MIT | [License](https://github.com/jawah/charset_normalizer/blob/master/LICENSE) | @@ -66,10 +65,10 @@ This document lists third-party open source and other software packages used in | python-semanticversion | 2.10.0 | BSD | [License](https://pypi.org/project/semantic-version/) | | python-protobuf | 6.33.5 | BSD | [License](https://raw.githubusercontent.com/protocolbuffers/protobuf/main/LICENSE) | | python-pluggy | 1.6.0 | MIT | [License](https://pypi.org/project/pluggy/) | -| python-multipart | 0.0.22 | Apache 2.0 | [License](https://pypi.org/project/python-multipart/) | -| python-jsonschema | 4.26.0 | MIT | [License](https://github.com/python-jsonschema/jsonschema/blob/main/LICENSE) | +| python-multipart | 0.0.32 | Apache 2.0 | [License](https://pypi.org/project/python-multipart/) | +| python-jsonschema | 4.23.0 | MIT | [License](https://github.com/python-jsonschema/jsonschema/blob/main/LICENSE) | | python-hpack | 4.1.0 | MIT | [License](https://raw.githubusercontent.com/python-hyper/hpack/master/LICENSE) | -| python-dotenv | 1.2.1 | BSD | [License](https://github.com/theskumar/python-dotenv/blob/main/LICENSE) | +| python-dotenv | 1.2.2 | BSD | [License](https://github.com/theskumar/python-dotenv/blob/main/LICENSE) | | python-distlib | 0.4.0 | Python Software Foundation License 2.0 | [License](https://github.com/pypa/distlib/blob/master/LICENSE.txt) | | python-dill | 0.3.8 | BSD | [License](https://github.com/uqfoundation/dill/blob/master/LICENSE) | | python-decorator | 5.2.1 | BSD | [License](https://github.com/micheles/decorator/blob/master/LICENSE) | @@ -82,7 +81,7 @@ This document lists third-party open source and other software packages used in | pytest-xdist | 3.8.0 | MIT | [License](https://github.com/pytest-dev/pytest-xdist/blob/master/LICENSE) | | pytest-timeout | 2.4.0 | MIT | [License](https://github.com/pytest-dev/pytest-timeout/blob/main/LICENSE) | | pytest-cov | 7.0.0 | MIT | [License](https://github.com/pytest-dev/pytest-cov/blob/master/LICENSE) | -| pytest | 9.0.2 | MIT | [License](https://github.com/pytest-dev/pytest/blob/main/LICENSE) | +| pytest | 9.0.3 | MIT | [License](https://github.com/pytest-dev/pytest/blob/main/LICENSE) | | pypi/setuptools | 80.10.2 | MIT | [License](https://github.com/pypa/setuptools/blob/main/LICENSE) | | pypa/sampleproject | 0.0.32 | MIT | [License](https://github.com/pypa/sampleproject/blob/main/LICENSE.txt) | | pyinvoke | 2.2.1 | BSD | [License](https://github.com/pyinvoke/invoke/blob/main/LICENSE) | @@ -92,11 +91,11 @@ This document lists third-party open source and other software packages used in | pydantic | 2.12.5 | MIT | [License](https://github.com/pydantic/pydantic/blob/main/LICENSE) | | pycparser | 3.0 | BSD | [License](https://github.com/eliben/pycparser/blob/master/LICENSE) | | pyca/pynacl | 1.6.2 | Apache 2.0 | [License](https://raw.githubusercontent.com/pyca/pynacl/main/LICENSE) | -| pyca/cryptography | 42.0.8 | BSD | [License](https://raw.githubusercontent.com/pyca/cryptography/main/LICENSE) | -| pyca/cryptography | 42.0.8 | Apache 2.0 | [License](https://raw.githubusercontent.com/pyca/cryptography/main/LICENSE) | +| pyca/cryptography | 49.0.0 | BSD | [License](https://raw.githubusercontent.com/pyca/cryptography/main/LICENSE) | +| pyca/cryptography | 49.0.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/pyca/cryptography/main/LICENSE) | | pyasn1-modules | 0.4.2 | BSD | [License](https://raw.githubusercontent.com/pyasn1/pyasn1-modules/master/LICENSE.txt) | -| pyasn1 | 0.6.2 | BSD | [License](https://raw.githubusercontent.com/pyasn1/pyasn1/master/LICENSE.txt) | -| pyarrow | 23.0.0 | Apache 2.0 | [License](https://github.com/apache/arrow/blob/main/LICENSE) | +| pyasn1 | 0.6.4 | BSD | [License](https://raw.githubusercontent.com/pyasn1/pyasn1/master/LICENSE.txt) | +| pyarrow | 25.0.0 | Apache 2.0 | [License](https://github.com/apache/arrow/blob/main/LICENSE) | | py3-google-auth | 2.48.0 | Apache 2.0 | [License](https://github.com/googleapis/google-auth-library-python/blob/main/LICENSE) | | py-spy | 0.4.1 | MIT | [License](https://raw.githubusercontent.com/benfred/py-spy/master/LICENSE) | | py-filelock | 3.20.3 | MIT | [License](https://py-filelock.readthedocs.io/en/latest/license.html) | @@ -108,7 +107,7 @@ This document lists third-party open source and other software packages used in | propcache | 0.4.1 | Apache 2.0 | [License](https://github.com/aio-libs/propcache/blob/master/LICENSE) | | prompt-toolkit/python-prompt-toolkit | 3.0.52 | BSD | [License](https://github.com/prompt-toolkit/python-prompt-toolkit/blob/main/LICENSE) | | prometheus-fastapi-instrumentator | 7.0.0 | ISC (Internet Software Consortium) | [License](https://raw.githubusercontent.com/trallnag/prometheus-fastapi-instrumentator/master/LICENSE) | -| projectray | 2.53.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/ray-project/ray/master/LICENSE) | +| projectray | 2.56.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/ray-project/ray/master/LICENSE) | | pre-commit/pre-commit | 4.5.1 | MIT | [License](https://github.com/pre-commit/pre-commit/blob/main/LICENSE) | | portalocker | 3.2.0 | BSD | [License](https://github.com/wolph/portalocker/blob/develop/LICENSE) | | platformdirs | 4.5.1 | MIT | [License](https://raw.githubusercontent.com/tox-dev/platformdirs/main/LICENSE) | @@ -118,13 +117,13 @@ This document lists third-party open source and other software packages used in | parso | 0.8.5 | MIT | [License](https://github.com/davidhalter/parso/blob/main/LICENSE) | | paramiko | 4.0.0 | LGPL (Library or Lesser GPL) | [License](https://github.com/paramiko/paramiko/blob/main/LICENSE) | | pandas-python | 3.0.0 | BSD | [License](https://github.com/pandas-dev/pandas/blob/main/LICENSE) | -| orjson | 3.11.3 | Apache 2.0 | [License](https://github.com/ijl/orjson/blob/master/LICENSE-APACHE,https://github.com/ijl/orjson/blob/master/LICENSE-MIT) | +| orjson | 3.11.7 | Apache 2.0 | [License](https://github.com/ijl/orjson/blob/master/LICENSE-APACHE,https://github.com/ijl/orjson/blob/master/LICENSE-MIT) | | opentelemetry-semantic-conventions | 0.60b1 | Apache 2.0 | [License](https://raw.githubusercontent.com/open-telemetry/opentelemetry-python/main/LICENSE) | | opentelemetry-sdk | 1.39.1 | Apache 2.0 | [License](https://raw.githubusercontent.com/open-telemetry/opentelemetry-python/refs/heads/main/LICENSE) | | opentelemetry-python | 1.39.1 | Apache 2.0 | [License](https://github.com/open-telemetry/opentelemetry-python/blob/main/LICENSE) | | opentelemetry packages | 0.60b1 | Apache 2.0 | [License](https://github.com/open-telemetry/opentelemetry-python/blob/main/LICENSE) | | opencensus | 0.11.4 | Apache 2.0 | [License](https://github.com/census-instrumentation/opencensus-python/blob/master/LICENSE) | -| openai | 2.16.0 | Apache 2.0 | [License](https://github.com/openai/openai-python/blob/main/LICENSE) | +| openai | 2.24.0 | Apache 2.0 | [License](https://github.com/openai/openai-python/blob/main/LICENSE) | | omegaconf | 2.3.0 | BSD | [License](https://github.com/omry/omegaconf/blob/master/LICENSE) | | nvidia-nvtx-cu12 | 12.8.90 | Apache 2.0 | [License](https://pypi.org/project/nvidia-nvtx-cu12/) | | nvidia-nvshmem-cu12 | 3.4.5 | Nvidia Proprietary License | [License](https://raw.githubusercontent.com/NVIDIA/nvshmem/devel/License.txt) | @@ -142,24 +141,24 @@ This document lists third-party open source and other software packages used in | nvidia-cuda-cupti-cu12 | 12.8.90 | Nvidia Proprietary License | [License](https://docs.nvidia.com/cuda/eula/index.html) | | nvidia-cublas-cu12 | 12.8.4.1 | Nvidia Proprietary License | [License](https://docs.nvidia.com/cuda/eula/index.html) | | nodeenv | 1.10.0 | BSD | [License](https://github.com/ekalinin/nodeenv/blob/master/LICENSE) | -| nemo-run | 0.7.0 | Apache 2.0 | [License](https://github.com/NVIDIA/NeMo-Run/blob/main/LICENSE) | +| nemo-run | 0.9.0rc0.dev0 | Apache 2.0 | [License](https://github.com/NVIDIA/NeMo-Run/blob/main/LICENSE) | | nemo-evaluator-launcher | 0.1.46 | Apache 2.0 | [License](https://pypi.org/project/cuda-bindings/) | | mypy-extensions | 1.1.0 | MIT | [License](https://github.com/python/mypy_extensions/blob/master/LICENSE) | | mypy | 1.19.1 | MIT | [License](https://pypi.org/project/mypy/) | | multipledispatch | 1.0.0 | BSD | [License](https://github.com/mrocklin/multipledispatch/blob/master/LICENSE.txt) | | multidict | 6.7.1 | Apache 2.0 | [License](https://raw.githubusercontent.com/aio-libs/multidict/master/LICENSE) | | multi_process | 0.70.16 | BSD | [License](https://github.com/uqfoundation/multiprocess/blob/master/LICENSE) | -| msgpack | 1.1.2 | Apache 2.0 | [License](https://raw.githubusercontent.com/msgpack/msgpack-python/main/COPYING) | +| msgpack | 1.2.1 | Apache 2.0 | [License](https://raw.githubusercontent.com/msgpack/msgpack-python/main/COPYING) | | mpmath | 1.3.0 | BSD | [License](https://github.com/fredrik-johansson/mpmath/blob/master/LICENSE) | | mdurl | 0.1.2 | MIT | [License](https://raw.githubusercontent.com/executablebooks/mdurl/master/LICENSE) | -| mcp | 1.26.0 | MIT | [License](https://github.com/modelcontextprotocol/python-sdk/blob/main/LICENSE) | +| mcp | 1.28.1 | MIT | [License](https://github.com/modelcontextprotocol/python-sdk/blob/main/LICENSE) | | matplotlib-inline | 0.2.1 | BSD | [License](https://github.com/ipython/matplotlib-inline/blob/master/LICENSE) | | math-verify | 0.9.0 | Apache 2.0 | [License](https://github.com/huggingface/math-verify/blob/main/LICENCE) | | markdown-it-py | 4.0.0 | MIT | [License](https://github.com/executablebooks/markdown-it-py/blob/master/LICENSE) | | marisa-trie | 1.3.1 | MIT | [License](https://github.com/pytries/marisa-trie/blob/master/LICENSE) | -| lxml | 6.0.2 | BSD | [License](https://github.com/lxml/lxml/blob/main/LICENSE) | +| lxml | 6.1.1 | BSD | [License](https://github.com/lxml/lxml/blob/main/LICENSE) | | loguru | 0.7.3 | MIT | [License](https://github.com/Delgan/loguru/blob/master/LICENSE) | -| litellm | 1.81.6 | MIT | [License](https://github.com/BerriAI/litellm/blob/main/LICENSE) | +| litellm | 1.84.10 | MIT | [License](https://github.com/BerriAI/litellm/blob/main/LICENSE) | | librt | 0.7.8 | MIT | [License](https://github.com/mypyc/librt/blob/main/LICENSE) | | libcst | 1.8.6 | MIT | [License](https://github.com/Instagram/LibCST/blob/main/LICENSE) | | leptonai | 0.27.0 | Apache 2.0 | [License](https://sourceforge.net/projects/lepton-ai.mirror/) | @@ -176,22 +175,22 @@ This document lists third-party open source and other software packages used in | iso639-lang | 2.6.3 | MIT | [License](https://pypi.org/project/iso639-lang/) | | ipython-pygments-lexers | 1.1.1 | BSD | [License](https://github.com/ipython/ipython-pygments-lexers/blob/main/LICENSE) | | inquirerpy | 0.3.4 | MIT | [License](https://github.com/kazhala/InquirerPy/blob/master/LICENSE) | -| importlib_metadata | 8.7.1 | Apache 2.0 | [License](https://pypi.org/project/importlib-metadata/) | +| importlib_metadata | 8.5.0 | Apache 2.0 | [License](https://pypi.org/project/importlib-metadata/) | | idna | 3.11 | BSD | [License](https://github.com/kjd/idna/blob/master/LICENSE.md) | | identify | 2.6.16 | MIT | [License](https://github.com/pre-commit/identify/blob/main/LICENSE) | | hyperframe | 6.1.0 | MIT | [License](https://github.com/python-hyper/hyperframe/blob/master/LICENSE) | | hydra-core | 1.3.2 | MIT | [License](https://github.com/facebookresearch/hydra/blob/main/LICENSE) | -| huggingface-hub | 1.3.7 | Apache 2.0 | [License](https://github.com/huggingface/huggingface_hub/blob/main/LICENSE) | +| huggingface-hub | 1.16.1 | Apache 2.0 | [License](https://github.com/huggingface/huggingface_hub/blob/main/LICENSE) | | httpx-sse | 0.4.3 | MIT | [License](https://raw.githubusercontent.com/florimondmanca/httpx-sse/master/LICENSE) | -| httpx | 0.27.2 | BSD | [License](https://github.com/encode/httpx/blob/master/LICENSE.md) | +| httpx | 0.28.1 | BSD | [License](https://github.com/encode/httpx/blob/master/LICENSE.md) | | httpcore | 1.0.9 | BSD | [License](https://raw.githubusercontent.com/encode/httpcore/master/LICENSE.md) | -| hf-xet | 1.2.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/huggingface/xet-core/main/LICENSE) | +| hf-xet | 1.5.1 | Apache 2.0 | [License](https://raw.githubusercontent.com/huggingface/xet-core/main/LICENSE) | | h2 | 4.3.0 | MIT | [License](https://github.com/python-hyper/h2/blob/master/LICENSE) | | h11 | 0.16.0 | MIT | [License](https://github.com/python-hyper/h11/blob/master/LICENSE.txt) | | grpc | 1.76.0 | Apache 2.0 | [License](https://github.com/grpc/grpc/blob/master/LICENSE) | | groovy | 0.1.2 | MIT | [License](https://github.com/abidlabs/groovy/blob/main/LICENSE) | -| gradio-client | 2.0.3 | Apache 2.0 | [License](https://github.com/gradio-app/gradio/blob/main/LICENSE) | -| gradio | 6.5.1 | Apache 2.0 | [License](https://github.com/gradio-app/gradio/blob/main/LICENSE) | +| gradio-client | 2.5.0 | Apache 2.0 | [License](https://github.com/gradio-app/gradio/blob/main/LICENSE) | +| gradio | 6.20.0 | Apache 2.0 | [License](https://github.com/gradio-app/gradio/blob/main/LICENSE) | | googleapis-common-protos | 1.72.0 | Apache 2.0 | [License](https://github.com/gradio-app/gradio/blob/main/LICENSE) | | google/fiddle | 0.3.0 | Apache 2.0 | [License](https://github.com/gradio-app/gradio/blob/main/LICENSE) | | google-api-core | 2.29.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/googleapis/python-api-common-protos/main/LICENSE) | @@ -210,7 +209,7 @@ This document lists third-party open source and other software packages used in | executing | 2.2.1 | MIT | [License](https://github.com/alexmojaki/executing/blob/master/LICENSE) | | execnet | 2.1.2 | MIT | [License](https://pypi.org/project/execnet/) | | exceptiongroup | 1.3.1 | MIT | [License](https://github.com/agronholm/exceptiongroup/blob/main/LICENSE) | -| evalplus | 0.3.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/evalplus/evalplus/refs/heads/master/LICENSE) | +| evalplus | 0.3.0.dev27 | Apache 2.0 | [License](https://raw.githubusercontent.com/evalplus/evalplus/refs/heads/master/LICENSE) | | encode/uvicorn | 0.40.0 | BSD | [License](https://github.com/encode/uvicorn/blob/master/LICENSE.md) | | edit-distance | 0.8.1 | MIT | [License](https://github.com/roy-ht/editdistance/blob/master/LICENSE) | | docstring-parser | 0.17.0 | MIT | [License](https://raw.githubusercontent.com/rr-/docstring_parser/master/LICENSE.md) | @@ -218,8 +217,8 @@ This document lists third-party open source and other software packages used in | distro | 1.9.0 | Apache 2.0 | [License](https://github.com/python-distro/distro/blob/master/LICENSE) | | diskcache | 5.6.3 | Apache 2.0 | [License](http://www.apache.org/licenses/LICENSE-2.0) | | datasets | 3.6.0 | Apache 2.0 | [License](https://github.com/huggingface/datasets/blob/main/LICENSE) | -| cuda-pathfinder | 1.3.3 | Apache 2.0 | [License](https://www.apache.org/licenses/LICENSE-2.0) | -| cuda-bindings | 12.9.4 | NVIDIA CUDA Toolkit License Agreement | [License](https://pypi.org/project/cuda-bindings/) | +| cuda-pathfinder | 1.5.6 | Apache 2.0 | [License](https://www.apache.org/licenses/LICENSE-2.0) | +| cuda-bindings | 13.3.1 | NVIDIA CUDA Toolkit License Agreement | [License](https://pypi.org/project/cuda-bindings/) | | contextlib2 | 21.6.0 | Apache 2.0 | [License](https://policies.python.org/pypi.org/Acceptable-Use-Policy/#credits-license) | | colorama | 0.4.6 | BSD | [License](https://raw.githubusercontent.com/tartley/colorama/master/LICENSE.txt) | | cloudpickle | 3.0.0 | BSD | [License](https://github.com/cloudpipe/cloudpickle/blob/master/LICENSE) | @@ -239,8 +238,8 @@ This document lists third-party open source and other software packages used in | annotated-doc | 0.0.4 | MIT | [License](https://github.com/fastapi/annotated-doc/blob/main/LICENSE) | | aiosignal | 1.4.0 | Apache 2.0 | [License](https://github.com/aio-libs/aiosignal/blob/master/LICENSE) | | aiohttp-cors | 0.8.1 | Apache 2.0 | [License](https://github.com/aio-libs/aiohttp-cors/blob/master/LICENSE) | -| aiohttp | 3.13.0 | MIT | [License](https://github.com/aio-libs/aiohttp/blob/main/LICENSE) | -| aiohttp | 3.13.0 | Apache 2.0 | [License](https://github.com/aio-libs/aiohttp/blob/main/LICENSE) | +| aiohttp | 3.14.1 | MIT | [License](https://github.com/aio-libs/aiohttp/blob/main/LICENSE) | +| aiohttp | 3.14.1 | Apache 2.0 | [License](https://github.com/aio-libs/aiohttp/blob/main/LICENSE) | | aiohappyeyeballs | 2.6.1 | Python Software Foundation License 2.0 | [License](https://github.com/aio-libs/aiohappyeyeballs/blob/main/LICENSE) | | aiofiles | 24.1.0 | Apache 2.0 | [License](https://github.com/Tinche/aiofiles/blob/main/LICENSE) | | absl-py | 2.4.0 | Apache 2.0 | [License](https://raw.githubusercontent.com/abseil/abseil-py/master/LICENSE) | @@ -249,20 +248,20 @@ This document lists third-party open source and other software packages used in | Traitlets | 5.14.3 | BSD | [License](https://github.com/ipython/traitlets/blob/main/LICENSE) | | TOML | 0.10.2 | MIT | [License](https://github.com/uiri/toml/blob/master/LICENSE) | | SymPy | 1.14.0 | BSD | [License](https://github.com/sympy/sympy/blob/main/LICENSE) | -| Starlette | 0.50.0 | BSD | [License](https://github.com/Kludex/starlette/blob/main/LICENSE.md) | +| Starlette | 1.3.1 | BSD | [License](https://github.com/Kludex/starlette/blob/main/LICENSE.md) | | SciPy | 1.17.0 | BSD | [License](https://github.com/scipy/scipy/blob/main/LICENSE.txt) | | RonnyPfannschmidt/iniconfig | 2.3.0 | MIT | [License](https://github.com/pytest-dev/iniconfig/blob/main/LICENSE) | | Python-RSA | 4.9.1 | Apache 2.0 | [License](https://github.com/sybrenstuvel/python-rsa/blob/main/LICENSE) | -| Python tzdata | 2025.2 | Apache 2.0 | [License](https://github.com/python/tzdata/blob/master/LICENSE) | +| Python tzdata | 2025.3 | Apache 2.0 | [License](https://github.com/python/tzdata/blob/master/LICENSE) | | Python six | 1.17.0 | MIT | [License](https://github.com/benjaminp/six/blob/master/LICENSE) | -| Pygments - Python syntax highlighter | 2.19.2 | BSD | [License](https://github.com/pygments/pygments/blob/master/LICENSE) | +| Pygments - Python syntax highlighter | 2.20.0 | BSD | [License](https://github.com/pygments/pygments/blob/master/LICENSE) | | PyYAML | 6.0.3 | MIT | [License](https://github.com/yaml/pyyaml/blob/main/LICENSE) | | PyTZ - Python Time Zone Library | 2025.2 | MIT | [License](https://github.com/stub42/pytz/blob/master/LICENSE.txt) | -| PyJWT | 2.11.0 | MIT | [License](https://spdx.org/licenses/MIT.html) | -| PillowPython | 12.1.0 | MIT | [License](https://github.com/python-pillow/Pillow/blob/main/LICENSE) | +| PyJWT | 2.13.0 | MIT | [License](https://spdx.org/licenses/MIT.html) | +| PillowPython | 12.3.0 | MIT | [License](https://github.com/python-pillow/Pillow/blob/main/LICENSE) | | Packaging | 26.0 | BSD | [License](https://github.com/pypa/packaging/blob/main/LICENSE) | | Packaging | 26.0 | Apache 2.0 | [License](https://github.com/pypa/packaging/blob/main/LICENSE) | -| OpenCensus | 0.1.3 | Apache 2.0 | [License](https://github.com/census-instrumentation/opencensus-python/blob/master/LICENSE) | +| opencensus-context | 0.1.3 | Apache 2.0 | [License](https://github.com/census-instrumentation/opencensus-python/blob/master/LICENSE) | | NumPy | 2.1.3 | BSD | [License](https://github.com/numpy/numpy/blob/main/LICENSE.txt) | | NetworkX | 3.6.1 | BSD | [License](https://github.com/networkx/networkx/blob/master/LICENSE.txt) | | MarkupSafe | 3.0.3 | BSD | [License](https://github.com/pallets/markupsafe/blob/main/LICENSE.txt) | @@ -271,6 +270,6 @@ This document lists third-party open source and other software packages used in | Flask | 3.1.2 | BSD | [License](https://raw.githubusercontent.com/pallets/flask/main/LICENSE.txt) | | Deprecated | 1.3.1 | MIT | [License](https://raw.githubusercontent.com/laurent-laporte-pro/deprecated/master/LICENSE.rst) | | Coverage | 7.13.2 | Apache 2.0 | [License](https://github.com/nedbat/coveragepy/blob/master/LICENSE.txt) | -| Click - Python Command Line Utility | 8.1.8 | BSD | [License](https://click.palletsprojects.com/en/stable/license/) | +| Click - Python Command Line Utility | 8.4.2 | BSD | [License](https://click.palletsprojects.com/en/stable/license/) | | BeautifulSoup4 | 4.14.3 | MIT | [License](https://www.crummy.com/software/BeautifulSoup/bs4/doc/) | | python-hyper/h2 | 4.3.0 | MIT | [License](https://github.com/python-hyper/h2/blob/master/LICENSE) | diff --git a/NOTICE b/NOTICE index e9e43ba..fdb7efe 100644 --- a/NOTICE +++ b/NOTICE @@ -19,7 +19,7 @@ For full license texts, see the links below or the documentation of each component. - nemo-skills (NVIDIA NeMo-Skills) - https://github.com/NVIDIA/NeMo-Skills + https://github.com/NVIDIA-NeMo/Skills License: Apache-2.0 - omegaconf diff --git a/README.md b/README.md index 2640664..59de1dc 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,8 @@ uv run pre-commit install > ⚠️ **Developers:** Always run `uv run pre-commit install` after cloning. This enables automatic code quality checks on every commit. +> πŸ”Œ **Air-gapped / no internet on the install host?** Skip the local install and drive `nflow` from the prebuilt `nvflow-client` container (CLI + venv baked in, no `uv sync`). See [docs/remote-launch.md](docs/remote-launch.md). + ### Activating the Virtual Environment (Optional) By default, use `uv run ` to run commands in the project's virtual environment. If you prefer to activate the environment directly: @@ -123,11 +125,13 @@ pytest ## πŸ”§ Cluster Setup -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. +To run workflows on a Slurm cluster you need to: (1) build the NVFlow +container images from the Dockerfiles in [`dockerfiles/`](dockerfiles/) +(`nemo-rl`, `nemo-gym`, `nemo-skills`, `vllm`, `vllm-grpo`; only `sglang` is +pulled as-is), (2) convert them to `.sqsh` for Slurm, and (3) write a cluster +config (`cluster_configs/my_cluster.yaml`). Every image bakes the packages and +virtual environments its stages need, so no stage resolves dependencies at job +runtime (see [INSTALL.md](INSTALL.md)). > **See [INSTALL.md](INSTALL.md)** for the complete setup guide > (build, sanity-check, `.sqsh` conversion, model staging, cluster @@ -142,6 +146,9 @@ Once cluster setup is complete, set the config directory: export NEMO_SKILLS_CONFIG_DIR=/path/to/nvflow/cluster_configs ``` +> **Airgapped / no local install?** Drive `nflow` from the `nvflow-client` +> container (no `uv sync` needed) over an SSH tunnel β€” see [docs/remote-launch.md](docs/remote-launch.md). + ## πŸš€ Quick Start ### CLI Invocation @@ -200,16 +207,14 @@ nvflow/ β”‚ β”œβ”€β”€ stages/sdg/ # Example SDG stage β”‚ β”œβ”€β”€ prompts/ # Prompt templates β”‚ └── workflows/ # Example workflows - └── finance/ # Finance reasoning recipe - β”œβ”€β”€ stages/ # Stage implementations - β”‚ β”œβ”€β”€ download/ # SEC filing download - β”‚ β”œβ”€β”€ evaluation/ # Evaluation stages - β”‚ β”œβ”€β”€ rl/ # GRPO RL training stages - β”‚ β”œβ”€β”€ sdg/ # SDG stages - β”‚ β”œβ”€β”€ sft/ # SFT training stages - β”‚ └── shared/ # Shared stages (data_transformation, train_validation_split) - β”œβ”€β”€ prompts/ # Prompt templates - └── workflows/ # Workflow configs + β”œβ”€β”€ finance/ # Finance reasoning recipe + β”‚ β”œβ”€β”€ stages/ # Stage implementations + β”‚ β”œβ”€β”€ prompts/ # Prompt templates + β”‚ └── workflows/ # Workflow configs + └── multimodal/ # Multimodal HopChain recipe + β”œβ”€β”€ stages/ # Image filtering and SDG stages + β”œβ”€β”€ prompts/ # Vision-language prompt templates + └── workflows/ # HopChain workflow configs ``` ## πŸ“ Creating a Stage @@ -261,7 +266,7 @@ End-to-end pipeline for generating synthetic financial Q&A data from SEC filings **Quick Links:** - [Quick Start (~3 hour demo)](docs/recipes/finance/quick-start.md) - Get started quickly with 7 companies - [Workflow Guides](docs/recipes/finance/workflows/) - Detailed guides for all 6 workflows -- [Stage Reference](docs/recipes/finance/stages/) - Technical specifications for 42 stages +- [Stage Reference](docs/recipes/finance/stages/) - Technical specifications for all 37 stages **Pipeline:** ``` @@ -269,12 +274,31 @@ download-sec β†’ template-sdg / document-sdg β†’ sft β†’ eval β†’ grpo ``` **Features:** -- 42 stages across 6 workflows +- 37 stages across 6 workflows (1 + 6 + 7 + 6 + 7 + 10) - Two SDG approaches (template-based & document-grounded) - Multiple model support (GPT-OSS-120B, Qwen3, Nemotron) -- Produces 80K+ synthetic Q&A pairs +- Produces 300K+ synthetic Q&A pairs - Complete training and evaluation pipeline +### Multimodal HopChain Recipe + +**πŸ“š [Complete Multimodal Recipe Documentation β†’](docs/recipes/multimodal/README.md)** + +HopChain-inspired multimodal synthetic data generation for multi-hop +vision-language reasoning. + +**Pipeline:** +``` +image-filter β†’ identify-categories β†’ localize-instances β†’ sample-combinations + β†’ generate-queries β†’ verify β†’ judge/reconcile β†’ difficulty-filter β†’ sft-traces +``` + +**Features:** +- Two workflows: image filtering and SDG +- SAM-backed instance localization +- Structural verification and optional external LLM judges +- Optional SFT reasoning-trace generation and filtering + ## πŸ“š CLI Commands ```bash @@ -314,5 +338,5 @@ Apache-2.0 ## πŸ™ Acknowledgments Built on: -- [NeMo-Skills](https://github.com/NVIDIA/NeMo-Skills) +- [NeMo-Skills](https://github.com/NVIDIA-NeMo/Skills) - [NeMo-RL](https://github.com/NVIDIA-NeMo/RL) diff --git a/cluster_configs/containers.yaml b/cluster_configs/containers.yaml index 14281b9..6110e94 100644 --- a/cluster_configs/containers.yaml +++ b/cluster_configs/containers.yaml @@ -11,46 +11,73 @@ # 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 +# YAML KEYS (nemo-rl, nemo-gym, nemo-skills, 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/nvflow-nemo-skills:0229040) +# - Simple string: image reference (e.g., your-registry/nvflow-nemo-skills:v1.1.2) # - 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) +# Built locally (custom images; all deps pre-installed, no runtime downloads): +# nemo-rl -> dockerfiles/Dockerfile.nemo-rl (base: nvcr.io/nvidia/nemo-rl:v0.7.0) +# nemo-skills -> dockerfiles/Dockerfile.nemo-skills (base: ubuntu:22.04, NeMo-Skills @ e06c9b90) +# vllm -> dockerfiles/Dockerfile.vllm (base: vllm/vllm-openai:v0.22.0) +# vllm-grpo -> dockerfiles/Dockerfile.vllm --build-arg VLLM_VERSION=v0.20.0 +# nemo-gym -> dockerfiles/Dockerfile.nemo-gym (base: python:3.12-slim, CPU-only Gym client) +# Pulled as-is (no custom Dockerfile): +# sglang -> lmsysorg/sglang:v0.5.10.post1 +# +# Optional launcher image for airgapped users (airgap-only; NOT a worker -- see +# the nvflow-client entry below): +# nvflow-client -> dockerfiles/Dockerfile.nvflow (base: digest-pinned ubuntu:24.04; multi-arch amd64+arm64) containers: # --------------------------------------------------------------------------- # Required: Built locally from dockerfiles/, then pushed to your registry # (see INSTALL.md Step 1 and Step 2). # --------------------------------------------------------------------------- - # Tested: nvflow-nemo-rl:v0.6.0 (extends nvcr.io/nvidia/nemo-rl:v0.6.0) - nemo-rl: /nvflow-nemo-rl:v0.6.0 + # SFT/GRPO trainer. Extends the NeMo-RL v0.7.0 base with the NeMo-Gym + # per-component venvs baked in, so no `uv` resolve happens at job runtime. + # dockerfiles/Dockerfile.nemo-rl (upstream Gym main 33ef60369). The stock base + # also works if your compute nodes have internet during training. + nemo-rl: /nvflow-nemo-rl:v0.7.0 + + # CPU-only Gym client for the Gym-only stages (DG-SDG gym stages + GRPO + # prepare_data/prefetch_cache). dockerfiles/Dockerfile.nemo-gym (base: + # python:3.12-slim; upstream Gym main 33ef60369; per-component venvs baked). + nemo-gym: /nvflow-nemo-gym:0.4.0 - # Tested: nvflow-nemo-skills:0229040 (NeMo-Skills @ commit 0229040) - nemo-skills: /nvflow-nemo-skills:0229040 + # Tested: nvflow-nemo-skills:v1.1.2 (NeMo-Skills @ commit e06c9b90) + nemo-skills: /nvflow-nemo-skills:v1.1.2 - # 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:v0.22.0 (extends vllm/vllm-openai:v0.22.0; + # pre-cached tiktoken + openai_harmony; multi-arch amd64 + arm64) + vllm: /nvflow-vllm:v0.22.0 - # 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 + # Tested: nvflow-vllm:v0.20.0 (extends vllm/vllm-openai:v0.20.0; pinned to + # match NeMo-RL v0.7.0 colocated vLLM). Same repo as `vllm`, different tag. + vllm-grpo: /nvflow-vllm:v0.20.0 # --------------------------------------------------------------------------- - # Required: Pulled as-is from Docker Hub (no custom Dockerfile) + # Required: Pulled as-is (no custom Dockerfile) # --------------------------------------------------------------------------- # Tested: lmsysorg/sglang:v0.5.10.post1 sglang: lmsysorg/sglang:v0.5.10.post1 + # --------------------------------------------------------------------------- + # OPTIONAL -- launcher image, ONLY for users in an airgapped environment who + # cannot `uv sync` / pip-install. It is NOT a worker and is NOT referenced by + # my_cluster.yaml `containers:`; it bundles the `nflow` CLI + baked venv to + # drive NVFlow over an ssh_tunnel with no host install and no user-side internet. + # The default install remains `uv sync` (README / INSTALL.md). + # Built multi-arch (amd64 + arm64) from dockerfiles/Dockerfile.nvflow; see + # docs/remote-launch.md. + # + # Tested: nvflow-client:v1.1.2. Substitute the tag you built and pushed. + # --------------------------------------------------------------------------- + nvflow-client: /nvflow-client: + # --------------------------------------------------------------------------- # 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 86e14a3..387339a 100644 --- a/cluster_configs/template-slurm.yaml +++ b/cluster_configs/template-slurm.yaml @@ -5,14 +5,15 @@ # # Then update all values with your settings. # -# 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) +# Container versions tested with this release: +# nemo-skills: nvflow-nemo-skills (NeMo-Skills @ e06c9b90, tag v1.1.2) +# vllm: nvflow-vllm (base vllm/vllm-openai v0.22.0, standalone SDG/eval) +# vllm-grpo: nvflow-vllm at the v0.20.0 tag (base vllm/vllm-openai v0.20.0, GRPO rollouts/judge) # sglang: lmsysorg/sglang v0.5.10.post1 -# nemo-rl: nvflow-nemo-rl (base nvcr.io/nvidia/nemo-rl:v0.6.0, pre-built venvs + Gym) +# nemo-rl: nvflow-nemo-rl (base nvcr.io/nvidia/nemo-rl v0.7.0; SFT/GRPO trainer, Gym venvs baked) +# nemo-gym: nvflow-nemo-gym (base python:3.12-slim, CPU-only; per-component Gym venvs β€” DG-SDG gym stages) # -# Reference: https://github.com/NVIDIA/NeMo-Skills +# Reference: https://github.com/NVIDIA-NeMo/Skills executor: slurm @@ -80,11 +81,13 @@ extra_sandbox_args: # After converting containers to .sqsh format (see INSTALL.md), paste paths here. containers: # 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 + # setup_containers.sh names each file -.sqsh. + nemo-skills: /nemo-skills-v1.1.2.sqsh # Orchestration client (eval, SDG, data prep) + vllm: /vllm-v0.22.0.sqsh # vLLM v0.22.0 standalone (SDG, eval) + vllm-grpo: /vllm-grpo-v0.20.0.sqsh # vLLM v0.20.0 standalone (GRPO rollouts, judge); same image repo as vllm, different tag + sglang: /sglang-v0.5.10.post1.sqsh # sglang inference server (SDG stages 3-5) + nemo-rl: /nemo-rl-v0.7.0.sqsh # NeMo-RL v0.7.0 for SFT and GRPO training + nemo-gym: /nemo-gym-0.4.0.sqsh # CPU-only Gym client for the Gym-only stages (DG-SDG gym stages; no NeMo-RL trainer) # Optional containers (not currently used by NVFlow recipes) # trtllm: /trtllm.sqsh # megatron: /megatron.sqsh @@ -97,17 +100,22 @@ containers: # Map cluster paths to container paths mounts: - :/hf_models # HuggingFace models - - :/workspace # Your workspace + # /workspace is a WRITABLE DATA dir (outputs + caches) -- NOT the source repo. + # Recipe code and checked-in assets (prompts, dataset descriptors, overlays) ship + # via the nemo-run packaged snapshot at /nemo_run/code on every worker, so the + # repo is never mounted. This mount only needs to hold writable runtime data: + # /workspace/outputs/** stage outputs, checkpoints, SEC cache, eval-datasets + # /workspace/cache/huggingface HF_HOME (see env_vars below) + # Point it at a dedicated data dir (e.g. a sibling of your repo checkout), not the + # checkout itself. + - :/workspace # Writable data dir (outputs + cache) # 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 + # No Gym source mount is needed: the nemo-rl and nemo-gym images both bake + # their NeMo-Gym venvs, so nothing is built at job runtime. Bind-mounting a + # host clone over /opt/nemo-rl/3rdparty/Gym-workspace/Gym shadows those baked + # venvs and breaks the GRPO stages -- do it only for dev-mode work. # ============================================================================= # Timeouts (per partition) @@ -135,13 +143,17 @@ env_vars: - VIRTUAL_ENV_PROMPT= # Unset venv prompt # --- 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. + # Prevents accidental network calls at runtime. Clear HF_*_OFFLINE temporarily + # for one-time stages that download external data (download_sec_filings, + # create_seed_data). - HF_HUB_OFFLINE=1 - HF_DATASETS_OFFLINE=1 - TRANSFORMERS_OFFLINE=1 - - UV_OFFLINE=true + # UV_OFFLINE: left UNSET (global flag). The images bake every venv they need, + # so no stage resolves packages at runtime either way. Leaving it unset keeps + # the escape hatch for dev-mode work on components outside the baked set. + # Set true only for a strict-airgap cluster. + # - UV_OFFLINE=true # Pre-cached tiktoken encodings (baked into vllm/vllm-grpo containers; # set here for uniform coverage across all container types). @@ -153,10 +165,6 @@ env_vars: - 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-gym b/dockerfiles/Dockerfile.nemo-gym new file mode 100644 index 0000000..9bb233c --- /dev/null +++ b/dockerfiles/Dockerfile.nemo-gym @@ -0,0 +1,102 @@ +# ============================================================================= +# NVFlow NeMo-Gym container (CPU-only). For Gym-only stages (prepare_data / +# prefetch_cache), the eval rollout client, and SDG rollouts. GRPO/training +# uses Dockerfile.nemo-rl. +# +# Bakes one venv per component (`gym env start ... +dry_run=true`) for +# equivalence_llm_judge, finance_sec_search and format_verification, reused +# offline at runtime. They live in /opt/gym-venvs, outside /opt/Gym, so a dev +# source mount overlays source only. Non-baked components fall back to an +# on-demand `uv` build, which needs a writable venv dir and network. +# +# Build (GYM_REF is a pinned SHA, so a cached clone layer cannot go stale). +# Single-arch, host platform; see docker_instructions.md for multi-arch: +# docker build -f dockerfiles/Dockerfile.nemo-gym -t nvflow-nemo-gym:0.4.0 . +# ============================================================================= +ARG PYTHON_VERSION=3.12 +ARG GYM_VENV_DIR=/opt/gym-venvs +ARG GYM_CLI_VENV=/opt/gym-cli-venv + +# --- builder ----------------------------------------------------------------- +FROM python:${PYTHON_VERSION}-slim AS builder +ARG PYTHON_VERSION +ARG GYM_VENV_DIR +ARG GYM_CLI_VENV +ENV UV_INSTALL_DIR=/usr/local/bin UV_PYTHON_PREFERENCE=only-system + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git curl ca-certificates build-essential && rm -rf /var/lib/apt/lists/* +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +# Upstream Gym main; the finance fork is retired (its fixes landed via PR #2055). +ARG GYM_REPO_URL=https://github.com/NVIDIA-NeMo/Gym.git +ARG GYM_REF=33ef60369f76557e6a6dd828c0bd5f5529624a92 +RUN git clone "${GYM_REPO_URL}" /opt/Gym && cd /opt/Gym && git checkout --detach "${GYM_REF}" && \ + test "$(git rev-parse HEAD)" = "${GYM_REF}" && \ + echo "Gym baked at ${GYM_REF}" + +# CLI venv outside /opt/Gym so a source mount can't shadow it. The `gym` wrapper +# calls the venv python directly: uv's editable entry-point script is unreliable +# under podman/enroot. +RUN uv venv "${GYM_CLI_VENV}" --python "$(command -v python${PYTHON_VERSION})" && \ + cd /opt/Gym && uv pip install --python "${GYM_CLI_VENV}/bin/python" -e "." +RUN printf '#!/bin/sh\nexec "%s/bin/python" -c "from nemo_gym.cli.main import main; main()" "$@"\n' \ + "${GYM_CLI_VENV}" > /usr/local/bin/gym && chmod +x /usr/local/bin/gym + +# Bake per-component venvs. Inline values only satisfy config resolution during +# the dry-run; the recipe sets real values at runtime. +RUN cd /opt/Gym && \ + gym env start --resources-server equivalence_llm_judge --model-type vllm_model \ + +dry_run=true +uv_venv_dir="${GYM_VENV_DIR}" +skip_venv_if_present=false \ + +policy_base_url=http://unset/v1 +policy_api_key=unset +policy_model_name=unset && \ + gym env start --resources-server finance_sec_search --model-type vllm_model \ + +dry_run=true +uv_venv_dir="${GYM_VENV_DIR}" +skip_venv_if_present=false \ + +policy_base_url=http://unset/v1 +policy_api_key=unset +policy_model_name=unset \ + +search_judge_model_base_url=https://api.openai.com/v1 \ + +search_judge_model_api_key=unset +search_judge_model_name=gpt-5-mini +tavily_api_key=null && \ + gym env start --resources-server format_verification/freeform_formatting --model-type vllm_model \ + +dry_run=true +uv_venv_dir="${GYM_VENV_DIR}" +skip_venv_if_present=false \ + +policy_base_url=http://unset/v1 +policy_api_key=unset +policy_model_name=unset + +# Gym calls uvicorn.run(timeout_worker_healthcheck=) (uvicorn>=0.37) but declares +# no floor; assert so a resolver regression fails here, not at the first rollout. +RUN "${GYM_VENV_DIR}/resources_servers/finance_sec_search/.venv/bin/python" -c \ + "import inspect, uvicorn; \ +assert 'timeout_worker_healthcheck' in inspect.signature(uvicorn.run).parameters, uvicorn.__version__; \ +print('uvicorn', uvicorn.__version__, 'OK')" + +# Record the ref before dropping history, so a running container can still report +# which Gym it carries. +RUN echo "${GYM_REF}" > /opt/gym-commit && rm -rf /opt/Gym/.git + +# --- final ------------------------------------------------------------------- +FROM python:${PYTHON_VERSION}-slim +ARG GYM_VENV_DIR +ARG GYM_CLI_VENV +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates procps curl && rm -rf /var/lib/apt/lists/* + +COPY --from=builder /opt/gym-commit /opt/gym-commit +COPY --from=builder /opt/Gym /opt/Gym +COPY --from=builder ${GYM_CLI_VENV} ${GYM_CLI_VENV} +COPY --from=builder ${GYM_VENV_DIR} ${GYM_VENV_DIR} +COPY --from=builder /usr/local/bin/gym /usr/local/bin/gym + +# Let a non-root runtime create new component venvs: open the top + category +# dirs; baked component dirs and venv files stay untouched. +RUN chmod a+rwX ${GYM_VENV_DIR} ${GYM_VENV_DIR}/*/ + +# uv: on-demand venv build for non-baked components (baked ones never use it). +COPY --from=builder /usr/local/bin/uv /usr/local/bin/uv + +# DG-SDG helper deps in the base interpreter for nvflow `python3 -m ...` jobs. +RUN pip install --no-cache-dir orjson pyyaml + +WORKDIR / + +# =========================================================================== +# Security hardening (Trivy/NSPECT scans, 2026-07-08) +# =========================================================================== +# rm all ray_dist.jar copies (shipped 3x): jackson-databind RCE CVE-2026-54512/CVE-2026-54513 (Ray-Java unused); fail build if one survives +RUN find /opt /usr/local -name 'ray_dist.jar' -type f -delete 2>/dev/null; \ + ! find /opt /usr/local -name 'ray_dist.jar' -type f 2>/dev/null | grep -q . diff --git a/dockerfiles/Dockerfile.nemo-rl b/dockerfiles/Dockerfile.nemo-rl index 56c21ad..bfd2293 100644 --- a/dockerfiles/Dockerfile.nemo-rl +++ b/dockerfiles/Dockerfile.nemo-rl @@ -1,131 +1,123 @@ # ============================================================================= -# 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. +# NVFlow NeMo-RL trainer (airgapped) for GRPO/SFT `training`. +# +# The base bakes the RL environment but leaves the Gym venvs unbuilt (upstream +# gates that prefetch behind NEMO_GYM_PREFETCH_CONFIGS). This image bakes them, +# so no `uv` resolve happens at job runtime. +# +# Run with `--no-container-mount-home`, which the NeMo-Skills launcher always +# passes: the venvs resolve through /root, which enroot otherwise shadows with +# $HOME. # -# Build: -# docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:latest . +# Build (single-arch, host platform; see docker_instructions.md for multi-arch): +# docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:v0.7.0 . # ============================================================================= - -ARG BASE_IMAGE=nvcr.io/nvidia/nemo-rl:v0.6.0 +ARG BASE_IMAGE=nvcr.io/nvidia/nemo-rl:v0.7.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 +ARG GYM_REF=33ef60369f76557e6a6dd828c0bd5f5529624a92 +ARG NEMO_GYM_CUDA=cu130 +ARG NEMO_GYM_VLLM_VERSION=0.20.0 +ARG TARGETARCH +ARG GYM_SRC=/opt/nemo-rl/3rdparty/Gym-workspace/Gym +ARG GYM_VENV=/opt/ray_venvs/nemo_rl.environments.nemo_gym.NemoGym -# --- 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 +# scripts/convert_checkpoint_to_hf.sh cd's to /opt/NeMo-RL (wrong case). +RUN ln -sf /opt/nemo-rl /opt/NeMo-RL -# --- 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 +# Versions every process joining the Ray cluster must share. Gym's ray floor is +# only >=2.55.1, which would not stop a resolver from moving it. +RUN /opt/nemo_rl_venv/bin/python -c \ + "import numpy, ray; print(f'numpy=={numpy.__version__}'); print(f'ray=={ray.__version__}')" \ + > /opt/nvflow-pins.txt && \ + cat /opt/nvflow-pins.txt + +# Advance in place: generate_fingerprint.py hashes submodule SHAs. +WORKDIR ${GYM_SRC} +RUN git fetch --depth 1 origin ${GYM_REF} && \ + git checkout --detach ${GYM_REF} && \ + test "$(git rev-parse HEAD)" = "${GYM_REF}" + +# Pick up deps Gym declared since the base was built (editable, so source follows). +RUN uv pip install --python ${GYM_VENV}/bin/python \ + --constraint /opt/nvflow-pins.txt -e . && \ + ${GYM_VENV}/bin/python -c "import nemo_gym" + +# Bake one venv per Gym component. Driving the CLI from the actor venv is what +# makes Gym pin each component to that interpreter's ray== and python_version(). +# Gym's vllm==0.20.0 defaults to a cu12 wheel; override to cu130 to match torch. +RUN <<"EOF" bash -eux +case "${TARGETARCH:-amd64}" in arm64) WHEEL_ARCH=aarch64 ;; *) WHEEL_ARCH=x86_64 ;; esac +printf 'vllm @ https://github.com/vllm-project/vllm/releases/download/v%s/vllm-%s-cp38-abi3-manylinux_2_35_%s.whl\n' \ + "${NEMO_GYM_VLLM_VERSION}" "${NEMO_GYM_VLLM_VERSION}" "${WHEEL_ARCH}" > /tmp/gym-vllm.txt +export UV_TORCH_BACKEND="${NEMO_GYM_CUDA}" UV_OVERRIDE=/tmp/gym-vllm.txt UV_LINK_MODE=symlink + +# Empty to start, so skip_venv_if_present can only reuse venvs from this build +# (the shared vllm_model/agent ones), never a stale one. +test -z "$(ls -A "${NEMO_GYM_VENV_DIR}" 2>/dev/null)" + +BAKE="${GYM_VENV}/bin/gym env start --model-type vllm_model +dry_run=true" +BAKE="${BAKE} +uv_venv_dir=${NEMO_GYM_VENV_DIR} +skip_venv_if_present=true" +BAKE="${BAKE} +policy_base_url=http://unset/v1 +policy_api_key=unset +policy_model_name=unset" + +${BAKE} --resources-server equivalence_llm_judge +${BAKE} --resources-server format_verification/freeform_formatting +${BAKE} --resources-server finance_sec_search \ + +search_judge_model_base_url=https://api.openai.com/v1 \ + +search_judge_model_api_key=unset +search_judge_model_name=gpt-5-mini \ + +tavily_api_key=null + +rm -f /tmp/gym-vllm.txt + +# Gym calls uvicorn.run(timeout_worker_healthcheck=) (uvicorn>=0.37) but declares +# no floor; assert so a resolver regression fails here, not at the first rollout. +SEC_VENV="${NEMO_GYM_VENV_DIR}/resources_servers/finance_sec_search/.venv" +test -d "$SEC_VENV" +"$SEC_VENV/bin/python" -c "import inspect, uvicorn; \ +assert 'timeout_worker_healthcheck' in inspect.signature(uvicorn.run).parameters, uvicorn.__version__; \ +print('uvicorn', uvicorn.__version__, 'OK')" +EOF + +WORKDIR /opt/nemo-rl +RUN python tools/generate_fingerprint.py > /opt/nemo_rl_container_fingerprint + +# =========================================================================== +# Security hardening (mirrors Dockerfile.vllm) +# =========================================================================== +# Keep headers installed: triton and TransformerEngine JIT-compile at run time. +RUN apt-get update && apt-get upgrade -y && rm -rf /var/lib/apt/lists/* + +# ray_dist.jar: jackson-databind RCE CVE-2026-54512/CVE-2026-54513, Ray-Java +# unused. Most copies are venv symlinks into the uv cache, so the cache holds the +# only real file and must be searched -- but never delete the cache itself, which +# the worker venvs symlink into for everything else. +RUN find /usr/local /opt /root/.cache/uv -name 'ray_dist.jar' -delete 2>/dev/null; \ + ! find /usr/local /opt /root/.cache/uv -name 'ray_dist.jar' 2>/dev/null | grep -q . && \ + /opt/nemo_rl_venv/bin/python -c "import ray; print('ray OK', ray.__version__)" + +# Ray refuses to join a cluster on a different Ray or Python version. numpy is +# gated only in the actor venvs, which exchange pickled arrays; Gym components +# talk HTTP/orjson, so a difference there is reported, not fatal. +RUN <<"EOF" bash -eux +ref() { /opt/nemo_rl_venv/bin/python -c "import $1 as m, platform; print(m.__version__)"; } +REF_PY=$(/opt/nemo_rl_venv/bin/python -c "import platform; print(platform.python_version())") +REF_RAY=$(ref ray); REF_NP=$(ref numpy) +FOUND=0 +for py in /opt/ray_venvs/*/bin/python "${NEMO_GYM_VENV_DIR}"/*/*/.venv/bin/python; do + [ -x "$py" ] || continue + FOUND=$((FOUND + 1)) + got() { "$py" -c "import $1 as m; print(m.__version__)" 2>/dev/null || true; } + V=$("$py" -c "import platform; print(platform.python_version())") + [ "$V" = "$REF_PY" ] || { echo "python skew: $py is $V, want $REF_PY"; exit 1; } + RAY=$(got ray) + [ -z "$RAY" ] || [ "$RAY" = "$REF_RAY" ] || { echo "ray skew: $py has $RAY, want $REF_RAY"; exit 1; } + NP=$(got numpy) + case "$py" in + /opt/ray_venvs/*) [ -z "$NP" ] || [ "$NP" = "$REF_NP" ] || \ + { echo "numpy skew: $py has $NP, want $REF_NP"; exit 1; } ;; + *) [ -z "$NP" ] || [ "$NP" = "$REF_NP" ] || echo "note: gym venv $py numpy $NP vs $REF_NP" ;; + esac +done +[ "$FOUND" -gt 0 ] || { echo "no venvs inspected"; exit 1; } +echo "python ${REF_PY} / ray ${REF_RAY} consistent across ${FOUND} venvs; numpy ${REF_NP} in actor venvs" +EOF diff --git a/dockerfiles/Dockerfile.nemo-skills b/dockerfiles/Dockerfile.nemo-skills index a67c6eb..e390c3a 100644 --- a/dockerfiles/Dockerfile.nemo-skills +++ b/dockerfiles/Dockerfile.nemo-skills @@ -1,49 +1,75 @@ # ============================================================================= # 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 . +# 1. UPSTREAM (NeMo-Skills' Dockerfile at a pinned commit, kept diffable), +# 2. NVFLOW (packages the workflow steps import), 3. CVE (scan-driven only -- +# delete it for a stock image). Upstream: +# https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.nemo-skills # -# Upstream source: -# https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.nemo-skills +# Build (single-arch, host platform; see docker_instructions.md for multi-arch): +# docker build --no-cache \ +# -f dockerfiles/Dockerfile.nemo-skills -t nvflow-nemo-skills:v1.1.2 . # ============================================================================= +# Section 3 item; Docker needs builder stages first, so delete it with the COPY in +# section 3. wandb 0.28.1 (newest release) bundles wandb-core built with Go 1.26.4, +# grpc-go v1.82.0 and x/text v0.38.0; commit e118409 is the unreleased fix. +# Cross-compiled from BUILDPLATFORM to skip arm64 emulation. +ARG WANDB_CORE_COMMIT=e1184091520c9b44aa1096fdb27b2f4bf52f26d7 +FROM --platform=$BUILDPLATFORM golang:1.26.5 AS wandb-core-builder +ARG WANDB_CORE_COMMIT +ARG TARGETARCH +RUN git init /src/wandb && cd /src/wandb && \ + git remote add origin https://github.com/wandb/wandb.git && \ + git sparse-checkout init --cone && git sparse-checkout set core && \ + git fetch --depth 1 origin "${WANDB_CORE_COMMIT}" && git checkout --detach FETCH_HEAD +RUN cd /src/wandb/core && \ + CGO_ENABLED=0 GOOS=linux GOARCH=${TARGETARCH} go build \ + -tags "disable_grpc_modules parquet_read_only" \ + -ldflags "-s -w -X main.commit=${WANDB_CORE_COMMIT}" -mod=vendor \ + -o /wandb-core ./cmd/wandb-core && \ + go version -m /wandb-core | grep -F "go1.26.5" && \ + go version -m /wandb-core | grep -E "google\.golang\.org/grpc[[:space:]]+v1\.82\.1([[:space:]]|$)" && \ + go version -m /wandb-core | grep -E "golang\.org/x/text[[:space:]]+v0\.40\.0([[:space:]]|$)" -# Clone NeMo-Skills at a pinned commit (replaces build-context COPY commands) +# NeMo-Skills at a pinned commit, replacing upstream's build-context COPYs. FROM scratch AS nemo-skills-src -ARG NEMO_SKILLS_COMMIT=022904023ad7a83a87662a313cf72e7df5891d55 +ARG NEMO_SKILLS_COMMIT=e06c9b900177be3f60d6a3f99135bb5de9af9bed 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") +# 1. BEGIN UPSTREAM -- github.com/NVIDIA-NeMo/Skills @ e06c9b90 +# Deviations, so a diff after re-syncing shows only these: +# - COPY -> COPY --from=nemo-skills-src (pinned SHA instead of build context) +# - google-research not cloned (unused here; carries Criticals) +# - `cd /tmp` before the nltk download (upstream defect; see below) +# - ARG WANDB_CORE_COMMIT redeclared, for the CVE section's COPY +# Do not add --no-install-recommends below: it drops gpg-agent, which +# add-apt-repository needs for the apptainer PPA key. # =========================================================================== +# using ubuntu instead of debian for easier apptainer installation on arm64 FROM ubuntu:22.04 +ARG WANDB_CORE_COMMIT +# Install Python and other dependencies RUN apt-get update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y \ + apt-get install -y \ python3.10 \ python3-pip \ curl \ wget \ git \ git-lfs \ - ffmpeg \ - tzdata && \ + ffmpeg && \ 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 pip install --upgrade pip setuptools "uv>=0.11.10" +# Update package lists and install apptainer for arm64 +# https://apptainer.org/docs/admin/1.1/installation.html RUN apt update && \ apt install -y software-properties-common && \ add-apt-repository -y ppa:apptainer/ppa && \ @@ -52,8 +78,18 @@ RUN apt update && \ apt update && apt install -y apptainer-suid && \ rm -rf /var/cache/apt/archives /var/lib/apt/lists/* +# Apply security patches for PackageKit, pulled in transitively by software-properties-common. +# Ubuntu 22.04 has published 1.2.5-2ubuntu3.1 with the fix for the local privilege escalation CVE. +RUN apt-get update && \ + apt-get install --only-upgrade -y \ + packagekit \ + packagekit-tools \ + libpackagekit-glib2-18 \ + gir1.2-packagekitglib-1.0 && \ + rm -rf /var/cache/apt/archives /var/lib/apt/lists/* + +# for ifeval benchmark -- google-research clone skipped (see deviations) 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 @@ -61,6 +97,7 @@ RUN cd /opt/gorilla/berkeley-function-call-leaderboard && pip install --no-cache RUN apt remove -y python3-blinker +# ifbench ARG IFBENCH_COMMIT=c6767a19bd82ac0536cab950f2f8f6bcc6fabe7c ARG IFBENCH_REPO=https://github.com/allenai/IFBench.git ARG IFBENCH_DIR=/opt/benchmarks/IFBench @@ -68,24 +105,32 @@ RUN git init "$IFBENCH_DIR" && cd "$IFBENCH_DIR" && git remote add origin "$IFBE git fetch --depth 1 origin "${IFBENCH_COMMIT}" && git reset --hard FETCH_HEAD RUN cd ${IFBENCH_DIR} && pip install -r requirements.txt +# removing on-the-fly installation in ifbench to avoid conflicts from parallel jobs COPY --from=nemo-skills-src /dockerfiles/ifbench.patch /opt/benchmarks/IFBench/ifbench.patch RUN cd /opt/benchmarks/IFBench && git apply ifbench.patch +# nltk >=3.10.1 blocks imports resolving inside the CWD; at CWD=/ that matches +# every stdlib path, so `import nltk` needs the `cd`. Downloads use fixed paths. 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'); \ + cd /tmp && 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')" +# we aren't copying main nemo_skills folder as it will always be mounted from host +# but we do want to install all requirements in the container directly 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 +# installing sdp in container only 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 +# Install via `uv pip` from the project directory so [tool.uv].override-dependencies +# in pyproject.toml (which relaxes leptonai's httpx==0.27.2 pin so litellm 1.83.x +# can be installed) is picked up. Plain pip ignores [tool.uv] and the resolver fails. +RUN cd /opt/NeMo-Skills && uv pip install --system --no-cache-dir \ + -r core/requirements.txt -r requirements/pipeline.txt +# Fix http mismatch between lepton and dggs by manually downloading dggs here RUN pip install ddgs # =========================================================================== @@ -94,12 +139,14 @@ RUN pip install ddgs # =========================================================================== -# NVFlow Additional Packages -# =========================================================================== -# Pre-install packages used by NVFlow workflow steps so they are available -# at runtime without needing to download anything. +# 2. NVFLOW -- imported by workflow steps; baked so no job downloads at launch. # =========================================================================== +# tzdata populates /usr/share/zoneinfo for pyarrow/pandas libc tz lookups. +RUN apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends tzdata && \ + rm -rf /var/cache/apt/archives /var/lib/apt/lists/* + RUN pip install --no-cache-dir --ignore-requires-python \ jsonlines \ tiktoken \ @@ -110,15 +157,12 @@ RUN pip install --no-cache-dir --ignore-requires-python \ "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. +# SEC data prep: workflow-2 download_sec_filings, workflow-3 create_seed_data. RUN pip install --no-cache-dir --ignore-requires-python \ edgartools==5.20.2 \ sec-parser \ @@ -130,9 +174,98 @@ RUN pip install --no-cache-dir --ignore-requires-python \ 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')" + from bs4 import BeautifulSoup; from datasets import load_dataset; \ + print('NVFlow deps OK')" + + +# =========================================================================== +# 3. CVE -- scan findings only, not functionality; drop each once upstream ships +# the fix. Floors upstream already satisfies are asserted, never re-pinned -- +# re-pinning core packages is what silently broke tokenizers once. +# =========================================================================== + +# libssl3 -> High CVE-2026-45447. linux-libc-dev's ~176 header findings are inert +# but fix-available; nothing compiles at runtime. +RUN apt-get update && apt-get install -y --only-upgrade --no-install-recommends libssl3 && \ + apt-get purge -y linux-libc-dev && apt-get autoremove -y && \ + rm -rf /var/lib/apt/lists/* + +# cssutils imports more_itertools at module level, and pip earlier resolved that +# against an apt copy which autoremove above then took away. Own it with pip. +RUN pip install --no-cache-dir --ignore-installed more-itertools + +# starlette and lxml have no satisfiable version: the fixes (GHSA-86qp-5c8j-p5mr, +# -x746-7m8f-x49c, -wqp7-x3pw-xc5r, -jp82-jpqv-5vv3, -82w8-qh3p-5jfq; and +# GHSA-vfmq-68hx-4jfw) landed in 1.x and 6.1.0, above the caps in leptonai's +# instrumentator pin and in sec-parser. Forced, and allowlisted in the pip check. +# GitPython and datamodel-code-generator are not in upstream's graph, so they are +# requested here. Full requirement set, no --upgrade: no ceiling missed, nothing +# else moves. +RUN printf '%s\n' 'starlette>=1.3.1' 'lxml>=6.1.0' > /tmp/cve-overrides.txt && \ + cd /opt/NeMo-Skills && uv pip install --system --no-cache-dir \ + --override /tmp/cve-overrides.txt \ + -r core/requirements.txt -r requirements/pipeline.txt \ + 'GitPython>=3.1.55' 'datamodel-code-generator>=0.64.0' + +# Pairs with the wandb-core-builder stage; the commit is injected at link time. +COPY --from=wandb-core-builder /wandb-core /usr/local/lib/python3.10/dist-packages/wandb/bin/wandb-core +RUN /usr/local/lib/python3.10/dist-packages/wandb/bin/wandb-core --help 2>&1 | \ + grep -F "Commit SHA: ${WANDB_CORE_COMMIT}" + +# Benchmark trees finance never runs, and Ray's Java jar (jackson-databind +# findings; Ray is used only through Python). bfcl_eval was installed editable, so +# uninstall it rather than leave a dangling dist-info. +RUN uv pip uninstall --system bfcl_eval || true && \ + rm -rf /opt/gorilla /opt/benchmarks/IFBench /usr/local/lib/python3*/dist-packages/ray/jars && \ + ! python -c "import bfcl_eval" 2>/dev/null && \ + [ -z "$(find /usr/local/lib/python3*/dist-packages -maxdepth 1 -name '*bfcl*')" ] && \ + ! find /usr/local /opt -name 'ray_dist.jar' -type f 2>/dev/null | grep -q . && \ + python -c "import ray; print('bfcl/IFBench removed, ray OK:', ray.__version__)" + +# Asserted, not pinned, so a regression fails the build rather than being papered +# over. wandb is exact: the patched core binary must match its Python protocol. +RUN python -c "from importlib.metadata import version as v; from packaging.requirements import Requirement as R; \ + floors = ['starlette>=1.3.1', 'GitPython>=3.1.55', 'datamodel-code-generator>=0.64.0', \ + 'nltk>=3.10.0', 'wandb==0.28.1', 'litellm>=1.84.10', 'lxml>=6.1.0', \ + 'httpx>=0.28.1', 'urllib3>=2.6.3', 'msgpack>=1.2.1', 'setuptools>=78.1.1', \ + 'click>=8.2,<9', 'typer>=0.16,<0.27', 'mcp<2.0']; \ + bad = ['%s %s needs %s' % (r.name, v(r.name), r.specifier) for r in map(R, floors) \ + if not r.specifier.contains(v(r.name), prereleases=True)]; \ + assert not bad, bad; print('CVE floors OK')" + +# Gate what the upgrades put at risk, here rather than on the cluster: transformers +# enforces its tokenizers range at import, and lxml is forced past sec-parser's +# cap, so parse a document rather than only importing it. +RUN python -c "from transformers import AutoTokenizer; import transformers, tokenizers, lxml.etree, sec_parser; \ + import cssutils; from litellm import completion; from fastapi import FastAPI; FastAPI(); \ + els = sec_parser.Edgar10QParser().parse('

Item 2. Management Discussion

Revenue rose 12 percent.

'); \ + assert len(els) >= 2, els; \ + print('gates OK: transformers', transformers.__version__, 'tokenizers', tokenizers.__version__, \ + 'lxml', lxml.etree.__version__, '/', len(els), 'sec elements')" + +RUN WANDB_MODE=offline WANDB_SILENT=true WANDB_DIR=/tmp/wandb-smoke \ + python -c "import wandb; run = wandb.init(project='nvflow-security-smoke', name='offline'); \ + run.log({'metric': 1.0}); run.finish(); print('wandb offline smoke OK')" && \ + rm -rf /tmp/wandb-smoke + +# Whole-environment check, last so nothing can invalidate it. Matched on the exact +# pair, so a new break fails even in a listed package. The first three are +# upstream's own (its pyproject overrides, and sdp's numpy pin). +RUN printf '%s\n' \ + 'leptonai .* requirement httpx' \ + 'torchx .* requirement urllib3' \ + 'sdp .* requirement numpy' \ + 'prometheus-fastapi-instrumentator .* requirement starlette' \ + 'sec-parser .* requirement lxml' \ + 'No broken requirements found' \ + > /tmp/pipcheck-allow.txt; \ + pip check > /tmp/pipcheck.txt 2>&1 || true; \ + if grep -vEf /tmp/pipcheck-allow.txt /tmp/pipcheck.txt | grep -q '[^[:space:]]'; then \ + echo "ERROR: unexpected dependency inconsistency:"; cat /tmp/pipcheck.txt; exit 1; \ + fi; \ + cat /tmp/pipcheck.txt + +# uv's sdist cache leaves Git metadata that nSpect's global policy flags. +RUN rm -rf /root/.cache/uv /tmp/cve-overrides.txt /tmp/pipcheck.txt /tmp/pipcheck-allow.txt diff --git a/dockerfiles/Dockerfile.nvflow b/dockerfiles/Dockerfile.nvflow new file mode 100644 index 0000000..7da7d5a --- /dev/null +++ b/dockerfiles/Dockerfile.nvflow @@ -0,0 +1,151 @@ +# syntax=docker/dockerfile:1 +# ============================================================================= +# NVFlow client (launcher) image -- the `nflow` CLI + baked venv for airgap use. +# +# Lean, unprivileged orchestration image: `nflow` submits Slurm jobs over an SSH +# tunnel, so no Slurm client / munge / podman / enroot is needed inside. Worker +# images stay as .sqsh on the cluster (referenced in my_cluster.yaml), not here. +# +# Requires a source with the tunnel-aware launcher (feature/finance-rl-grpo or +# later); the v1.1.1 release cannot launch off-cluster over ssh_tunnel. +# +# Build from the repo root, on a committed tree (.baked_commit records HEAD). +# .dockerignore drops cache/, .venv and secrets; .git is kept for hatch-vcs +# versioning, then squashed to a single history-free commit so the image ships +# no repo history. Single-arch, host platform; see docker_instructions.md for +# multi-arch: +# docker build -f dockerfiles/Dockerfile.nvflow -t nvflow-client:v1.1.2 . +# +# Usage: see docs/remote-launch.md. +# ============================================================================= + +ARG PYTHON_VERSION=3.12 +# Pin the multi-arch Ubuntu 24.04 manifest so a rebuild cannot silently move to +# a new OS release. The previous floating python:3.12-slim tag moved to Debian +# 13 and introduced 17 CRITICAL + 40 HIGH OS-package findings per architecture. +ARG UBUNTU_IMAGE=ubuntu:24.04@sha256:4fbb8e6a8395de5a7550b33509421a2bafbc0aab6c06ba2cef9ebffbc7092d90 + +# --- builder: resolve the venv (build tools stay out of the final image) ----- +FROM ${UBUNTU_IMAGE} AS builder +ARG PYTHON_VERSION +ENV DEBIAN_FRONTEND=noninteractive \ + UV_INSTALL_DIR=/usr/local/bin \ + UV_CACHE_DIR=/opt/uv-cache \ + UV_PYTHON_PREFERENCE=only-system + +RUN apt-get update && apt-get install -y --no-install-recommends \ + "python${PYTHON_VERSION}" "python${PYTHON_VERSION}-venv" \ + git curl ca-certificates build-essential \ + && rm -rf /var/lib/apt/lists/* + +# uv on a system path (not /root/.local) so it survives enroot's $HOME remap. +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +COPY . /opt/nvflow +WORKDIR /opt/nvflow +# Record the baked commit for provenance, and hide it from git so the baked tree +# stays clean (nemo-run packages via `git archive` of HEAD regardless). +RUN git rev-parse HEAD > /opt/nvflow/.baked_commit 2>/dev/null || echo unknown > /opt/nvflow/.baked_commit && \ + echo '.baked_commit' >> /opt/nvflow/.git/info/exclude + +# --frozen pins to the committed uv.lock; --no-dev skips the dev group (debugger). +RUN uv venv .venv --python "$(command -v python${PYTHON_VERSION})" && \ + uv sync --frozen --no-dev +# Strip the vendored wandb "core" Go binary. wandb arrives transitively via +# nemo-skills, but the launcher never calls wandb.init()/sync (only the Python +# API is imported), so wandb-core is dead weight that vendors Go stdlib/grpc/ +# x-crypto CVEs (e.g. CVE-2025-68121, CVE-2026-33186) and dominates image scans. +# Deleting the binary keeps the wandb Python package importable for nemo-skills. +# Strip Ray's bundled Java jar (site-packages/ray/jars/ray_dist.jar). Ray arrives +# transitively via nemo-skills; the launcher drives Ray purely through the Python +# Ray Jobs API, so the Java jar is dead weight carrying the jackson-databind HIGH +# CVEs (CVE-2026-54512 / CVE-2026-54513) plus other vendored Java libs +# (guava/gson/jaxb). Removing the whole jars/ dir clears them without affecting Ray. +RUN rm -rf /opt/uv-cache && \ + rm -f /opt/nvflow/.venv/lib/python*/site-packages/wandb/bin/wandb-core && \ + rm -rf /opt/nvflow/.venv/lib/python*/site-packages/ray/jars && \ + find /opt/nvflow -depth -type d -name __pycache__ -exec rm -rf {} + && \ + find /opt/nvflow -type f -name '*.pyc' -delete +# Fail-fast: Ray must still import after the jar strip (Python Ray Jobs API intact). +RUN /opt/nvflow/.venv/bin/python -c "import ray; print('ray OK after jar strip:', ray.__version__)" + +# Ship a history-free repo. hatch-vcs versioning already ran during `uv sync` +# (static nvflow/_version.py written above), and nemo-run only needs +# `git archive HEAD` (current tree) -- a distributed image must not carry repo +# history. Replace the full .git with a single snapshot commit that reproduces +# the EXACT original tracked set: capture `git ls-files` first, then re-add it +# with --force so tracked-but-gitignored files (e.g. recipes/finance/data/ +# __init__.py, which drives stage auto-discovery) are preserved. Plain +# `git add -A` respects .gitignore and would silently drop them, changing what +# `git archive HEAD` ships to the compute nodes. The assertion guards this. +RUN cd /opt/nvflow && \ + git ls-files -z > /tmp/tracked && \ + rm -rf .git && \ + git init -q -b main && \ + git -c user.email=release@nvidia.com -c user.name=nvflow add --pathspec-from-file=/tmp/tracked --pathspec-file-nul --force && \ + git -c user.email=release@nvidia.com -c user.name=nvflow commit -q -m "nvflow baked release snapshot" && \ + git ls-files -z | sort -z > /tmp/after && sort -z /tmp/tracked > /tmp/before && \ + cmp -s /tmp/before /tmp/after && echo "snapshot tree == original tracked set" && \ + rm -f /tmp/tracked /tmp/before /tmp/after + +# --- final: slim runtime = base + venv + source(.git snapshot) + launcher ------ +FROM ${UBUNTU_IMAGE} +ARG PYTHON_VERSION +ARG CA_CERTIFICATES_VERSION=20260601~24.04.1 +ARG GIT_VERSION=1:2.43.0-1ubuntu7.3 +ARG OPENSSH_CLIENT_VERSION=1:9.6p1-3ubuntu13.18 +ARG PYTHON_DEB_VERSION=3.12.3-1ubuntu0.15 +ARG RSYNC_VERSION=3.2.7-1ubuntu1.5 +LABEL org.opencontainers.image.title="nvflow-client" \ + org.opencontainers.image.description="NVFlow launcher (nflow CLI) for airgap use" + +# Launcher runtime deps: git (nemo-run `git archive`), openssh-client (ssh_tunnel), +# rsync (code sync to job_dir), ca-certificates, and the interpreter backing the +# copied venv. Pin the security-updated Ubuntu packages: if an exact version +# leaves the archive, fail the build for an intentional refresh instead of +# silently accepting a vulnerable package set. +RUN apt-get update && apt-get install -y --no-install-recommends \ + "python${PYTHON_VERSION}=${PYTHON_DEB_VERSION}" \ + "git=${GIT_VERSION}" \ + "openssh-client=${OPENSSH_CLIENT_VERSION}" \ + "rsync=${RSYNC_VERSION}" \ + "ca-certificates=${CA_CERTIFICATES_VERSION}" \ + && rm -rf /var/lib/apt/lists/* + +# ssh_tunnel host-key handling. nemo-run authenticates the tunnel via paramiko +# (given the key explicitly), but then rsyncs code with plain `ssh -i `, +# which reads known_hosts from $HOME/.ssh β€” NOT from the mounted /opt/ssh. On a +# fresh container $HOME/.ssh is empty, so rsync dies with "Host key verification +# failed". Point ssh at the *mounted* known_hosts and auto-accept a first-ever +# connect (written back to the mounted dir, so it persists) β€” no manual priming. +# System-wide via the ssh_config Include, so it holds for any HOME/uid, enroot or +# docker, on-cluster or off. Deliberately no IdentityFile here: nemo-run passes +# -i explicitly, and this stays agnostic to the user's key name. +RUN printf 'Host *\n UserKnownHostsFile /opt/ssh/known_hosts\n StrictHostKeyChecking accept-new\n' \ + > /etc/ssh/ssh_config.d/10-nvflow-tunnel.conf + +COPY --from=builder /usr/local/bin/uv /usr/local/bin/uv +COPY --from=builder /opt/nvflow /opt/nvflow + +# System gitconfig (not /root, which enroot remaps) so `git archive` doesn't trip +# git's "dubious ownership" guard when the container runs as a mapped uid. +RUN git config --system --add safe.directory /opt/nvflow + +# enroot ignores ENV PATH; symlink the nflow console script onto the default PATH. +RUN ln -sf /opt/nvflow/.venv/bin/nflow /usr/local/bin/nflow + +# UV_OFFLINE + UV_NO_SYNC: `uv run nflow` runs in the baked venv with no network +# and no pre-run sync (a sync would try to fetch the skipped dev group and fail). +# NEMO_SKILLS_DISABLE_UNCOMMITTED_CHANGES_CHECK: the image bakes a fixed committed +# snapshot; nemo-run packages HEAD via `git archive`, so its uncommitted-changes +# gate is a false positive here (build artifacts like the venv make the tree look +# dirty). Disabling it lets the launcher package the baked commit unattended. +ENV VIRTUAL_ENV=/opt/nvflow/.venv \ + PATH=/opt/nvflow/.venv/bin:/usr/local/bin:$PATH \ + NEMO_SKILLS_CONFIG_DIR=/opt/nvflow/cluster_configs \ + UV_CACHE_DIR=/tmp/uv-cache \ + UV_OFFLINE=1 \ + UV_NO_SYNC=1 \ + NEMO_SKILLS_DISABLE_UNCOMMITTED_CHANGES_CHECK=1 + +WORKDIR /opt/nvflow diff --git a/dockerfiles/Dockerfile.vllm b/dockerfiles/Dockerfile.vllm index 46bb563..8135f10 100644 --- a/dockerfiles/Dockerfile.vllm +++ b/dockerfiles/Dockerfile.vllm @@ -1,27 +1,41 @@ # ============================================================================= # NVFlow vLLM Container # ============================================================================= -# Self-contained Dockerfile that builds the vLLM inference container with -# pre-cached tokenizer encodings. +# vLLM inference container: adds SAM 3.1 and pre-caches tokenizer encodings so +# it serves airgapped. VLLM_VERSION selects the base tag -- both vLLM images +# build from this file and ship in the nvflow-vllm repo under different tags. # -# Build: -# docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:latest . +# Builds below are single-arch (host platform); see docker_instructions.md for +# multi-arch. +# +# Build (SDG / eval): +# docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:v0.22.0 . +# +# Build (GRPO rollouts + judge; matches NeMo-RL v0.7.0's colocated vLLM): +# docker build --build-arg VLLM_VERSION=v0.20.0 \ +# -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:v0.20.0 . # # Upstream source: # https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.vllm # ============================================================================= +ARG VLLM_VERSION=v0.22.0 + +FROM scratch AS sam3-src +ARG SAM3_COMMIT=a51b9f498c84824a94702cc289ed75d9cc544c64 +ADD --keep-git-dir=true https://github.com/facebookresearch/sam3.git#${SAM3_COMMIT} / + # =========================================================================== # BEGIN UPSTREAM (NeMo-Skills Dockerfile.vllm) # =========================================================================== -ARG VLLM_VERSION=v0.18.1 FROM vllm/vllm-openai:${VLLM_VERSION} +RUN pip install ray RUN pip install "vllm[audio]" +# Required by vLLM for Qwen-VL model family (runtime dependency, not directly imported) RUN pip install qwen-vl-utils -RUN pip install ray # =========================================================================== # END UPSTREAM @@ -31,19 +45,36 @@ RUN pip install ray # =========================================================================== # NVFlow Additional Layers # =========================================================================== -# Pre-cache tokenizer encodings so no downloads are needed at runtime. +# Add SAM 3.1 without disturbing vLLM's tested Torch/NumPy stack: install only +# its missing deps + source with --no-deps (SAM's numpy<2 pin would downgrade). # =========================================================================== +COPY --from=sam3-src / /opt/sam3 +RUN pip install --no-cache-dir --no-deps \ + ftfy==6.1.1 \ + iopath==0.1.10 \ + portalocker==3.2.0 \ + pycocotools==2.0.11 \ + wcwidth==0.2.14 && \ + pip install --no-cache-dir --no-deps --no-build-isolation -e /opt/sam3 && \ + rm -rf /opt/sam3/.git + +# `vllm[audio]` and `ray` above are unpinned; assert the base tag's vLLM version +# survived. Ray is only reported -- these servers run standalone. +ARG VLLM_VERSION +RUN python3 -c "import numpy, sam3, torch, vllm, ray; \ +from sam3.model_builder import build_sam3_image_model; \ +assert vllm.__version__.startswith('${VLLM_VERSION#v}'), f'vllm {vllm.__version__} != ${VLLM_VERSION#v}'; \ +print(f'SAM 3.1 + vLLM imports OK: numpy={numpy.__version__}, torch={torch.__version__}, vllm={vllm.__version__}, ray={ray.__version__}')" + +# 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. +# Fetch tiktoken encodings explicitly -- the Rust tiktoken-rs client in +# openai_harmony fails to download under QEMU arm64. Keeps the image airgapped. 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 \ @@ -54,3 +85,13 @@ RUN python3 -c "\ from openai_harmony import load_harmony_encoding, HarmonyEncodingName; \ load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS); \ print('openai_harmony encoding loaded OK')" + +# =========================================================================== +# Security hardening (Trivy/NSPECT wave scans, 2026-07-07) +# =========================================================================== +# apt upgrade for base-channel security fixes (linux-libc-dev/gnupg/openssl = 227 of 233 HIGH/CRIT Trivy); headers stay INSTALLED for vLLM triton JIT +RUN apt-get update && apt-get upgrade -y && rm -rf /var/lib/apt/lists/* + +# rm ray_dist.jar: jackson-databind RCE CVE-2026-54512/CVE-2026-54513 (Ray-Java unused); fail build if one survives +RUN find /usr/local /opt -name 'ray_dist.jar' -type f -delete 2>/dev/null; \ + ! find /usr/local /opt -name 'ray_dist.jar' -type f 2>/dev/null | grep -q . diff --git a/dockerfiles/Dockerfile.vllm-grpo b/dockerfiles/Dockerfile.vllm-grpo deleted file mode 100644 index 3ade4f7..0000000 --- a/dockerfiles/Dockerfile.vllm-grpo +++ /dev/null @@ -1,55 +0,0 @@ -# ============================================================================= -# 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 index 88abbfd..baa0e10 100644 --- a/dockerfiles/README.md +++ b/dockerfiles/README.md @@ -1,52 +1,48 @@ # 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. +NVFlow uses six worker images plus an optional launcher. Five +(`nemo-rl`, `nemo-gym`, `nemo-skills`, `vllm`, `vllm-grpo`) are **built locally** +from the self-contained Dockerfiles in this directory, as is the optional +`nvflow-client` launcher; only `sglang` is **pulled as-is**. 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 +> **Air-gap.** All six images run fully offline, including the `training` stage: +> `nemo-rl` bakes one NeMo-Gym venv per component at build time, so nothing is +> resolved or downloaded at job runtime and no Gym source mount is needed. +> `UV_OFFLINE` is left **unset** by policy β€” it keeps a dev-mode escape hatch, not +> because any stage needs the network. See +> [`docs/development/nemo-rl-gym.md`](../docs/development/nemo-rl-gym.md) for the +> trainer/Gym details. -```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 . +## Building -# 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 -``` +The full build commands β€” single-arch, multi-arch (`buildx` + QEMU), the +`sglang` pull, sanity checks, and `.sqsh` conversion β€” are in +**[docker_instructions.md](docker_instructions.md)** (the authoritative +build/deploy reference). The images and their version pins are below. ## Images | Image | Base | Purpose | |-------|------|---------| -| `nvflow-nemo-rl` | `nvcr.io/nvidia/nemo-rl:v0.6.0` | SFT, GRPO training, collect_rollouts, compute_rewards | +| `nvflow-nemo-rl` | `nvcr.io/nvidia/nemo-rl:v0.7.0` | SFT and GRPO `training`; NeMo-Gym venvs baked per component | | `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 | +| `nvflow-vllm` | `vllm/vllm-openai:v0.22.0` | Standalone vLLM inference (SDG, eval) β€” multi-arch (amd64 + arm64) | +| `nvflow-vllm` (`v0.20.0*` tag) | `vllm/vllm-openai:v0.20.0` | Standalone vLLM inference (GRPO rollouts, judge) β€” multi-arch. Same `Dockerfile.vllm`, built with `--build-arg VLLM_VERSION=v0.20.0` | +| `nvflow-nemo-gym` | `python:3.12-slim` | CPU-only Gym-only GRPO stages (prepare_data, prefetch_cache, collect_rollouts, compute_rewards); finance Gym venvs baked | +| `nvflow-client` | pinned `ubuntu:24.04` | Optional launcher: drive `nflow` over an SSH tunnel (airgap/off-cluster); Python 3.12 + CLI + venv baked | | `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 | +| `NEMO_SKILLS_COMMIT` | `e06c9b90…` (image tag `v1.1.2`) | Should match `Dockerfile.nemo-skills` and `pyproject.toml` | +| `GYM_REF` (nemo-gym, nemo-rl) | `33ef60369…` | A commit on upstream NeMo-Gym `main`; keep both images on the same one | +| `VLLM_VERSION` (vllm) | `v0.22.0` | [vLLM releases](https://github.com/vllm-project/vllm/releases) | +| `VLLM_VERSION` (vllm-grpo) | `v0.20.0` | Pinned to match NeMo-RL v0.7.0 colocated vLLM | diff --git a/dockerfiles/docker_instructions.md b/dockerfiles/docker_instructions.md index fc2a989..762c6ad 100644 --- a/dockerfiles/docker_instructions.md +++ b/dockerfiles/docker_instructions.md @@ -1,31 +1,33 @@ # 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. +air-gapped Slurm clusters. Four of them (`nemo-rl`, `nemo-skills`, `vllm`, +`vllm-grpo`) are produced by running `docker build` against the self-contained +Dockerfiles in this directory; only `sglang` is pulled as-is. All custom 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-rl` | `nvcr.io/nvidia/nemo-rl:v0.7.0` | SFT and GRPO `training`; bakes one Gym venv per component so the trainer needs no network | | `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 | +| `nvflow-vllm` | `vllm/vllm-openai:v0.22.0` | Standalone vLLM (SDG, eval) β€” multi-arch (amd64 + arm64) | +| `nvflow-vllm` (`v0.20.0*` tag) | `vllm/vllm-openai:v0.20.0` | Standalone vLLM (GRPO rollouts, judge) β€” multi-arch. Same `Dockerfile.vllm` and repo as above; `--build-arg VLLM_VERSION=v0.20.0` | | `sglang` | `lmsysorg/sglang:v0.5.10.post1` | SGLang inference server (pulled as-is, no custom Dockerfile) | +Two more images build the same way but are only needed for specific paths: +`nvflow-nemo-gym` (CPU Gym worker, `Dockerfile.nemo-gym`, pinned by `GYM_REF`) for GRPO / DG-SDG, and `nvflow-client` (optional airgap-only launcher, `Dockerfile.nvflow`, multi-arch) for driving NVFlow over an `ssh_tunnel`. Both are covered in [containers.md](../docs/maintainers/containers.md). + ## 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 | +| `NEMO_SKILLS_COMMIT` | `e06c9b90…` (image tag `v1.1.2`) | Must match `Dockerfile.nemo-skills` and `pyproject.toml` | +| `GYM_REF` (nemo-gym, nemo-rl) | `33ef60369…` | A commit on upstream NeMo-Gym `main`; keep both images on the same one | +| `VLLM_VERSION` (vllm) | `v0.22.0` | [vLLM releases](https://github.com/vllm-project/vllm/releases) | +| `VLLM_VERSION` (vllm-grpo) | `v0.20.0` | Pinned to match NeMo-RL v0.7.0 colocated vLLM | ## 1. Build @@ -45,10 +47,21 @@ below produce amd64 images. ```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 . +docker build --no-cache \ + -f dockerfiles/Dockerfile.nemo-skills -t nvflow-nemo-skills:v1.1.2 . + +# Airgapped trainer. Tag tracks the base version it extends. `docker build` pulls +# that base from nvcr.io, so `docker login nvcr.io` must have run first. +docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:v0.7.0 . + +# CPU-only Gym worker. Tag tracks the baked GYM_REF; bump it when GYM_REF moves. +docker build -f dockerfiles/Dockerfile.nemo-gym -t nvflow-nemo-gym:0.4.0 . + +# One Dockerfile builds both vLLM images; VLLM_VERSION picks the base tag, and +# both ship in the nvflow-vllm repo. +docker build -f dockerfiles/Dockerfile.vllm -t nvflow-vllm:v0.22.0 . +docker build -f dockerfiles/Dockerfile.vllm \ + --build-arg VLLM_VERSION=v0.20.0 -t nvflow-vllm:v0.20.0 . # sglang β€” pulled directly, no custom Dockerfile docker pull lmsysorg/sglang:v0.5.10.post1 @@ -70,7 +83,7 @@ Then build with an explicit `--platform`: # 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 \ + -t nvflow-vllm:v0.22.0 \ --load . ``` @@ -80,13 +93,15 @@ 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: +The custom images are all built multi-arch (see below); this single-platform +recipe is for testing one arch in isolation, and works for any of them by +swapping `-f`. From an arm64 host 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 \ + -t nvflow-vllm:v0.22.0-arm64 \ --load . ``` @@ -94,25 +109,61 @@ docker buildx build --platform linux/arm64 \ ### 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: +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. + +Needs QEMU (above) and a `docker-container` builder β€” the default `docker` +driver cannot build multiple platforms: + +```bash +docker buildx create --name nvflow --driver docker-container --use +docker buildx inspect --bootstrap +``` ```bash REGISTRY= +# Airgapped trainer. Tag tracks the base version it extends. +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.nemo-rl \ + -t $REGISTRY/nvflow-nemo-rl:v0.7.0 \ + --provenance=false --sbom=false --push . + +# CPU-only Gym worker. Tag tracks the baked GYM_REF; bump it when GYM_REF moves. +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.nemo-gym \ + -t $REGISTRY/nvflow-nemo-gym:0.4.0 \ + --provenance=false --sbom=false --push . + +docker buildx build --platform linux/amd64,linux/arm64 \ + -f dockerfiles/Dockerfile.vllm \ + -t $REGISTRY/nvflow-vllm:v0.22.0 \ + --provenance=false --sbom=false --push . + +# Same Dockerfile and repo; VLLM_VERSION picks the base tag. docker buildx build --platform linux/amd64,linux/arm64 \ -f dockerfiles/Dockerfile.vllm \ - -t $REGISTRY/nvflow-vllm:v0.18.1 \ + --build-arg VLLM_VERSION=v0.20.0 \ + -t $REGISTRY/nvflow-vllm:v0.20.0 \ + --provenance=false --sbom=false --push . + +# --no-cache is required: ARG CACHEBUST gates the dependency-override layer, so +# a warm cache reuses stale resolutions and skips the security floors. +docker buildx build --platform linux/amd64,linux/arm64 --no-cache \ + -f dockerfiles/Dockerfile.nemo-skills \ + -t $REGISTRY/nvflow-nemo-skills:v1.1.2 \ --provenance=false --sbom=false --push . +# Launcher, not a worker. Build from a committed tree: .baked_commit records +# `git rev-parse HEAD`, so uncommitted changes ship under the wrong provenance. docker buildx build --platform linux/amd64,linux/arm64 \ - -f dockerfiles/Dockerfile.vllm-grpo \ - -t $REGISTRY/nvflow-vllm-grpo:v0.17.1 \ + -f dockerfiles/Dockerfile.nvflow \ + -t $REGISTRY/nvflow-client:v1.1.2 \ --provenance=false --sbom=false --push . # Verify both architectures are in the manifest list -docker buildx imagetools inspect $REGISTRY/nvflow-vllm:v0.18.1 +docker buildx imagetools inspect $REGISTRY/nvflow-vllm:v0.22.0 ``` `--provenance=false --sbom=false` keeps the manifest list compatible with @@ -133,50 +184,46 @@ image will not work in production. ### 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/... +> Required for the default release: `nvflow-nemo-rl` is built from +> `Dockerfile.nemo-rl`, and these checks are what prove its baked Gym venvs are +> usable offline. See +> [`docs/development/nemo-rl-gym.md`](../docs/development/nemo-rl-gym.md) for the +> image internals. -# 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: ..." +```bash +IMAGE=nvflow-nemo-rl:v0.7.0 + +# A. every Gym component venv is baked +docker run --rm $IMAGE bash -c \ + 'find /opt/gym_venvs -maxdepth 3 -name .venv | sort' +# Expect: 7 paths β€” resources_servers/{equivalence_llm_judge,finance_sec_search, +# format_verification}, responses_api_agents/{finance_agent,simple_agent}, +# responses_api_models/{openai_model,vllm_model} + +# B. Gym imports with no network and no uv resolve +docker run --rm --network=none -e UV_OFFLINE=true $IMAGE bash -c \ + '/opt/ray_venvs/nemo_rl.environments.nemo_gym.NemoGym/bin/python -c \ + "import nemo_gym; print(\"nemo_gym OK\")"' +# Expect: nemo_gym OK + +# C. /opt/NeMo-RL symlink (scripts/convert_checkpoint_to_hf.sh cd's to it) +docker run --rm $IMAGE bash -c \ + 'cd /opt/NeMo-RL && ls examples/converters/convert_dcp_to_hf.py' +# Expect: examples/converters/convert_dcp_to_hf.py # 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 " + /opt/gym_venvs/resources_servers/finance_sec_search/.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 +# Expect: uvicorn 0.52.x OK ``` ### nemo-skills ```bash -IMAGE=nvflow-nemo-skills:0229040 +IMAGE=nvflow-nemo-skills:v1.1.2 # A. tiktoken pre-cache loads offline docker run --rm --network=none -e HF_HUB_OFFLINE=1 $IMAGE bash -c ' @@ -202,7 +249,7 @@ docker run --rm $IMAGE bash -c ' Run the same set against both images: ```bash -for IMAGE in nvflow-vllm:v0.18.1 nvflow-vllm-grpo:v0.17.1; do +for IMAGE in nvflow-vllm:v0.22.0 nvflow-vllm:v0.20.0; do echo "=== $IMAGE ===" # A. tiktoken encoding files present @@ -230,15 +277,17 @@ 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 +docker tag nvflow-nemo-rl:v0.7.0 $REGISTRY/nvflow-nemo-rl:v0.7.0 +docker tag nvflow-nemo-gym:0.4.0 $REGISTRY/nvflow-nemo-gym:0.4.0 +docker tag nvflow-nemo-skills:v1.1.2 $REGISTRY/nvflow-nemo-skills:v1.1.2 +docker tag nvflow-vllm:v0.22.0 $REGISTRY/nvflow-vllm:v0.22.0 +docker tag nvflow-vllm:v0.20.0 $REGISTRY/nvflow-vllm:v0.20.0 + +docker push $REGISTRY/nvflow-nemo-rl:v0.7.0 +docker push $REGISTRY/nvflow-nemo-gym:0.4.0 +docker push $REGISTRY/nvflow-nemo-skills:v1.1.2 +docker push $REGISTRY/nvflow-vllm:v0.22.0 +docker push $REGISTRY/nvflow-vllm:v0.20.0 ``` Then on the cluster (typically a CPU partition): @@ -247,11 +296,14 @@ Then on the cluster (typically a CPU partition): CONTAINER_DIR= REGISTRY= +# Name the output -.sqsh to match what scripts/setup_containers.sh +# produces, so either staging method drops in to the same my_cluster.yaml. enroot import \ - --output $CONTAINER_DIR/nvflow-nemo-rl-v0.6.0.sqsh \ - "docker://$REGISTRY/nvflow-nemo-rl:v0.6.0" + --output $CONTAINER_DIR/nemo-skills-v1.1.2.sqsh \ + "docker://$REGISTRY/nvflow-nemo-skills:v1.1.2" -# Repeat for nemo-skills, vllm, vllm-grpo, and (optionally) sglang. +# Repeat for vllm, vllm-grpo, nemo-gym, nemo-rl. sglang imports directly: +# docker://lmsysorg/sglang:v0.5.10.post1 ``` If the cluster authenticates to your registry, drop credentials into @@ -269,14 +321,14 @@ 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 +docker save nvflow-nemo-skills:v1.1.2 | gzip > nvflow-nemo-skills-v1.1.2.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 +gunzip -c nvflow-nemo-skills-v1.1.2.tar.gz | docker load enroot import \ - --output $CONTAINER_DIR/nvflow-nemo-rl-v0.6.0.sqsh \ - dockerd://nvflow-nemo-rl:v0.6.0 + --output $CONTAINER_DIR/nemo-skills-v1.1.2.sqsh \ + dockerd://nvflow-nemo-skills:v1.1.2 ``` > `enroot import` natively supports only `docker://` (remote registry), @@ -291,14 +343,14 @@ enroot import \ 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" + enroot import --output vllm-v0.22.0.sqsh \ + "docker://nvcr.io#/nvflow-vllm:v0.22.0" ``` - **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 + mv "nvflow-nemo-skills:v1.1.2.sqsh" nemo-skills-v1.1.2.sqsh ``` ## 4. Cluster config (`my_cluster.yaml`) @@ -312,10 +364,12 @@ operation. ```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 + # Filenames are -.sqsh, as produced by setup_containers.sh. + nemo-rl: /nemo-rl-v0.7.0.sqsh + nemo-skills: /nemo-skills-v1.1.2.sqsh + vllm: /vllm-v0.22.0.sqsh + vllm-grpo: /vllm-grpo-v0.20.0.sqsh + nemo-gym: /nemo-gym-0.4.0.sqsh # sglang: /sglang-v0.5.10.post1.sqsh ``` @@ -330,8 +384,10 @@ env_vars: - HF_DATASETS_OFFLINE=1 - TRANSFORMERS_OFFLINE=1 - # Disable uv package and Python interpreter downloads. - - UV_OFFLINE=true + # Disable uv package and Python interpreter downloads. Left UNSET: every image + # bakes the venvs it needs, so no stage resolves at runtime either way, and + # unset keeps a dev-mode escape hatch. + # - 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). @@ -340,35 +396,29 @@ env_vars: - 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: +### NeMo-RL / NeMo-Gym: trainer image and Gym source -```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 -``` +The `training` stage runs on `nvflow-nemo-rl`, built here from +`Dockerfile.nemo-rl`. It extends the NeMo-RL base with the Gym source at +`GYM_REF` and one prebuilt venv per Gym component under `/opt/gym_venvs`, so +nothing resolves at runtime and **no Gym source mount is required**. -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. +Do not bind-mount a host Gym or NeMo-RL clone over +`/opt/nemo-rl/3rdparty/Gym-workspace/Gym` in production β€” it shadows the baked +source and venvs and breaks the GRPO stages. That mount is a dev-mode-only tool, +and it is the one case where `uv` resolves at runtime, so it needs `UV_OFFLINE` +unset plus a reachable pypi mirror. -### Don't enable this in offline mode +The Gym-only GRPO stages (`prepare_data`, `prefetch_cache`, `collect_rollouts`, +`compute_rewards`) run on the separate `nvflow-nemo-gym` image, also with baked +per-component venvs and no mount. -```yaml -# - NRL_FORCE_REBUILD_VENVS=true # forces Ray workers to re-resolve via uv - # (requires internet; will fail under air-gap) -``` +For a fully-airgapped trainer with no runtime `uv` resolve, build the custom +image from [`Dockerfile.nemo-rl`](Dockerfile.nemo-rl) (bakes the Gym venvs) and +drop the Gym mount. That image needs no network at job time on its own, so +`UV_OFFLINE` still stays **unset** by policy: leaving it unset is what lets a +developer mount local Gym source and have `uv` resolve it. See +[`docs/development/nemo-rl-gym.md`](../docs/development/nemo-rl-gym.md). ## Notes for one-time / connected-node operations @@ -384,8 +434,8 @@ air-gapped cluster: | `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. +(`HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE`, `TRANSFORMERS_OFFLINE`). `UV_OFFLINE` +stays unset as always; none of these stages invoke `uv`. Note: `huggingface_hub` interprets `TRANSFORMERS_OFFLINE=1` as `HF_HUB_OFFLINE=1`, so all three need to be off (or unset) for HF dataset diff --git a/docs/architecture/ARCHITECTURE.md b/docs/ARCHITECTURE.md similarity index 88% rename from docs/architecture/ARCHITECTURE.md rename to docs/ARCHITECTURE.md index 1d58a20..dfc3fa3 100644 --- a/docs/architecture/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -1,7 +1,5 @@ # NVFlow Architecture -> **Version:** 1.0 -> **Last Updated:** January 21, 2026 > **Purpose:** Comprehensive architectural overview of the NVFlow orchestration framework --- @@ -361,11 +359,11 @@ sequenceDiagram NemoSkills->>Slurm: Submit job with dependencies Slurm-->>NemoSkills: Job ID NemoSkills-->>Stage: Job submitted - Stage-->>WorkflowRunner: Stage complete + Stage-->>WorkflowRunner: Stage submitted end - WorkflowRunner-->>CLI: All stages complete - CLI-->>User: βœ… Workflow Complete! + WorkflowRunner-->>CLI: All stages submitted + CLI-->>User: βœ… Workflow Submitted ``` ### Dependency Resolution @@ -437,8 +435,9 @@ graph TB TBS --> DataPrep DGS -.-> DataPrep DataPrep --> SFT + DataPrep --> GRPO SFT --> Eval - SFT --> GRPO + SFT -.-> GRPO GRPO --> Eval style SEC fill:#e3f2fd @@ -450,6 +449,8 @@ graph TB style GRPO fill:#fff3e0 ``` +Dashed edges are optional. GRPO starts from the base HuggingFace checkpoint in the shipped configs (`hf_model_path: /hf_models/Qwen/Qwen3-4B`), so SFT is not a prerequisite for it β€” point `hf_model_path` at an SFT checkpoint only if you want to chain the two. + ### Workflow Breakdown #### **Workflow 1: Download SEC Filings** @@ -479,14 +480,14 @@ Models: GPT-OSS-120B, Qwen3-14B ``` Stages: 1. dg_sdg_preprocess - Preprocess filings - 2. document_grounded_qa_generation - Generate verified Q&A - 3. genselect_answers - Self-consistency check - 4. evaluate_answers - Quality evaluation - 5. aggregate_answers - Combine results - 6. difficulty_estimation - Stratify by difficulty - 7. document_grounded_data - Prepare training data - -Output: ~800K Q&A pairs (stratified) + 2. generate_verified_questions - Question generation + verification + 3. generate_answers - Answer generation (multi-rollout) + 4. gym_genselect_answers - Self-consistency selection + 5. evaluate_answers - Quality evaluation + 6. aggregate_answers - Combine multi-seed results + 7. dgsdg_post_process - Clean + rename β†’ final_result.jsonl + +Output: ~800K Q&A pairs in final_result.jsonl GPU: 8 GPUs Models: Qwen3 family (14B-235B) ``` @@ -497,13 +498,15 @@ Stages: 1. data_transformation - Convert to training format 2. prepare_for_sft - Format for NeMo 3. train_validation_split - Split dataset - 4. sequence_length_grouping - [Optional] Group by length + 4. sequence_length_grouping - Group by length 5. training - Multi-node training - 6. convert_to_messages - [Optional] Post-processing + 6. eval - Evaluate checkpoints on benchmarks + + Qwen3 configs insert convert_to_messages between training and eval. GPU: 256 GPUs (32 nodes Γ— 8 GPUs) Model: Qwen3-14B -Parallelism: TP=2, PP=1, CP=2 +Parallelism: TP=4, PP=1, CP=8 ``` #### **Workflow 5: Evaluation** @@ -521,7 +524,7 @@ Benchmarks: Financial reasoning tasks #### **Workflow 6: GRPO RL Training** ``` Stages: - 1. validate_questions - Validate format + deduplicate + 1. validate_questions - Regex prefilter + LLM validity classifier 2. data_transformation - SDG cleanup to model-agnostic schema 3. apply_prompt_template - Apply prompt template + extract answer 4. convert_to_responses_api - Convert to NeMo-Gym Responses API format @@ -534,7 +537,7 @@ Stages: Output: RL-trained model + eval results GPU: 16 GPUs (2 nodes for demo), 64 GPUs (8 nodes for production) -Model: Qwen3-4B dense (demo, FSDP v2), Qwen3-30B-A3B MoE (production, Megatron) +Model: Qwen3-4B dense (demo β€” equivalence_llm_judge: FSDP v2 @ 32K; finance_sec_search: Megatron TP2Γ—CP8 @ 131K), Qwen3-30B-A3B MoE (production, Megatron) ``` ### Finance Recipe Component Diagram @@ -544,12 +547,14 @@ graph TB subgraph "Finance Recipe Structure" direction TB - subgraph "Stages (42 total)" + subgraph "Stage modules (23 total)" direction LR - SDG[SDG Stages
12 stages] - SFT[SFT Stages
4 stages] - Eval[Eval Stages
2 stages] - RL[RL Stages
10 stages] + SDG[SDG
6 modules] + RL[RL
8 modules] + SFT[SFT
4 modules] + Eval[Eval
2 modules] + DL[Download
1 module] + Shared[Shared
2 modules] end subgraph "Workflows (6 total)" @@ -557,7 +562,7 @@ graph TB W2[template-sdg
6 stages] W3[document-sdg
7 stages] W4[sft
6 stages] - W5[eval
9 stages] + W5[eval
7 stages] W6[grpo
10 stages] end @@ -567,12 +572,14 @@ graph TB P3[Evaluation Prompts] end - W1 -.-> SDG + W1 -.-> DL W2 -.-> SDG W3 -.-> SDG W4 -.-> SFT + W4 -.-> Shared W5 -.-> Eval W6 -.-> RL + W6 -.-> Shared SDG -.-> P1 SDG -.-> P2 @@ -583,6 +590,8 @@ graph TB style SFT fill:#bbdefb style Eval fill:#f8bbd0 style RL fill:#fff3e0 + style DL fill:#ede7f6 + style Shared fill:#eceff1 style W1 fill:#e1f5ff style W2 fill:#e1f5ff style W3 fill:#e1f5ff @@ -669,35 +678,41 @@ graph TB ``` β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ Air-gapped container images (.sqsh format) β”‚ -β”‚ (built locally from dockerfiles/Dockerfile.* β€” see INSTALL.md)β”‚ +β”‚ Container images (.sqsh format) β”‚ +β”‚ (nemo-rl/nemo-skills/vllm/vllm-grpo/nemo-gym built from β”‚ +β”‚ dockerfiles/*; only sglang pulled as-is β€” see INSTALL.md) β”‚ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”‚ β”‚ β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ β”‚ β”‚ nvflow-nemo-skills β”‚ β”‚ nvflow-vllm β”‚ β”‚ -β”‚ β”‚ (0229040) β”‚ β”‚ (v0.18.1) β”‚ β”‚ +β”‚ β”‚ (v1.1.2) β”‚ β”‚ (v0.22.0) β”‚ β”‚ β”‚ β”‚ 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 β”‚ β”‚ +β”‚ β”‚ (v0.5.10.post1) β”‚ β”‚ (v0.7.0) β”‚ β”‚ +β”‚ β”‚ SDG inference β”‚ β”‚ SFT/GRPO trainer β”‚ β”‚ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ β”‚ β”‚ -β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ -β”‚ β”‚ nvflow-vllm-grpo β”‚ β”‚ -β”‚ β”‚ (v0.17.1) β”‚ β”‚ -β”‚ β”‚ GRPO rollout/judge β”‚ β”‚ -β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ +β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ +β”‚ β”‚ nvflow-vllm β”‚ β”‚ nvflow-nemo-gym β”‚ β”‚ +β”‚ β”‚ (v0.20.0 tag) β”‚ β”‚ (0.4.0) β”‚ β”‚ +β”‚ β”‚ GRPO rollout/judge β”‚ β”‚ CPU Gym-only stages β”‚ β”‚ +β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +The five `nvflow-*` images each require a `docker build` on a connected host before use, then run fully offline; only `sglang` can be pulled directly. See [`dockerfiles/docker_instructions.md`](../dockerfiles/docker_instructions.md). + +``` ↓ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ Shared Filesystem Mounts β”‚ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ /workspace β†’ /lustre/.../workspace β”‚ +β”‚ /workspace β†’ writable data (outputs, HF cache) β”‚ β”‚ /hf_models β†’ /lustre/.../models/hf_models β”‚ -β”‚ /outputs β†’ /lustre/.../outputs β”‚ +β”‚ /nemo_run/code β†’ recipe code (nemo-run packaged) β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ ``` @@ -876,7 +891,7 @@ graph TB ### Storage Layout ``` -/lustre/fsw/.../workspace/nvflow/ +/workspace/nvflow/ β”‚ β”œβ”€β”€ cluster_configs/ # Cluster configuration β”‚ β”œβ”€β”€ containers.yaml # Container definitions @@ -921,26 +936,8 @@ graph TB ## Summary -### Key Architectural Highlights - -1. **Modular Design**: Clear separation between framework, recipes, and infrastructure -2. **Hierarchical Organization**: Recipe β†’ Workflow β†’ Stage provides natural organization -3. **Declarative Configuration**: YAML-based configs with inheritance support -4. **Flexible Execution**: CLI, Python API, and programmatic interfaces -5. **Cluster Native**: First-class Slurm integration with dependency management -6. **Extensible**: Easy to add new recipes, workflows, and stages -7. **Built on NeMo**: Leverages NVIDIA's NeMo ecosystem (Skills, RL, Framework) - -### Design Benefits - -- **Reproducibility**: Version-controlled configs and deterministic execution -- **Reusability**: Stages can be shared across workflows and recipes -- **Scalability**: Seamless scaling from local development to multi-node clusters -- **Maintainability**: Clear structure and separation of concerns -- **Discoverability**: Registry pattern enables stage discovery and documentation - ---- +Three ideas carry most of the design: -**Document Version:** 1.0 -**Generated:** January 21, 2026 -**Repository:** nvflow +- **Recipe β†’ Workflow β†’ Stage.** Stages are the unit of reuse and are discovered through the registry, so they can be shared across workflows and recipes, and listed without being hardcoded anywhere. +- **Declarative, inheritable YAML.** A model config inherits a workflow base and patches it, which keeps runs version-controlled and reproducible. +- **Slurm-native submission.** Stages submit jobs with dependencies rather than executing inline, which is what lets the same config scale from a demo to a multi-node production run. diff --git a/docs/architecture/ARCHITECTURE_INDEX.md b/docs/architecture/ARCHITECTURE_INDEX.md deleted file mode 100644 index dc4dbd1..0000000 --- a/docs/architecture/ARCHITECTURE_INDEX.md +++ /dev/null @@ -1,435 +0,0 @@ -# NVFlow - Architecture Documentation Index - -> **Navigation guide for all architecture documentation** -> **Start here to find the right documentation for your needs** - ---- - -## πŸ“š Documentation Overview - -The NVFlow architecture is documented across multiple files, each serving a specific purpose. This index helps you find the right documentation quickly. - ---- - -## 🎯 Quick Navigation - -### I want to... - -| Goal | Document | Time | -|------|----------|------| -| **Get a quick overview** | [ARCHITECTURE_QUICK_REFERENCE.md](#quick-reference) | 5 min | -| **Understand the system deeply** | [ARCHITECTURE.md](#comprehensive-architecture) | 30 min | -| **View visual diagrams** | [diagrams/](#visual-diagrams) | 10 min | -| **Learn about diagrams** | [DIAGRAMS_SUMMARY.md](#diagrams-summary) | 10 min | -| **Get started with NVFlow** | [README.md](#main-readme) | 15 min | -| **Set up the cluster** | [INSTALL.md](#installation-guide) | 30 min | -| **Learn the finance recipe** | [docs/recipes/finance/](#finance-recipe-docs) | 45 min | - ---- - -## πŸ“– Document Descriptions - -### Quick Reference -**File:** [ARCHITECTURE_QUICK_REFERENCE.md](./ARCHITECTURE_QUICK_REFERENCE.md) -**Size:** ~5 KB -**Reading Time:** 5 minutes -**Best For:** Quick onboarding, cheat sheet, reference card - -**Contents:** -- One-page architecture overview -- Core components table -- Common commands -- Quick stage creation guide -- Key features checklist -- Documentation map - -**When to Use:** -- First time learning about NVFlow -- Need a quick reminder of concepts -- Looking for specific commands -- Want a printable reference - ---- - -### Comprehensive Architecture -**File:** [ARCHITECTURE.md](./ARCHITECTURE.md) -**Size:** ~26 KB -**Reading Time:** 30 minutes -**Best For:** Deep understanding, system design, contribution - -**Contents:** -1. High-level architecture with diagrams -2. System overview and design principles -3. Core framework components (detailed) -4. Hierarchical organization (Recipe β†’ Workflow β†’ Stage) -5. Complete execution flow with sequence diagrams -6. Finance recipe architecture (all 42 stages) -7. Deployment architecture and topology -8. Technology stack and integrations -9. Data flow diagrams - -**When to Use:** -- Need comprehensive system understanding -- Planning to contribute to the codebase -- Designing new recipes or workflows -- Troubleshooting complex issues -- Presenting architecture to stakeholders - ---- - -### Diagrams Summary -**File:** [DIAGRAMS_SUMMARY.md](./DIAGRAMS_SUMMARY.md) -**Size:** ~12 KB -**Reading Time:** 10 minutes -**Best For:** Understanding available diagrams, diagram usage guide - -**Contents:** -- Overview of all 5 diagrams -- Diagram details and use cases -- Audience-specific recommendations -- Question-to-diagram mapping -- Diagram statistics -- Rendering examples -- Update guidelines - -**When to Use:** -- Want to know what diagrams are available -- Need to choose the right diagram -- Want to render diagrams in different formats -- Planning to create new diagrams - ---- - -### Visual Diagrams -**Location:** [diagrams/](../diagrams/) -**Format:** Mermaid (.mmd files) -**Count:** 5 diagrams + README -**Best For:** Visual learners, presentations, documentation - -**Available Diagrams:** - -1. **[architecture-overview.mmd](../diagrams/architecture-overview.mmd)** - - High-level system architecture - - All major components and relationships - - 5 layers: UI, Core, Recipes, Infrastructure, Storage - -2. **[finance-pipeline.mmd](../diagrams/finance-pipeline.mmd)** - - Complete finance recipe pipeline - - All 6 workflows with 42 stages - - Data flow from SEC filings to evaluation - -3. **[execution-flow.mmd](../diagrams/execution-flow.mmd)** - - Runtime execution sequence diagram - - User command to job completion - - 4 phases: Init, Validate, Execute, Monitor - -4. **[component-architecture.mmd](../diagrams/component-architecture.mmd)** - - Class diagram of core framework - - BaseStage, StageRegistry, WorkflowRunner - - Relationships and dependencies - -5. **[deployment-architecture.mmd](../diagrams/deployment-architecture.mmd)** - - Infrastructure and deployment topology - - Local machine to Slurm cluster - - Compute nodes, storage, containers - -**Viewing Options:** -- Online: https://mermaid.live/ -- VS Code: Mermaid Preview extension -- CLI: `mmdc -i diagram.mmd -o diagram.png` -- GitHub: Native rendering - -**When to Use:** -- Need visual understanding -- Creating presentations -- Onboarding new team members -- Documentation in other systems - ---- - -### Main README -**File:** [README.md](../../README.md) -**Size:** ~10 KB -**Reading Time:** 15 minutes -**Best For:** Getting started, understanding concepts, running workflows - -**Contents:** -- Project overview and key features -- Core concepts (Recipe, Workflow, Stage) -- Folder structure explanation -- Installation instructions -- Quick start examples -- CLI commands reference -- Development guide - -**When to Use:** -- First time using NVFlow -- Need to understand basic concepts -- Want to run your first workflow -- Looking for CLI command syntax - ---- - -### Installation Guide -**File:** [INSTALL.md](../../INSTALL.md) -**Size:** ~8 KB -**Reading Time:** 30 minutes (including setup) -**Best For:** Cluster setup, container configuration, troubleshooting - -**Contents:** -1. Prerequisites (uv, yq, enroot) -2. Container setup (automated script) -3. Model download instructions -4. Cluster configuration -5. Verification steps -6. Troubleshooting guide - -**When to Use:** -- Setting up NVFlow for the first time -- Configuring a new cluster -- Troubleshooting installation issues -- Understanding container requirements - ---- - -### Finance Recipe Docs -**Location:** [docs/recipes/finance/](../recipes/finance/) -**Size:** Multiple files (~20 KB total) -**Reading Time:** 45 minutes -**Best For:** Understanding finance recipe, running production pipelines - -**Main Files:** - -1. **[README.md](../recipes/finance/README.md)** - Recipe overview - - 6 workflows, 42 stages - - Pipeline architecture - - Getting started guide - - Command reference - -2. **[quick-start.md](../recipes/finance/quick-start.md)** - 30-min demo - - Hands-on tutorial with 7 companies - - Step-by-step instructions - - Expected outputs - -3. **Workflow Guides** (in `workflows/`) - - 01-download-sec.md - - 02-template-based-sdg.md - - 03-document-grounded-sdg.md - - 04-sft.md - - 05-eval.md - - 06-grpo.md - - 06-finance-agent-eval.md - -4. **Stage Reference** (in `stages/`) - - Technical specifications for all 42 stages - - Input/output formats - - Configuration options - -5. **[troubleshooting.md](../recipes/finance/troubleshooting.md)** - - Common issues and solutions - - Debugging tips - -**When to Use:** -- Running the finance recipe -- Understanding SDG approaches -- Training financial reasoning models -- Troubleshooting finance-specific issues - ---- - -## πŸ—ΊοΈ Documentation Map (Visual) - -``` -NVFlow Documentation -β”‚ -β”œβ”€ πŸ“˜ Getting Started -β”‚ β”œβ”€ README.md ...................... Project overview & quick start -β”‚ β”œβ”€ INSTALL.md ..................... Cluster setup guide -β”‚ └─ ARCHITECTURE_QUICK_REFERENCE.md One-page cheat sheet -β”‚ -β”œβ”€ πŸ—οΈ Architecture -β”‚ β”œβ”€ ARCHITECTURE.md ................ Comprehensive architecture (26 KB) -β”‚ β”œβ”€ DIAGRAMS_SUMMARY.md ............ Diagram usage guide -β”‚ β”œβ”€ ARCHITECTURE_INDEX.md .......... This file -β”‚ └─ diagrams/ ...................... Visual diagrams (5 files) -β”‚ β”œβ”€ architecture-overview.mmd -β”‚ β”œβ”€ finance-pipeline.mmd -β”‚ β”œβ”€ execution-flow.mmd -β”‚ β”œβ”€ component-architecture.mmd -β”‚ β”œβ”€ deployment-architecture.mmd -β”‚ └─ README.md -β”‚ -β”œβ”€ 🍴 Recipes -β”‚ β”œβ”€ docs/recipes/finance/ .......... Finance recipe (production) - β”‚ β”‚ β”œβ”€ README.md ................... Recipe overview - β”‚ β”‚ β”œβ”€ quick-start.md .............. 30-min demo - β”‚ β”‚ β”œβ”€ workflows/ .................. 7 workflow guides - β”‚ β”‚ β”œβ”€ stages/ ..................... 42 stage specifications -β”‚ β”‚ └─ troubleshooting.md .......... Common issues -β”‚ β”‚ -β”‚ └─ nvflow/recipes/example/ ..... Example recipe (learning) -β”‚ -β”œβ”€ πŸ’» Code & development docs -β”‚ β”œβ”€ docs/development/console-ui.md .. Console UI guide (stage terminal output) -β”‚ └─ tests/README.md ................ Testing guide -β”‚ -└─ πŸ”§ Configuration - β”œβ”€ cluster_configs/ ............... Cluster configuration files - └─ pyproject.toml ................. Project dependencies -``` - ---- - -## πŸ‘₯ Audience-Specific Paths - -### For New Users -1. Start: [README.md](../../README.md) - Understand what NVFlow is -2. Quick ref: [ARCHITECTURE_QUICK_REFERENCE.md](./ARCHITECTURE_QUICK_REFERENCE.md) - Key concepts -3. Visual: [diagrams/architecture-overview.mmd](../diagrams/architecture-overview.mmd) - See the big picture -4. Try it: [docs/recipes/finance/quick-start.md](../recipes/finance/quick-start.md) - Run first workflow - -### For Data Scientists -1. Overview: [README.md](../../README.md) - Core concepts -2. Pipeline: [diagrams/finance-pipeline.mmd](../diagrams/finance-pipeline.mmd) - See data flow -3. Recipe: [docs/recipes/finance/README.md](../recipes/finance/README.md) - Finance pipeline -4. Run: [docs/recipes/finance/quick-start.md](../recipes/finance/quick-start.md) - Hands-on demo - -### For ML Engineers -1. Setup: [INSTALL.md](../../INSTALL.md) - Cluster configuration -2. Architecture: [ARCHITECTURE.md](./ARCHITECTURE.md) - System design -3. Execution: [diagrams/execution-flow.mmd](../diagrams/execution-flow.mmd) - Runtime behavior -4. Troubleshoot: [docs/recipes/finance/troubleshooting.md](../recipes/finance/troubleshooting.md) - -### For Software Engineers / Stage Authors -1. Components: [diagrams/component-architecture.mmd](../diagrams/component-architecture.mmd) - Class structure -2. Deep dive: [ARCHITECTURE.md](./ARCHITECTURE.md) - Design patterns -3. Code: Browse `nvflow/core/` - Framework implementation -4. Extend: [README.md](../../README.md#-creating-a-stage) - Create new stages -5. Console UI: [docs/development/console-ui.md](../development/console-ui.md) - Terminal output in stage `execute()` methods - -### For DevOps/Infrastructure -1. Setup: [INSTALL.md](../../INSTALL.md) - Installation guide -2. Deployment: [diagrams/deployment-architecture.mmd](../diagrams/deployment-architecture.mmd) - Topology -3. Cluster: [ARCHITECTURE.md](./ARCHITECTURE.md#7-deployment-architecture) - Infrastructure details -4. Config: `cluster_configs/` - Configuration files - -### For System Architects -1. Overview: [ARCHITECTURE_QUICK_REFERENCE.md](./ARCHITECTURE_QUICK_REFERENCE.md) - Quick scan -2. Complete: [ARCHITECTURE.md](./ARCHITECTURE.md) - Full architecture -3. All diagrams: [diagrams/](../diagrams/) - Visual representations -4. Design: [ARCHITECTURE.md](./ARCHITECTURE.md#2-system-overview) - Design principles - ---- - -## πŸ” Finding Specific Information - -### Concepts & Terminology -- **Recipe, Workflow, Stage:** [README.md](../../README.md#-core-concepts) -- **Hierarchical organization:** [ARCHITECTURE.md](./ARCHITECTURE.md#4-hierarchical-organization) -- **Design patterns:** [ARCHITECTURE.md](./ARCHITECTURE.md#key-design-patterns) - -### How-To Guides -- **Create a stage:** [README.md](../../README.md#-creating-a-stage) -- **Console output in stages:** [docs/development/console-ui.md](../development/console-ui.md) - Use `console.status()`, `console.detail()`, etc. -- **Run a workflow:** [README.md](../../README.md#-quick-start) -- **Set up cluster:** [INSTALL.md](../../INSTALL.md) -- **Run finance recipe:** [docs/recipes/finance/quick-start.md](../recipes/finance/quick-start.md) - -### Technical Reference -- **CLI commands:** [README.md](../../README.md#-cli-commands) -- **Core components:** [ARCHITECTURE.md](./ARCHITECTURE.md#3-core-framework-components) -- **Finance stages:** [docs/recipes/finance/stages/](../recipes/finance/stages/) -- **API reference:** Code docstrings in `nvflow/core/` - -### Visual Diagrams -- **System overview:** [diagrams/architecture-overview.mmd](../diagrams/architecture-overview.mmd) -- **Data pipeline:** [diagrams/finance-pipeline.mmd](../diagrams/finance-pipeline.mmd) -- **Execution flow:** [diagrams/execution-flow.mmd](../diagrams/execution-flow.mmd) -- **Class structure:** [diagrams/component-architecture.mmd](../diagrams/component-architecture.mmd) -- **Infrastructure:** [diagrams/deployment-architecture.mmd](../diagrams/deployment-architecture.mmd) - ---- - -## πŸ“Š Documentation Statistics - -| Metric | Count | -|--------|-------| -| Total documentation files | 20+ | -| Architecture documents | 4 | -| Visual diagrams | 5 | -| Recipe guides | 10+ | -| Total pages (estimated) | 100+ | -| Total size | ~100 KB | - ---- - -## πŸ”„ Documentation Maintenance - -### When to Update - -| Change Type | Documents to Update | -|-------------|-------------------| -| New recipe | Architecture overview, diagrams | -| New stage | Recipe docs, pipeline diagram | -| Core framework change | ARCHITECTURE.md, component diagram | -| Infrastructure change | INSTALL.md, deployment diagram | -| New workflow | Recipe README, workflow guide | - -### Update Checklist - -- [ ] Update relevant markdown files -- [ ] Update diagrams if visual changes -- [ ] Test diagram rendering -- [ ] Update this index if new docs added -- [ ] Update README if major changes -- [ ] Verify all links still work - ---- - -## πŸ“ž Getting Help - -- **Documentation issues:** Check this index for the right document -- **Architecture questions:** See [ARCHITECTURE.md](./ARCHITECTURE.md) -- **Setup problems:** See [INSTALL.md](../../INSTALL.md) troubleshooting -- **Recipe issues:** See recipe-specific troubleshooting guides -- **Code questions:** Check code docstrings and comments - ---- - -## 🀝 Contributing to Documentation - -1. **For typos/small fixes:** Edit the relevant file directly -2. **For new diagrams:** Add to `diagrams/` and update `DIAGRAMS_SUMMARY.md` -3. **For new sections:** Update relevant docs and this index -4. **For new recipes:** Create recipe docs following finance recipe structure - -**Style Guide:** -- Use clear, concise language -- Include code examples where helpful -- Add diagrams for complex concepts -- Keep this index updated -- Test all commands before documenting - ---- - -## πŸ“„ License - -All documentation is part of the NVFlow project and follows the Apache-2.0 license. - ---- - -**Version:** 1.0 -**Last Updated:** January 21, 2026 -**Maintained by:** NVFlow Team - ---- - -## πŸš€ Next Steps - -1. **New to NVFlow?** β†’ Start with [README.md](../../README.md) -2. **Need quick reference?** β†’ See [ARCHITECTURE_QUICK_REFERENCE.md](./ARCHITECTURE_QUICK_REFERENCE.md) -3. **Want deep understanding?** β†’ Read [ARCHITECTURE.md](./ARCHITECTURE.md) -4. **Visual learner?** β†’ Browse [diagrams/](../diagrams/) -5. **Ready to run?** β†’ Follow [docs/recipes/finance/quick-start.md](../recipes/finance/quick-start.md) - -**Happy learning! πŸŽ‰** diff --git a/docs/architecture/ARCHITECTURE_QUICK_REFERENCE.md b/docs/architecture/ARCHITECTURE_QUICK_REFERENCE.md deleted file mode 100644 index b40b399..0000000 --- a/docs/architecture/ARCHITECTURE_QUICK_REFERENCE.md +++ /dev/null @@ -1,276 +0,0 @@ -# NVFlow - Architecture Quick Reference - -> **One-page overview of NVFlow architecture** -> **For:** Quick onboarding and reference -> **See also:** [ARCHITECTURE.md](./ARCHITECTURE.md) for comprehensive details - ---- - -## πŸ—οΈ System Architecture (3 Layers) - -``` -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ USER LAYER: CLI, Python API, Scripts β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ FRAMEWORK LAYER: WorkflowRunner, StageRegistry β”‚ -β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ EXECUTION LAYER: NeMo-Skills, Slurm, Containers β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ -``` - ---- - -## πŸ“¦ Core Components - -| Component | Purpose | Key Methods | -|-----------|---------|-------------| -| **BaseStage** | Abstract base for all stages | `execute()`, `validate_config()` | -| **StageRegistry** | Hierarchical stage registry | `register()`, `get()`, `list_*()` | -| **WorkflowRunner** | Orchestrates workflow execution | `run()`, `validate_config()` | -| **Console** | Rich terminal UI | `header()`, `info()`, `success()` | - ---- - -## 🎯 Hierarchical Organization - -``` -Recipe (Domain: finance, healthcare, retail) - ↓ -Workflow (Pipeline: download, sdg, sft, eval, grpo) - ↓ -Stage (Task: generate_answers, training, evaluate) -``` - -**Example Path:** `finance.sft.sft` β†’ `SFTStage` class - ---- - -## πŸ“Š Finance Recipe Pipeline (6 Workflows, 42 Stages) - -``` -1. download-sec (1 stage) - └─ Download SEC filings β†’ ~100GB JSON - -2. template-based-sdg (6 stages) [PRODUCTION] - └─ Seed β†’ Questions β†’ Context β†’ Answers β†’ Filter β†’ ~300K Q&A - -3. document-grounded-sdg (7 stages) [EXPERIMENTAL] - └─ Preprocess β†’ Generate β†’ Evaluate β†’ ~800K Q&A - -4. sft (4 stages + 2 shared) - └─ Transform β†’ Prepare β†’ Split β†’ Train β†’ Checkpoints - -5. eval (2 stages, dynamically expanded) - └─ Prepare β†’ Evaluate checkpoints β†’ Compare β†’ Results - -6. grpo (10 stages: 9 active + 1 optional) - └─ GRPO reinforcement learning workflow -``` - ---- - -## πŸš€ Execution Flow (4 Phases) - -``` -1. INIT: Load config β†’ Resolve inheritance β†’ Extract context -2. VALIDATE: Check registry β†’ Validate stages -3. EXECUTE: For each stage β†’ Submit to Slurm β†’ Track dependencies -4. MONITOR: Check status β†’ View logs β†’ Collect results -``` - ---- - -## πŸ’» Technology Stack - -```yaml -Core: - - Python 3.12+, OmegaConf, Typer, Rich - -Execution: - - NeMo-Skills (SDG & pipelines) - - Slurm (cluster scheduling) - - Enroot (containers) - -Models: - - vLLM, SGLang (inference) - - HuggingFace (model loading) -``` - ---- - -## πŸ—‚οΈ Directory Structure - -``` -nvflow/ -β”œβ”€β”€ core/ # Framework (BaseStage, Registry, Runner) -β”œβ”€β”€ cli/ # CLI interface (nflow commands) -└── recipes/ # Domain-specific implementations - β”œβ”€β”€ finance/ # 42 stages, 6 workflows - β”‚ β”œβ”€β”€ stages/ # Stage implementations - β”‚ β”œβ”€β”€ workflows/ # YAML configs - β”‚ └── prompts/ # Prompt templates - └── example/ # Learning & testing -``` - ---- - -## πŸ”§ Common Commands - -```bash -# List all stages -nflow list-stages --recipe finance - -# Get stage info -nflow stage-info finance.sft.sft - -# Run single stage -nflow run sft --config workflow.yaml - -# Run all stages -nflow run-all --config workflow.yaml - -# Validate config -nflow validate --config workflow.yaml -``` - ---- - -## πŸ“ Creating a New Stage (3 Steps) - -```python -# 1. Create stage file: nvflow/recipes/finance/stages/sdg/my_stage.py -from nvflow.core import BaseStage, StageRegistry - -# 2. Implement with decorator -@StageRegistry.register( - recipe="finance", - workflow="my_workflow", - stage="my_stage" -) -class MyStage(BaseStage): - workflow = "my_workflow" - - def execute(self, config, cluster, expname, run_after=None): - # Your implementation - pass -``` - -```yaml -# 3. Add to workflow YAML -recipe: finance -workflow: - name: my_workflow -pipeline_stages: - - my_stage -stages: - my_stage: - # Your config -``` - ---- - -## 🎨 Design Patterns - -| Pattern | Usage | Example | -|---------|-------|---------| -| **Template Method** | BaseStage defines interface | `execute()` method | -| **Registry** | Stage discovery | `StageRegistry.get()` | -| **Decorator** | Stage registration | `@StageRegistry.register()` | -| **Strategy** | Execution modes | Local vs. Slurm | -| **Dependency Injection** | Config passing | `execute(config, cluster, ...)` | - ---- - -## 🏭 Deployment Topology - -``` -Local Machine Slurm Cluster -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ nflow CLI │───SSH───► β”‚ Login Node β”‚ -β”‚ Config YAML β”‚ β”‚ ↓ β”‚ -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ Slurm Scheduler β”‚ - β”‚ ↓ β”‚ - β”‚ Compute Nodes β”‚ - β”‚ β€’ 8Γ— H100 GPUs β”‚ - β”‚ β€’ Enroot containers β”‚ - β”‚ β€’ Shared storage β”‚ - β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ -``` - ---- - -## πŸ“Š Data Flow (Finance Recipe) - -``` -SEC API β†’ Filings (100GB) β†’ SDG (300K Q&A) β†’ -Data Prep β†’ Training β†’ Checkpoints β†’ Evaluation β†’ Results -``` - ---- - -## πŸ”‘ Key Features - -βœ… **Modular**: Reusable stages across workflows -βœ… **Declarative**: YAML-based configuration -βœ… **Scalable**: Local to multi-node clusters -βœ… **Reproducible**: Version-controlled configs -βœ… **Extensible**: Easy to add recipes/stages -βœ… **Built on NeMo**: Leverages NVIDIA ecosystem - ---- - -## πŸ“š Documentation Map - -| Document | Purpose | Audience | -|----------|---------|----------| -| **README.md** | Getting started | All users | -| **INSTALL.md** | Cluster setup | DevOps, ML Engineers | -| **ARCHITECTURE.md** | Deep dive (26 KB) | Architects, Contributors | -| **ARCHITECTURE_QUICK_REFERENCE.md** | This page | Quick reference | -| **DIAGRAMS_SUMMARY.md** | Diagram guide | Visual learners | -| **diagrams/*.mmd** | Visual diagrams | All users | -| **docs/recipes/finance/** | Finance recipe | Data Scientists | - ---- - -## 🎯 Use Case: Finance Recipe - -**Goal:** Generate synthetic financial Q&A data and train reasoning models - -**Input:** SEC filings (10-K, 10-Q, 8-K) - -**Process:** -1. Download filings (S&P 500) -2. Generate 300K Q&A pairs (template-based SDG) -3. Prepare training data -4. Fine-tune Qwen3-14B (256 GPUs) -5. Evaluate on benchmarks - -**Output:** Fine-tuned financial reasoning model + evaluation metrics - -**Scale:** ~100GB data β†’ 300K Q&A β†’ 256 GPU training β†’ Production model - ---- - -## πŸ”— Quick Links - -- **Full Architecture:** [ARCHITECTURE.md](./ARCHITECTURE.md) -- **Diagrams:** [diagrams/](../diagrams/) -- **Finance Recipe:** [docs/recipes/finance/README.md](../recipes/finance/README.md) -- **Quick Start:** [docs/recipes/finance/quick-start.md](../recipes/finance/quick-start.md) -- **NeMo-Skills:** https://github.com/NVIDIA/NeMo-Skills - ---- - -## πŸ’‘ Tips - -1. **Start with example recipe** for learning -2. **Use `nflow list-stages`** to discover stages -3. **Check `nflow stage-info`** for stage details -4. **Validate configs** before running: `nflow validate` -5. **Monitor jobs** with `squeue` and log files -6. **Pre-download models** to avoid GPU time waste - ---- - -**Version:** 1.0 | **Updated:** Jan 21, 2026 | **License:** Apache-2.0 diff --git a/docs/architecture/DIAGRAMS_SUMMARY.md b/docs/architecture/DIAGRAMS_SUMMARY.md deleted file mode 100644 index 04c20e2..0000000 --- a/docs/architecture/DIAGRAMS_SUMMARY.md +++ /dev/null @@ -1,310 +0,0 @@ -# NVFlow - Architecture Diagrams Summary - -> **Created:** January 21, 2026 -> **Purpose:** Quick reference guide for all architectural diagrams - ---- - -## πŸ“¦ What's Included - -A comprehensive set of architectural diagrams and documentation for the NVFlow orchestration framework has been created: - -### πŸ“„ Main Documentation -- **[ARCHITECTURE.md](./ARCHITECTURE.md)** - Complete architectural documentation (26 KB) - - High-level architecture overview - - System components and design patterns - - Hierarchical organization (Recipe β†’ Workflow β†’ Stage) - - Execution flow and dependency management - - Finance recipe detailed architecture - - Deployment topology - - Technology stack - - Data flow diagrams - -### πŸ“Š Mermaid Diagrams (in `diagrams/` folder) - -1. **[architecture-overview.mmd](../diagrams/architecture-overview.mmd)** - High-level system architecture - - Shows all major components and their relationships - - User interfaces β†’ Core framework β†’ Recipes β†’ Infrastructure β†’ Storage - -2. **[finance-pipeline.mmd](../diagrams/finance-pipeline.mmd)** - Finance recipe end-to-end pipeline - - Complete data flow from SEC filings to model evaluation - - All 6 workflows with 42 stages visualized - - Production vs. experimental paths - -3. **[execution-flow.mmd](../diagrams/execution-flow.mmd)** - Runtime execution sequence - - Step-by-step workflow execution - - User command β†’ Config loading β†’ Stage execution β†’ Job submission - - Background Slurm job processing - -4. **[component-architecture.mmd](../diagrams/component-architecture.mmd)** - Class diagram - - Core framework classes and relationships - - BaseStage, StageRegistry, WorkflowRunner - - Concrete stage implementations - - External dependencies - -5. **[deployment-architecture.mmd](../diagrams/deployment-architecture.mmd)** - Infrastructure view - - Local development environment - - Slurm cluster topology - - Compute nodes, storage, containers - - Network connections and data flow - -### πŸ“– Diagram Documentation -- **[diagrams/README.md](../diagrams/README.md)** - Guide for viewing and editing diagrams - - Description of each diagram - - Multiple viewing options (online, VS Code, CLI, GitHub) - - Mermaid syntax reference - - Style guide and contribution guidelines - ---- - -## πŸš€ Quick Start Guide - -### Viewing the Architecture - -**Option 1: Read the comprehensive documentation** -```bash -cat ARCHITECTURE.md -# or open in your favorite markdown viewer -code ARCHITECTURE.md -``` - -**Option 2: View diagrams online** -1. Visit https://mermaid.live/ -2. Open any `.mmd` file from `diagrams/` -3. Copy-paste the content -4. View and export as needed - -**Option 3: Generate PNG images** -```bash -cd diagrams/ - -# Install Mermaid CLI if not already installed -npm install -g @mermaid-js/mermaid-cli - -# Generate all diagrams as PNG -mmdc -i architecture-overview.mmd -o architecture-overview.png -mmdc -i finance-pipeline.mmd -o finance-pipeline.png -mmdc -i execution-flow.mmd -o execution-flow.png -mmdc -i component-architecture.mmd -o component-architecture.png -mmdc -i deployment-architecture.mmd -o deployment-architecture.png -``` - -**Option 4: VS Code with Mermaid Preview** -1. Install "Mermaid Preview" extension -2. Open any `.mmd` file -3. Right-click β†’ "Mermaid: Preview" - ---- - -## 🎯 Which Diagram Should I Use? - -### For Different Audiences - -| Audience | Recommended Diagrams | Purpose | -|----------|---------------------|---------| -| **New Users** | `architecture-overview.mmd` | Get a high-level understanding of the system | -| **Data Scientists** | `finance-pipeline.mmd` | Understand the ML pipeline and data flow | -| **ML Engineers** | `execution-flow.mmd`, `finance-pipeline.mmd` | Learn how to run and debug workflows | -| **Software Engineers** | `component-architecture.mmd` | Understand code structure and extend the framework | -| **DevOps/Infrastructure** | `deployment-architecture.mmd` | Set up cluster and infrastructure | -| **System Architects** | All diagrams + `ARCHITECTURE.md` | Comprehensive system understanding | -| **Contributors** | `component-architecture.mmd`, `ARCHITECTURE.md` | Contribute new stages and recipes | - -### For Different Questions - -| Question | Diagram to Check | -|----------|------------------| -| "What does NVFlow do?" | `architecture-overview.mmd` | -| "How do I build an ML pipeline?" | `finance-pipeline.mmd` | -| "How does stage execution work?" | `execution-flow.mmd` | -| "How do I create a new stage?" | `component-architecture.mmd` | -| "What infrastructure do I need?" | `deployment-architecture.mmd` | -| "How are stages organized?" | `component-architecture.mmd` | -| "How does the finance recipe work?" | `finance-pipeline.mmd` | -| "How does NVFlow integrate with Slurm?" | `deployment-architecture.mmd`, `execution-flow.mmd` | - ---- - -## πŸ“‹ Diagram Details - -### 1. Architecture Overview -``` -Components Shown: -βœ“ User Interfaces (CLI, Python API, Scripts) -βœ“ Core Framework (WorkflowRunner, StageRegistry, BaseStage, Console) -βœ“ Recipe Layer (Finance, Example, Custom recipes) -βœ“ External Dependencies (NeMo-Skills, NeMo-RL, Slurm, Containers) -βœ“ Storage Layer (Data, Models, Outputs) - -Use Case: Understanding system boundaries and component relationships -``` - -### 2. Finance Pipeline -``` -Coverage: -βœ“ Complete 6-workflow pipeline (42 stages total) -βœ“ Data acquisition (SEC filings download) -βœ“ SDG (Template-based & Document-grounded approaches) -βœ“ Data preparation (Transformation, formatting, splitting) -βœ“ Training (Multi-node SFT with Qwen3-14B) -βœ“ Evaluation (Benchmarks and baselines) - -Use Case: Understanding the end-to-end ML pipeline -``` - -### 3. Execution Flow -``` -Phases Covered: -βœ“ Initialization (Config loading, validation) -βœ“ Validation (Registry checks, config validation) -βœ“ Execution (Stage execution loop, job submission) -βœ“ Monitoring (Job status, log viewing) - -Use Case: Debugging and understanding runtime behavior -``` - -### 4. Component Architecture -``` -Classes Documented: -βœ“ BaseStage (abstract base class) -βœ“ StageRegistry (hierarchical registry) -βœ“ WorkflowRunner (orchestrator) -βœ“ Concrete stages (SFT, Generate, Download, Evaluate) -βœ“ CLI (user interface) -βœ“ External dependencies (NeMo-Skills, OmegaConf) - -Use Case: Code navigation and extension -``` - -### 5. Deployment Architecture -``` -Infrastructure Components: -βœ“ Local development machine (NVFlow installation) -βœ“ SSH tunnel (secure connection) -βœ“ Slurm cluster (login node, scheduler, compute nodes) -βœ“ GPU compute nodes (H100 GPUs, containers) -βœ“ Shared storage (Lustre/NFS filesystem) -βœ“ Container runtime (Enroot, .sqsh images) - -Use Case: Cluster setup and deployment planning -``` - ---- - -## 🎨 Diagram Rendering Examples - -### In Markdown (GitHub) -````markdown -```mermaid -graph TB - A[NVFlow] --> B[Recipes] - A --> C[Workflows] - A --> D[Stages] -``` -```` - -### In Python Documentation -```python -""" -Architecture: - Recipe β†’ Workflow β†’ Stage - - See: diagrams/architecture-overview.mmd -""" -``` - -### In Presentations -- Export diagrams to PNG/SVG using `mmdc` CLI -- Import into PowerPoint/Keynote/Google Slides -- High resolution for professional presentations - ---- - -## πŸ“Š Diagram Statistics - -| Metric | Count | -|--------|-------| -| Total Diagrams | 5 | -| Total Documentation Pages | 2 (ARCHITECTURE.md + diagrams/README.md) | -| Components Visualized | 50+ | -| Workflows Documented | 6 | -| Stages Documented | 27 | -| Architecture Layers | 5 | - ---- - -## πŸ”„ Keeping Diagrams Updated - -When updating the codebase: - -1. **Adding a new recipe:** - - Update `architecture-overview.mmd` (Recipe Layer section) - - Consider creating a new pipeline diagram (like `finance-pipeline.mmd`) - -2. **Adding a new stage:** - - Update recipe-specific pipeline diagram - - Update `component-architecture.mmd` if it's a new pattern - -3. **Changing core framework:** - - Update `component-architecture.mmd` - - Update `execution-flow.mmd` if execution logic changes - - Update `ARCHITECTURE.md` with detailed explanations - -4. **Infrastructure changes:** - - Update `deployment-architecture.mmd` - - Update cluster setup documentation - -5. **Major architectural changes:** - - Review and update all diagrams - - Update `ARCHITECTURE.md` comprehensively - ---- - -## πŸ“– Related Documentation - -- **[README.md](../../README.md)** - Main project documentation -- **[INSTALL.md](../../INSTALL.md)** - Installation and setup guide -- **[docs/recipes/finance/README.md](../recipes/finance/README.md)** - Finance recipe documentation -- **[docs/recipes/finance/quick-start.md](../recipes/finance/quick-start.md)** - Quick start guide - ---- - -## 🀝 Contributing - -To contribute to the architecture documentation: - -1. **For diagram updates:** - - Edit the `.mmd` files in `diagrams/` - - Test rendering before committing - - Follow the style guide in `diagrams/README.md` - -2. **For documentation updates:** - - Edit `ARCHITECTURE.md` for comprehensive changes - - Keep diagrams and text synchronized - - Use consistent terminology - -3. **For new diagrams:** - - Create new `.mmd` file in `diagrams/` - - Add description to `diagrams/README.md` - - Update this summary file - ---- - -## πŸ“„ License - -All architecture diagrams and documentation are part of the NVFlow project and follow the Apache-2.0 license. - ---- - -## πŸ™ Acknowledgments - -Built on the NVIDIA NeMo ecosystem: -- [NeMo-Skills](https://github.com/NVIDIA/NeMo-Skills) -- [NeMo-RL](https://github.com/NVIDIA-NeMo/RL) -- [NeMo Framework](https://github.com/NVIDIA/NeMo) - ---- - -**Version:** 1.0 -**Last Updated:** January 21, 2026 -**Maintained by:** NVFlow Team diff --git a/docs/cluster-configuration.md b/docs/cluster-configuration.md index 5ca331b..19636b6 100644 --- a/docs/cluster-configuration.md +++ b/docs/cluster-configuration.md @@ -216,8 +216,10 @@ containers: # Required nemo-skills: /path/to/containers/nemo-skills.sqsh vllm: /path/to/containers/vllm.sqsh + vllm-grpo: /path/to/containers/vllm-grpo.sqsh # GRPO rollouts / judge sglang: /path/to/containers/sglang.sqsh - nemo-rl: /path/to/containers/nemo-rl.sqsh + nemo-rl: /path/to/containers/nemo-rl.sqsh # SFT/GRPO training + nemo-gym: /path/to/containers/nemo-gym.sqsh # CPU Gym-only GRPO stages ``` **Details:** @@ -237,33 +239,39 @@ Maps host file system paths to container paths. ```yaml mounts: - :/hf_models # HuggingFace models - - :/workspace # Your workspace + - :/workspace # Writable data dir (outputs + cache) # Add more mounts as needed: # - /lustre/data:/data ``` **Format:** `:` +> **`/workspace` holds writable data, not source code.** Recipe code and +> checked-in assets (prompts, dataset descriptors, Gym overlays) ship to workers +> via the nemo-run packaged snapshot at `/nemo_run/code` (also the job's working +> directory), so the nvflow repo is **not** mounted. Point `/workspace` at a +> dedicated writable data directory holding `/workspace/outputs/**` (stage +> outputs, checkpoints, SEC cache, eval-datasets) and `/workspace/cache/**` +> (`HF_HOME`) β€” not your repo checkout. On an on-cluster launcher (no +> `ssh_tunnel`), keep the launcher's cwd at the repo root so resume/skip +> detection can map `/workspace/outputs/...` back to the host outputs dir. + **Common mounts:** | Host Path | Container Path | Purpose | |-----------|----------------|---------| -| Your workspace directory | `/workspace` | Code, configs, outputs | +| Writable data directory | `/workspace` | Outputs, checkpoints, caches (code ships via `/nemo_run/code`) | | Shared model storage | `/hf_models` | Pre-trained models | | 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 / NeMo-Gym: trainer image and Gym source -- 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 +SFT and GRPO `training` run on the `nvflow-nemo-rl` image, built from [`dockerfiles/Dockerfile.nemo-rl`](../dockerfiles/Dockerfile.nemo-rl). It bakes the Gym source and one venv per Gym component, so nothing is resolved at job runtime and **no Gym mount is required**. -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 Gym-only GRPO stages (`prepare_data`, `prefetch_cache`, `collect_rollouts`, `compute_rewards`) run on the CPU-only `nvflow-nemo-gym` image, also with baked venvs (`&gym_install_cpu` in `base.yaml`). -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. +Do not bind-mount Gym or NeMo-RL source over the image in production β€” it shadows the baked tree and invalidates the container fingerprint, forcing a runtime rebuild. To iterate on Gym source in dev mode, mount your clone at `/opt/nemo-rl/3rdparty/Gym-workspace/Gym` and leave `UV_OFFLINE` unset so the editable install can resolve. See [`docs/development/nemo-rl-gym.md`](development/nemo-rl-gym.md). ### Model-Specific Cluster Configs @@ -272,14 +280,14 @@ Some models require additional cluster-level differences (e.g. different timeout | 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 | +| `my_cluster_nemotron.yaml` | Nemotron-3-Nano (MoE) | Use only if Nemotron needs different mounts/env -- the `nemo-rl` image handles MoE without a NeMo-RL source overlay | **How it works:** - `base.yaml` (SFT workflow) sets `cluster: my_cluster` as the default - A model config can override with `cluster: my_cluster_nemotron` - Keep both configs in sync when making infrastructure changes -> **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. +> **Note:** No NeMo-RL or Gym source overlay is mounted by default. Nemotron-3-Nano MoE support needs no host overlay, and both `nemo-rl` and `nemo-gym` ship with Gym baked in. See the [SFT Workflow Guide](recipes/finance/workflows/04-sft.md) for the current setup. --- @@ -345,12 +353,10 @@ env_vars: - HF_HUB_OFFLINE=1 - HF_DATASETS_OFFLINE=1 - TRANSFORMERS_OFFLINE=1 - - UV_OFFLINE=true + # - UV_OFFLINE=true # keep unset to allow runtime uv builds; set only for strict airgap - 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_... - OPENAI_API_KEY=sk-... @@ -377,18 +383,20 @@ These variables prevent the runtime from making outbound network calls and from | `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 | +| `UV_OFFLINE` | *unset* | Global flag; **left unset** so components beyond the baked set can be built on demand (see [trainer image and Gym source](#nemo-rl--nemo-gym-trainer-image-and-gym-source)). All GRPO/SFT venvs are baked, so nothing is built at runtime in practice. eval / SDG / SFT never invoke `uv` | | `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). +> **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`. See [INSTALL.md β†’ One-Time Connected-Node Stages](../INSTALL.md#one-time-connected-node-stages-datasets). -#### NeMo-RL / GRPO Variables (Dev Mode Only) +#### NeMo-RL / GRPO training venv -| Variable | Value | Purpose | -|----------|-------|---------| -| `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 | +GRPO `training` runs on the `nemo-rl` image, which bakes Gym and one venv per Gym component. Nothing is built at runtime: NeMo-RL matches `/opt/nemo_rl_container_fingerprint` and reuses the baked venvs. Do not bind-mount Gym or NeMo-RL source over the image -- that shadows the baked tree, invalidates the fingerprint, and forces a rebuild. The Gym-only stages run on the self-contained `nvflow-nemo-gym` image, also with baked venvs. + +`UV_OFFLINE` is left unset so components outside the baked set can still be built on demand. Note the consequence: a fingerprint miss will silently rebuild over the cluster proxy rather than fail, so verify airgap behaviour by checking training logs for venv-build activity, not by the job succeeding. eval / SDG / SFT never invoke `uv`. + +See [trainer image and Gym source](#nemo-rl--nemo-gym-trainer-image-and-gym-source) and [`docs/development/nemo-rl-gym.md`](development/nemo-rl-gym.md). #### API Keys (Secrets) diff --git a/docs/development/nemo-rl-gym.md b/docs/development/nemo-rl-gym.md new file mode 100644 index 0000000..811a202 --- /dev/null +++ b/docs/development/nemo-rl-gym.md @@ -0,0 +1,45 @@ +# NeMo-RL / NeMo-Gym: trainer image & Gym venvs (advanced) + +> Audience: **advanced / dev**. For a normal GRPO run you do **not** need this page β€” follow INSTALL.md and the quick-start. This page explains how GRPO `training` gets NeMo-RL and NeMo-Gym, and how to iterate on Gym source. (SFT `training` runs on the same image but never touches Gym.) + +## How the trainer gets NeMo-RL and Gym + +GRPO `training` runs on `nvflow-nemo-rl`, built from [`dockerfiles/Dockerfile.nemo-rl`](../../dockerfiles/Dockerfile.nemo-rl). The NeMo-RL base supplies Transformer Engine and the prebuilt NeMo-RL / Ray venvs but leaves the Gym venvs unbuilt, because upstream gates that prefetch behind `NEMO_GYM_PREFETCH_CONFIGS`. Our image closes exactly that gap and changes nothing else: + +- The Gym submodule is advanced in place to `GYM_REF` and reinstalled editable into the Gym actor venv. +- One venv is baked **per Gym component** under `/opt/gym_venvs`, by driving `gym env start … +dry_run=true` from that actor venv. Driving it this way is what makes Gym pin each component to the container's own interpreter and Ray version. +- `training.py` sets `env.nemo_gym.skip_venv_if_present = True` (`nvflow/recipes/finance/stages/rl/training.py:157`), so NeMo-RL reuses the baked venvs rather than building. +- The nemo-skills `installation_command` for the trainer is a no-op (`"true"`) β€” no Gym CLI setup is needed inside the trainer container. + +The result is that **nothing resolves at job runtime and no Gym mount is required.** + +The Gym-only stages (`prepare_data`, `prefetch_cache`, `collect_rollouts`, `compute_rewards`) run on the CPU-only `nvflow-nemo-gym` image instead, which bakes the Gym CLI (`/opt/gym-cli-venv`) and its own per-component venvs (`/opt/gym-venvs`). They share the `&gym_install_cpu` command in `nvflow/recipes/finance/workflows/grpo/base.yaml`, which only puts the baked CLI on `PATH` β€” no build, no network. See [`docs/maintainers/containers.md`](../maintainers/containers.md) for both builds. + +### Why two venv directories + +`nvflow-nemo-gym` bakes to `/opt/gym-venvs` (hyphen; `gym_uv_venv_dir` in the workflow YAML); the trainer bakes to `/opt/gym_venvs` (underscore; `NEMO_GYM_VENV_DIR`, inherited from the base). + +Aligning the paths would not make the venvs interchangeable. The images differ in Python (3.12 vs 3.13) and Ray (2.56.1 vs 2.55.1), and both are hard constraints: a Gym server joining the trainer's Ray cluster is version-checked on Ray and on Python down to the patch level, and a venv is bound to its interpreter. Each image bakes where its own runtime looks. + +## Dev iteration on Gym source + +To work against a modified Gym, bind-mount your clone over the trainer's Gym path: + +```yaml +mounts: + - /Gym:/opt/nemo-rl/3rdparty/Gym-workspace/Gym +``` + +This is the one configuration where `uv` resolves at runtime, so it needs `UV_OFFLINE` unset and a reachable pypi mirror. `skip_venv_if_present=True` still applies, so remove the stale venv if you want a rebuild. + +**Do not use this mount in production.** It shadows the baked source and venvs and invalidates the container fingerprint, which turns a fully offline run into one that silently rebuilds over the cluster proxy. + +## Why `UV_OFFLINE` stays unset + +The images need no resolve, so setting it would change nothing in a normal run. It is left unset deliberately, to keep the dev-iteration path above working. + +The trade-off is worth stating: because it is unset, a fingerprint miss **rebuilds instead of failing loudly**. So verify airgap behaviour by checking training logs for venv-build activity, not by the job succeeding. + +## Regression guard + +The `&gym_install_cpu` command and the Gym env-start wiring are covered by `tests/test_grpo_gym_install.py` β€” run `uv run pytest tests/test_grpo_gym_install.py -v` before changing either. diff --git a/docs/development/sdg/document_grounded/ADDING_A_DOMAIN.md b/docs/development/sdg/document_grounded/ADDING_A_DOMAIN.md new file mode 100644 index 0000000..c52aae8 --- /dev/null +++ b/docs/development/sdg/document_grounded/ADDING_A_DOMAIN.md @@ -0,0 +1,714 @@ +# Adding a New Domain to DG-SDG + +> Turn a directory of your own documents into a fine-tuning dataset +> (single `final_result.jsonl` consumed by both SFT and GRPO) by adding a +> new "recipe" to nvflow's Document-Grounded SDG pipeline. +> +> **Audience**: anyone β€” a coworker or an AI agent β€” who can read this doc +> (plus the code it links to), gather the domain-specific info, and build a +> new recipe end-to-end. It should be self-contained enough that handing it +> over is all it takes. The pipeline runs on a Slurm cluster. +> +> **Worked example**: a `legal` recipe with court opinions at +> `/data/legal/cases///.html`. Substitute your own +> domain name wherever you see `legal` / ``. + +The section numbers below mirror the phases in the diagram: + +![DG-SDG: adding a new domain](dgsdg-add-new-domain.png) + +## What you write, in 1 picture + +``` +nvflow/recipes// +β”œβ”€β”€ prompts/ +β”‚ β”œβ”€β”€ document_grounded_generate_questions.yaml Β§2 +β”‚ β”œβ”€β”€ document_grounded_verify_questions.yaml Β§2 +β”‚ β”œβ”€β”€ _qa_template.yaml Β§2 (reused in Β§3) +β”‚ β”œβ”€β”€ evaluate_answers.yaml Β§3 +β”‚ └── genselect_answers.yaml Β§2 (cp from finance verbatim) +β”œβ”€β”€ utils/sdg/ +β”‚ β”œβ”€β”€ _data_preprocess.py Β§1 +β”‚ β”œβ”€β”€ _callbacks.py Β§1 + Β§4 +β”‚ └── _question_prep.py Β§1 +β”‚ └── _postprocess.py Β§5 (thin wrapper) +β”œβ”€β”€ stages/sdg/ Β§5 (register shared generic DG-SDG stages) +β”‚ └── __init__.py Β§5 +β”œβ”€β”€ workflows/sdg/ +β”‚ β”œβ”€β”€ document-grounded-sdg.yaml Β§5 +β”‚ └── document-grounded-sdg-demo.yaml Β§5 +β”œβ”€β”€ __init__.py Β§5 +β”œβ”€β”€ stages/__init__.py Β§5 +β”œβ”€β”€ stages/sdg/__init__.py Β§5 +β”œβ”€β”€ utils/__init__.py Β§5 +β”œβ”€β”€ utils/sdg/__init__.py Β§5 +└── recipe.yaml Β§5 +``` + +--- + +## Β§0 Prerequisites + directory skeleton + +Before you start, verify: + +- nvflow repo checked out on the launcher/host; run commands from the repo root (`ls nvflow/recipes/finance` works). At runtime this code ships to workers via the nemo-run packaged snapshot (`/nemo_run/code`) β€” it is not mounted. +- Cluster config exists (`ls cluster_configs/my_cluster.yaml`) +- `nemo-gym` container available (`enroot list | grep nemo-gym`) β€” the gym-only client the shared DG-SDG generation stages run in (Gym source at `/opt/Gym`, per-component venvs baked at `/opt/gym-venvs`; no `nemo-rl` image or runtime `uv sync` needed for SDG) +- Model weights mounted (`ls /hf_models/openai/gpt-oss-120b` and `ls /hf_models/Qwen/Qwen3-235B-A22B-Instruct-2507`) +- `nflow --help` works +- Your raw documents are in one root directory +- A short snake_case domain name picked (this guide uses `legal`) + +Then create the skeleton: + +```bash +cd # repo root +export DOMAIN=legal # CHANGE ME + +mkdir -p nvflow/recipes/$DOMAIN/{prompts,utils/sdg,stages/sdg,workflows/sdg} +touch nvflow/recipes/$DOMAIN/__init__.py \ + nvflow/recipes/$DOMAIN/stages/__init__.py \ + nvflow/recipes/$DOMAIN/stages/sdg/__init__.py \ + nvflow/recipes/$DOMAIN/utils/__init__.py \ + nvflow/recipes/$DOMAIN/utils/sdg/__init__.py +``` + +--- + +## Β§1 Phase 1 Β· PREPARE INPUT + +Three Python files under `utils/sdg/`. They cover the diagram's Phase 1: +turn raw documents into JSONL records, then attach a `context` string to +each record so the LLM has something to read. + +### 1.1 `_data_preprocess.py` + +Walks your raw document tree and writes one JSONL file with one record per +chunk. **You** run this once manually (and the workflow re-runs it as +step 0). The fields you emit here become the input contract for +`context_builder` in 1.2. + +```python +#!/usr/bin/env python3 +"""Preprocess documents into per-chunk JSONL records.""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path +from typing import Any, Iterable + +from bs4 import BeautifulSoup + +try: + import tiktoken + _ENC = tiktoken.get_encoding("cl100k_base") +except ImportError: + _ENC = None + + +def _tokenize(text: str) -> list[str]: + return _ENC.encode(text) if _ENC is not None else text.split() + + +def _detokenize(tokens) -> str: + return _ENC.decode(tokens) if _ENC is not None else " ".join(tokens) + + +def chunk_text(text: str, max_tokens: int = 2000, overlap_tokens: int = 100) -> Iterable[str]: + toks = _tokenize(text) + step = max(max_tokens - overlap_tokens, 1) + for i in range(0, len(toks), step): + chunk = toks[i : i + max_tokens] + if chunk: + yield _detokenize(chunk) + if i + max_tokens >= len(toks): + break + + +def extract_metadata(html_path: Path) -> dict[str, Any]: + # CUSTOMIZE for your file layout. The keys returned here must be a + # superset of what context_builder reads in 1.2. + parts = html_path.parts + try: + court, year = parts[-3], parts[-2] + except IndexError: + court, year = "", "" + return { + "case_name": re.sub(r"[_\-]+", " ", html_path.stem).strip(), + "court": court, + "decision_year": year, + "section": "Opinion", + "doc_path": str(html_path), + } + + +def extract_body_text(html_path: Path) -> str: + soup = BeautifulSoup(html_path.read_text(encoding="utf-8", errors="ignore"), "html.parser") + for tag in soup(["script", "style", "nav", "footer", "header"]): + tag.decompose() + return soup.get_text(separator="\n", strip=True) + + +def main(input_dir: Path, output_dir: Path, max_tokens: int, overlap_tokens: int) -> int: + output_dir.mkdir(parents=True, exist_ok=True) + out_path = output_dir / f"{input_dir.name}-data.jsonl" + n_records = 0 + with open(out_path, "w", encoding="utf-8") as out: + for html_path in sorted(input_dir.rglob("*.html")): + meta = extract_metadata(html_path) + body = extract_body_text(html_path) + if not body.strip(): + continue + for chunk_idx, chunk_str in enumerate(chunk_text(body, max_tokens, overlap_tokens)): + rec = {**meta, "chunk_id": chunk_idx, "content": chunk_str} + out.write(json.dumps(rec, ensure_ascii=False) + "\n") + n_records += 1 + print(f"Wrote {n_records} records to {out_path}") + return n_records + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--input_dir", type=Path, required=True) + p.add_argument("--output_dir", type=Path, required=True) + p.add_argument("--max_tokens", type=int, default=2000) + p.add_argument("--overlap_tokens", type=int, default=100) + # The Stage 0 shim (generic_stage/sdg/document_grounded/dg_sdg_preprocess.py) ALWAYS + # passes these four extra flags too. You must accept them even if your + # domain doesn't sample by a distribution -- otherwise argparse aborts + # the Slurm job with "unrecognized arguments". Ignore the ones you don't + # use (finance reads multiple CSVs from --distribution_dir; see note below). + p.add_argument("--distribution_dir", type=Path, default=None) + p.add_argument("--total_samples", type=int, default=150000) + p.add_argument("--max_skip_count", type=int, default=20000) + p.add_argument("--seed", type=int, default=42) + args = p.parse_args() + main(args.input_dir, args.output_dir, args.max_tokens, args.overlap_tokens) +``` + +> **Stage 0 CLI contract β€” accept all 8 flags.** The generic shim invokes +> your module as +> `python3 -m _data_preprocess --input_dir … --output_dir … +> --distribution_dir … --max_tokens … --overlap_tokens … --total_samples … +> --max_skip_count … --seed …`. Your argparse must define every one of these +> (the four above plus the four sampling flags) or the job crashes before it +> does any work. `distribution_dir` is currently **required by the shim's +> `validate_config`**, so the workflow YAML must set +> `stages.dg_sdg_preprocess.distribution_dir` even if your CLI ignores it +> (point it at a dir with a placeholder CSV). +> +> **Multi-CSV input is supported.** `--distribution_dir` is a *directory*, not +> a single file, so a domain can read any number of CSVs from it. Finance +> reads four (`{10k,10q}_{1company,2company}_distribution.csv`); there is no +> generic constraint on count or naming β€” your CLI decides what to load. +> +> For non-HTML inputs replace `extract_body_text` with whatever extracts +> text from your format (`.read_text()`, `pypdf`, `pdfminer.six`, etc.). +> Chunking + metadata logic stays the same. +> +> **Finance counterpart** (for reference): `nvflow/recipes/finance/utils/sdg/dg_sdg_data_preprocess.py`. +> Keep the `_data_` infix to avoid confusion with the stage name +> `dg_sdg_preprocess`. + +### 1.2 `_callbacks.py` (context_builder) + +A pure function: `(record) β†’ str`. Called by the library once per record +during step-1 question generation. Empty string β‡’ skip the record. + +The fields you reference here MUST match what 1.1 writes. + +```python +"""Domain-specific callbacks for the DG-SDG recipe.""" + +from typing import Any + + +def legal_context_builder(record: dict[str, Any]) -> str: + case_name = record.get("case_name", "") + court = record.get("court", "") + year = record.get("decision_year", "") + section = record.get("section", "Opinion") + content = record.get("content", "") + + if not content: + return "" + + return ( + f"**{year} {court}: {case_name}**\n\n" + f"**Section: {section}**\n\n" + f"{content}\n" + ) + +# Β§4 will add is_legal_sft_eligible / is_legal_rl_eligible to this same file. +``` + +> **Finance counterpart**: `nvflow/recipes/finance/utils/sdg/sec_callbacks.py` +> (finance uses `sec_*` prefix not `finance_*`; the file is named after the +> SECQUE benchmark for historical reasons). + +### 1.3 `_question_prep.py` + +A 20-line CLI that bolts `context_builder` into the lib's generic helper. +This is the only place `context_builder` is actually invoked, and it's +what the workflow's step-1 entrypoint calls. + +```python +"""Thin CLI: construct_question_generate_input with the legal context_builder.""" + +import argparse +from pathlib import Path + +from nvflow.lib.sdg.document_grounded.preprocess import construct_question_generate_input +from nvflow.recipes.legal.utils.sdg.legal_callbacks import legal_context_builder + + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--input_folder", type=Path, required=True) + p.add_argument("--output_file", type=Path, required=True) + args = p.parse_args() + + construct_question_generate_input( + args.input_folder, + args.output_file, + context_builder=legal_context_builder, + ) +``` + +> **Finance counterpart**: `nvflow/recipes/finance/utils/sdg/sec_question_prep.py`. + +--- + +## Β§2 Phase 2 Β· GENERATE Q&A + +Four prompt YAMLs under `prompts/`. Two are domain-specific (you write +them), two come straight from finance (copy verbatim). + +> **JSON braces in YAML prompts**: literal `{` / `}` must be doubled +> (`{{` / `}}`) because Python `.format()` substitutes `{context}` etc. + +### 2.1 `document_grounded_generate_questions.yaml` (Q-gen) + +Generates ~12 questions per chunk in valid JSON the lib can parse. + +```yaml +# nvflow/recipes/legal/prompts/document_grounded_generate_questions.yaml +user: |- + You are a senior legal analyst. You will be given an excerpt from a court + opinion or other legal document. + + Your task is to propose the most important questions a legal researcher + should ask about the excerpt. Generate exactly 3 questions for EACH of + the following categories: + - Holding_and_Reasoning + - Procedural_History + - Legal_Standard_Applied + - Implications_and_Precedent + + Respond ONLY with a valid JSON object. No markdown, no commentary. + + Output Format: + {{ + "Holding_and_Reasoning": ["q1", "q2", "q3"], + "Procedural_History": ["q1", "q2", "q3"], + "Legal_Standard_Applied": ["q1", "q2", "q3"], + "Implications_and_Precedent": ["q1", "q2", "q3"] + }} + + Document: {context} +``` + +> If you change the schema (different categories / counts), you also need a +> custom `generation_parser` for the next stage β€” easier to keep this shape. + +### 2.2 `document_grounded_verify_questions.yaml` (Q-verify) + +Per-question Yes/No verdict. **The `system:` block is mandatory** β€” +without "respond ONLY Yes/No" the verifier regex silently drops ~30%+ of +valid questions. + +```yaml +# nvflow/recipes/legal/prompts/document_grounded_verify_questions.yaml +system: |- + You are a legal expert validating analytical questions. Decide if a given + question is valid, expert-level, and answerable using only the Reference + Text. + + Criteria for "Yes": + 1. The Reference Text contains the facts needed to answer. + 2. The question is non-trivial and assesses legal reasoning. + + Criteria for "No": + 1. The Reference Text lacks the specific data or context. + 2. The question is malformed or unrelated. + + Respond ONLY with "Yes" or "No". + +user: |- + **Reference Text:** + {context} + + **Question:** + {problem} + + Is this a valid expert-level question answerable from the text? +``` + +### 2.3 `_qa_template.yaml` (A-gen prompt) + +The single-turn answer prompt used by the answer-generation stage +(`generate_answers`, step-2). + +```yaml +# nvflow/recipes/legal/prompts/legal_qa_template.yaml +user: |- + You are a legal expert. Given a court-opinion excerpt and a question + written by a senior analyst, answer using ONLY the provided text. Do not + use external knowledge. Be concise but precise. If the text does not + support an answer, say so explicitly. + + Document: {context} + + Question: {problem} + + Answer: +``` + +> **Finance counterpart**: `nvflow/recipes/finance/prompts/secque_template.yaml` +> (the prod workflow YAML references it once, as the `generate_answers` stage's +> `++prompt_config`). + +### 2.4 `genselect_answers.yaml` (best-of-N picker) + +Generic, copy verbatim from finance: + +```bash +cp nvflow/recipes/finance/prompts/genselect_answers.yaml \ + nvflow/recipes/$DOMAIN/prompts/genselect_answers.yaml +``` + +--- + +## Β§3 Phase 3 Β· REFINE + +One prompt YAML you write: the judge for the `evaluate_answers` stage. +(`aggregate_answers` then folds the per-seed verdicts into a consensus +`answerable` and needs no prompt.) + +### 3.1 `evaluate_answers.yaml` (judge for seed-evaluation stage) + +Must emit a one-line JSON tag `{"answerable": "YES/NO", "correct": "YES/NO"}` +on the **last** line β€” `evaluate.parse_evaluation` regex looks for exactly +that shape. + +```yaml +# nvflow/recipes/legal/prompts/evaluate_answers.yaml +user: |- + You are evaluating an AI assistant's answer to a legal question grounded + in the provided court-opinion excerpt. + + You need to decide TWO things: + 1. ANSWERABLE: can the question be answered using only the excerpt? + 2. CORRECT: is the assistant's response appropriate? + + ANSWERABLE assessment: + - YES: excerpt contains the necessary facts / citations / reasoning. + - NO: excerpt lacks the necessary information. + + CORRECT assessment: + - When ANSWERABLE=YES: assistant gives an accurate, well-supported answer. + - When ANSWERABLE=NO: assistant correctly identifies info is missing. + + Provide your reasoning first, then end with this exact JSON tag on a new line: + {{"answerable": "YES/NO", "correct": "YES/NO"}} + + Document: {context} + + Question: {problem} + + Assistant's Answer: {generation} +``` + +--- + +## Β§4 Phase 4 Β· SHIP + +No subset-eligibility callbacks are needed. The pipeline emits a single +`final_result.jsonl` per run; downstream SFT / GRPO workflows pick +records by reading that file directly. If you later need a curated SFT +or GRPO subset, do it as a separate post-process step outside DG-SDG +(e.g. a small CLI in your recipe's `utils/`). + +The only domain-specific callback used by the shared stages is the +context-builder (`_context_builder`) wired into +`generate_verified_questions` via the `question_prep_script`, which you +already added in Β§1.2. + +--- + +## Β§5 Workflow wiring + launch + +Last lap: register shared DG-SDG stages, add one postprocess wrapper, +drop in 5 registration files, write the 2 workflow YAMLs, then launch. + +### 5.1 Register shared DG-SDG stages + add postprocess wrapper + +Add a thin domain wrapper around `nvflow.lib.sdg.document_grounded.postprocess`. +There is nothing domain-specific to inject by default β€” it exists only +so the workflow YAML can point at a recipe-owned path, leaving room for +domain-specific cleaning later: + +```python +# nvflow/recipes/legal/utils/sdg/legal_postprocess.py +import argparse +import os +import sys + +from nvflow.lib.sdg.document_grounded.postprocess import dgsdg_post_process + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Post-process DG-SDG data for the legal recipe." + ) + parser.add_argument("--input_file", required=True) + parser.add_argument("--output_dir", required=True) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + if not os.path.exists(args.input_file): + sys.exit(f"Input file not found: {args.input_file}") + + dgsdg_post_process(args.input_file, args.output_dir, seed=args.seed) +``` + +### 5.2 `__init__.py` Γ— 3 + `recipe.yaml` + +```python +# nvflow/recipes/legal/__init__.py +from . import stages # noqa: F401 +``` + +```python +# nvflow/recipes/legal/stages/__init__.py +from . import sdg # noqa: F401 +``` + +```python +# nvflow/recipes/legal/stages/sdg/__init__.py +from nvflow.generic_stage.sdg.document_grounded import register_for_recipe + +register_for_recipe("legal") +``` + +```yaml +# nvflow/recipes/legal/recipe.yaml +recipe: legal +description: "End-to-end pipeline for legal-domain model training and evaluation" + +workflow_order: + - document_grounded_sdg +``` + +### 5.3 Production workflow YAML + +Start from finance and edit paths: + +```bash +cp nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml \ + nvflow/recipes/$DOMAIN/workflows/sdg/document-grounded-sdg.yaml +``` + +Then edit (search for the strings on the left): + +| Find | Replace with | +| ------------------------------------------------------------- | ------------------------------------------------------------------------------------- | +| `recipe: finance` | `recipe: legal` | +| `description: "generate synthetic finance data ..."` | `description: "generate synthetic legal data ..."` | +| `base_data_dir: /workspace/outputs/finance/...` | `base_data_dir: /workspace/outputs/legal/workflow-document-grounded-sdg` | +| `filings_dir: /workspace/outputs/finance/...` | `filings_dir: /data/legal/cases` (your raw doc root) | +| `pipeline_stages: [- dg_sdg_preprocess, ...]` | Keep as-is. Order is `dg_sdg_preprocess β†’ generate_verified_questions β†’ generate_answers β†’ gym_genselect_answers β†’ evaluate_answers β†’ aggregate_answers β†’ dgsdg_post_process` (7 shared stages) | +| `stages: dg_sdg_preprocess:` block name | Keep block name; set `preprocess_module` to your `_data_preprocess` module path | +| `stages: generate_verified_questions:` block | Set `question_prep_script: nvflow/recipes//utils/sdg/_question_prep.py` | +| `stages: generate_answers:` block | Nothing domain-specific; inherits `gym_*` from prod. Tune `answer_preprocess_kwargs.threshold` if needed | +| `stages: gym_genselect_answers:` block | Set `prompt_template: nvflow/recipes//prompts/genselect_answers.yaml` | +| `stages: dgsdg_post_process:` block | Add `postprocess_script: nvflow/recipes//utils/sdg/_postprocess.py` | +| `++prompt_config=…/prompts/document_grounded_*.yaml` (Q-gen + Q-verify) | Repoint both to `nvflow/recipes/legal/prompts/…` | +| `++prompt_config=…/prompts/secque_template.yaml` (A-gen) | `…/prompts/legal_qa_template.yaml` | +| `prompt_template: …/prompts/evaluate_answers.yaml` (evaluate stage) | `nvflow/recipes/legal/prompts/evaluate_answers.yaml` | +| `finance_domain_keep_fields:` anchor + 6 stage refs | Rename anchor to `_domain_keep_fields:`, replace member fields with **every** domain key your callbacks / preprocess CLI write to JSONL that you want to survive to final training data. (Stages 1–6 carry the anchor; Stage 0 `dg_sdg_preprocess` does not trim.) See Β§5.6 for the full mechanics. | + +> **Don't touch** `gym_path`, `gym_container`, +> `gym_config_paths_format_verification`, `gym_agent_format_verification`, +> `verifier_passthrough`, `verifier_parse_vote` β€” they reference the gym-only +> container, upstream Gym envs, and SDG overlay YAMLs, all of which are +> domain-agnostic. +> +> **Don't touch** model paths under `args: model: /hf_models/…` unless you +> want different models. Finance defaults (gpt-oss-120b for Q-gen + A-gen, +> Qwen3-235B for Q-verify + judges) are strong general-purpose choices. + +### 5.4 Demo workflow YAML + +Inherit prod via `_base_:` and override just the "make it small + fast" +knobs: + +```yaml +# nvflow/recipes/legal/workflows/sdg/document-grounded-sdg-demo.yaml +recipe: legal +workflow: + name: "document_grounded_sdg" + type: "sdg" + description: "demo (smoke test) of document-grounded SDG for legal" + +cluster: my_cluster +_base_: document-grounded-sdg.yaml + +base_data_dir: /workspace/outputs/legal/demo/workflow-document-grounded-sdg-demo +filings_dir: /data/legal/cases # or a small subdir for smoke + +stages: + dg_sdg_preprocess: + max_tokens: 2000 + overlap_tokens: 200 + total_samples: 200 # prod uses 150_000 + + generate_verified_questions: + question_verify_kwargs: + args: + num_random_seeds: 3 + num_chunks: 4 + + generate_answers: + answer_preprocess_kwargs: + threshold: 0.5 + answer_generation_kwargs: + args: + num_random_seeds: 3 + + gym_genselect_answers: + num_chunks: 1 + num_random_seeds: 1 + + evaluate_answers: + num_chunks: 1 + num_random_seeds: 1 +``` + +> Domain paths (`preprocess_module`, `question_prep_script`, `postprocess_script`, +> `prompt_template`) are inherited from prod via `_base_:` deep-merge β€” no need +> to repeat them in demo. + +### 5.5 Launch + +```bash +uv run nflow run-all \ + --config nvflow/recipes/$DOMAIN/workflows/sdg/document-grounded-sdg-demo.yaml +``` + +That submits all 7 stages with `afterok` Slurm dependencies and returns +immediately. + +Output lands in `base_data_dir`: + +``` +$base_data_dir/ +β”œβ”€β”€ step-0-preprocess/jsonl/*.jsonl +β”œβ”€β”€ step-1-questions/ +β”‚ β”œβ”€β”€ generated/ # raw Q-gen rollouts +β”‚ └── verified/ # Q-verify rollouts (consumed by step-2) +β”œβ”€β”€ step-2-answers/ +β”‚ └── generated/output-rs*.jsonl # N candidate answers per question +β”œβ”€β”€ step-3-genselect/selected_answers.jsonl +β”œβ”€β”€ step-4-evaluate/output-rs*.jsonl +β”œβ”€β”€ step-5-aggregate/aggregated_answers.jsonl +└── step-6-post-process/ + └── final_result.jsonl # single cleaned + renamed dataset for SFT / GRPO +``` + +To launch production (after demo works): swap to +`document-grounded-sdg.yaml` (no `-demo` suffix). + +### 5.6 Stage boundary trim (`domain_keep_fields`) + +Every generic DG-SDG stage (Stages 1 through 6) projects its output JSONL +to an allowlist before the next stage reads it. The trim runs inside the +same Slurm job that produces the output, so there is **no extra dependency +to wire and no extra wall-time cost**. + +The allowlist is the set union: + +``` +STAGE_KEEP[stage] # generic fields the lib code produces / needs +| domain_keep_fields # extra fields your recipe writes that you want to survive +- ALWAYS_DROP # NeMo-Gym noise that we always strip +``` + +`STAGE_KEEP[stage]` and `ALWAYS_DROP` live in +[`nvflow/generic_stage/sdg/document_grounded/_schemas.py`](../../../../nvflow/generic_stage/sdg/document_grounded/_schemas.py) β€” +you should not need to edit either when adding a new domain. + +**What you write**: one YAML anchor in your prod workflow YAML and a +reference from each of the 6 trim-eligible stage blocks (Stage 0 +`dg_sdg_preprocess` does not trim because it manufactures the initial +JSONL from raw documents): + +```yaml +# ---- top of document-grounded-sdg.yaml ---- +legal_domain_keep_fields: &legal_domain_keep_fields + - case_id # every field your callbacks / preprocess CLI + - jurisdiction # write into JSONL that you want to survive + - filing_year # all the way to step-6-post-process + # ... (omit raw text fields like `content*` β€” see gotcha below) + +stages: + generate_verified_questions: + # ... existing keys ... + domain_keep_fields: *legal_domain_keep_fields + generate_answers: + domain_keep_fields: *legal_domain_keep_fields + gym_genselect_answers: + domain_keep_fields: *legal_domain_keep_fields + evaluate_answers: + domain_keep_fields: *legal_domain_keep_fields + aggregate_answers: + domain_keep_fields: *legal_domain_keep_fields + dgsdg_post_process: + domain_keep_fields: *legal_domain_keep_fields +``` + +The demo YAML inherits everything via `_base_:` deep-merge β€” no override +needed. + +**Generic `STAGE_KEEP` cheat-sheet** (for context β€” you don't need to +list these in `domain_keep_fields`): + +| Stage | Keep | +| ----------------------------- | ---- | +| `generate_verified_questions` | `context`, `problem`, `question_type`, `generation` | +| `generate_answers` | + `question_voting_pass_rate`, `question_voting_total`, `reasoning_content`, and the Responses-API original form of each candidate answer (`answer_response`, `answer_responses_create_params`) | +| `gym_genselect_answers` | + `reference_answer`, `reference_reasoning`, the picked answer's Responses-API original form (`reference_response`, `reference_responses_create_params`), `genselect_answers_metadata` (drops multi-candidate scaffolding) | +| `evaluate_answers` | as above + `evaluate_generation` | +| `aggregate_answers` | as above + `answerable` (drops `evaluate_generation`) | +| `dgsdg_post_process` | renames `reference_*` β†’ `answer` / `reasoning_content` / `response` / `responses_create_params`, adds `expected_answer` (mirrors `answer`), drops `generation` + `genselect_answers_metadata`; keeps voting stats | + +**Common gotchas:** + +- **Silent drop**: if your `context_builder` or postprocess wrapper writes + a field that's not in `domain_keep_fields`, it is **silently removed at + the first stage boundary**. The final training data won't have it. Add + the field to the anchor. +- **`content0/1/2` / raw text**: finance intentionally omits these. The + Q-prep callback folds them into `context`, so the raw markdown is + redundant after Stage 0. If your domain produces raw text that you want + to ship to SFT, either fold it into `context` in your callback or list + it in `domain_keep_fields`. +- **Per-record schema variance**: if your domain emits records with + different field shapes (finance has 1-company vs 2-company variants), + list the **union** of all variants in the anchor. The trim allowlist + treats missing fields as a no-op (no error). +- **Stage 0 has no trim**: it writes whatever your `_data_preprocess` + CLI writes. If you write junk fields, Stage 1's trim catches them, but + it's cleaner to write only the fields you intend to propagate. diff --git a/docs/development/sdg/document_grounded/dgsdg-add-new-domain.png b/docs/development/sdg/document_grounded/dgsdg-add-new-domain.png new file mode 100644 index 0000000..1bbf779 Binary files /dev/null and b/docs/development/sdg/document_grounded/dgsdg-add-new-domain.png differ diff --git a/docs/diagrams/README.md b/docs/diagrams/README.md index 0f49021..70c6fbc 100644 --- a/docs/diagrams/README.md +++ b/docs/diagrams/README.md @@ -102,7 +102,7 @@ Many modern IDEs (including Cursor) have built-in Mermaid preview support. Simpl ## πŸ“š Additional Documentation For detailed architectural descriptions and explanations, see: -- **[ARCHITECTURE.md](../architecture/ARCHITECTURE.md)** - Comprehensive architecture documentation with embedded diagrams +- **[ARCHITECTURE.md](../ARCHITECTURE.md)** - Comprehensive architecture documentation with embedded diagrams - **[README.md](../../README.md)** - Main project documentation - **[docs/recipes/finance/README.md](../recipes/finance/README.md)** - Finance recipe documentation @@ -144,13 +144,10 @@ When adding new diagrams: 3. Add a header comment explaining the diagram purpose 4. Update this README with a description 5. Test rendering in at least one viewer before committing -6. Consider updating [ARCHITECTURE.md](../architecture/ARCHITECTURE.md) if adding significant architectural information +6. Consider updating [ARCHITECTURE.md](../ARCHITECTURE.md) if adding significant architectural information ## πŸ“„ License These diagrams are part of the NVFlow project and follow the same Apache-2.0 license. ---- - -**Last Updated:** January 21, 2026 **Maintainer:** NVFlow Team diff --git a/docs/diagrams/finance-pipeline.mmd b/docs/diagrams/finance-pipeline.mmd index e35f3d1..4235303 100644 --- a/docs/diagrams/finance-pipeline.mmd +++ b/docs/diagrams/finance-pipeline.mmd @@ -26,14 +26,14 @@ graph TB end subgraph DGS["Document-Grounded SDG (Experimental)"] - DGS1["1. Preprocess Filings
Extract sections"] - DGS2["2. Generate Q&A
With verification"] - DGS3["3. GenSelect
Self-consistency"] - DGS4["4. Evaluate Quality
Judge-based scoring"] - DGS5["5. Aggregate Results
Combine datasets"] - DGS6["6. Difficulty Estimation
Stratify by difficulty"] - DGS7["7. Prepare Training Data
Format conversion"] - DGS_OUT[("~800K Q&A pairs
Stratified by difficulty
Work in progress")] + DGS1["1. Preprocess Filings
Chunk SEC HTML"] + DGS2["2. Generate Verified Questions
Q-gen + Yes/No verify"] + DGS3["3. Generate Answers
N candidates per question"] + DGS4["4. GenSelect Answers
Best-of-N pick"] + DGS5["5. Evaluate Answers
Judge (multi-seed)"] + DGS6["6. Aggregate Answers
Consensus answerable"] + DGS7["7. Post-process
Clean + rename"] + DGS_OUT[("~800K Q&A pairs
Single final_result.jsonl
Work in progress")] DGS1 --> DGS2 --> DGS3 --> DGS4 --> DGS5 --> DGS6 --> DGS7 --> DGS_OUT end diff --git a/docs/maintainers/containers.md b/docs/maintainers/containers.md new file mode 100644 index 0000000..b91c84c --- /dev/null +++ b/docs/maintainers/containers.md @@ -0,0 +1,137 @@ +# Building & Staging the Cluster Containers (maintainers) + +> Audience: **maintainers / builders** who produce the `.sqsh` container images for a cluster. If a maintainer has already staged the `.sqsh` files on your cluster, you don't need this page β€” just set the container paths in your cluster config (see [INSTALL.md β†’ Setup Containers](../../INSTALL.md#setup-containers)) and continue. + +NVFlow uses five core containers converted to `.sqsh` format for running on Slurm clusters, plus a CPU-only `nemo-gym` worker needed only for GRPO / DG-SDG (see [Gym worker](#gym-worker-cpu-only) below). Of the five core, **four are built locally** from self-contained Dockerfiles in [`dockerfiles/`](../../dockerfiles/) (`nemo-rl`, `nemo-skills`, `vllm`, `vllm-grpo`); only `sglang` is **pulled as-is**. + +## Build host requirements + +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, so `docker build` can pull 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`. + +## Required containers (5) + +| Container | Source | Tested Version | Action | +|-----------|--------|----------------|--------| +| `nvflow-nemo-rl` | [`dockerfiles/Dockerfile.nemo-rl`](../../dockerfiles/Dockerfile.nemo-rl) | base `nvcr.io/nvidia/nemo-rl:v0.7.0`, Gym @ `33ef60369` | **Build** (Gym venvs baked) | +| `nvflow-nemo-skills` | [`dockerfiles/Dockerfile.nemo-skills`](../../dockerfiles/Dockerfile.nemo-skills) | NeMo-Skills @ `e06c9b90` (tag `v1.1.2`) | **Build** (see Step 1) | +| `nvflow-vllm` | [`dockerfiles/Dockerfile.vllm`](../../dockerfiles/Dockerfile.vllm) | base `vllm/vllm-openai:v0.22.0` | **Build** (SDG/eval) | +| `nvflow-vllm` (`v0.20.0*` tag) | [`dockerfiles/Dockerfile.vllm`](../../dockerfiles/Dockerfile.vllm) `--build-arg VLLM_VERSION=v0.20.0` | base `vllm/vllm-openai:v0.20.0` | **Build** (GRPO rollouts/judge) | +| `sglang` | Docker Hub | `lmsysorg/sglang:v0.5.10.post1` | **Pull** (no custom Dockerfile) | + +> **Note:** The four custom worker images (`nemo-rl`, `nemo-skills`, `vllm`, `vllm-grpo`) are **built**; only `sglang` is **pulled as-is**. The custom Dockerfiles bake in their source, pre-built venvs, and `tiktoken` / `openai_harmony` caches so they run offline under `enroot`/`pyxis` with no outbound network. + +**Optional containers** (not currently used by any NVFlow recipes): + +| Container | Source | Action | +|-----------|--------|--------| +| `megatron` | NeMo-Skills Dockerfiles | Build | +| `sandbox` | NeMo-Skills Dockerfiles | Build | +| `verl` | NeMo-Skills Dockerfiles | Build | +| `trtllm` | `nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc8` | Pull from NGC | + +### Gym worker (CPU-only) + +The **Gym-only stages** β€” GRPO `prepare_data` / `prefetch_cache` and the DG-SDG gym stages β€” run in a dedicated **CPU-only** worker, **`nvflow-nemo-gym`** ([`dockerfiles/Dockerfile.nemo-gym`](../../dockerfiles/Dockerfile.nemo-gym), base `python:3.12-slim`, upstream Gym main `33ef60369`). It bakes one venv **per Gym component** (`gym env start … +dry_run`; `equivalence_llm_judge` + `finance_sec_search` + `format_verification` prebuilt, others build on demand) into `/opt/gym-venvs`, with Gym source at `/opt/Gym`. Stage it if you run **GRPO or DG-SDG** (SFT-only / eval-only runs don't need it). It is referenced by `my_cluster.yaml` `containers:` as **`nemo-gym`** and listed in [`cluster_configs/containers.yaml`](../../cluster_configs/containers.yaml). Build multi-arch (amd64 + arm64); `GYM_REF` is a pinned SHA, so layer caching is safe. + +### Launcher image (optional, airgap-only) + +Separate from the five **worker** containers above, the **`nvflow-client`** launcher image ([`dockerfiles/Dockerfile.nvflow`](../../dockerfiles/Dockerfile.nvflow), pinned Ubuntu 24.04 base with Python 3.12) bundles the `nflow` CLI + baked venv so **users in an airgapped environment who cannot `uv sync`** can drive NVFlow over an `ssh_tunnel`. It is a **launcher, not a worker**: it is *not* referenced by `my_cluster.yaml` `containers:` and is *not* required for a normal (`uv sync`) install. Build it **multi-arch (amd64 + arm64)** and match the client/cluster architecture. See [`docs/remote-launch.md`](../remote-launch.md) for usage. It is listed in [`cluster_configs/containers.yaml`](../../cluster_configs/containers.yaml) as `nvflow-client` (release-tag placeholder). + +## Step 1: Build Docker Images + +NVFlow ships self-contained Dockerfiles in [`dockerfiles/`](../../dockerfiles/) that pre-install all Python packages, pre-cache tokenizer encodings, and pre-build virtual environments. The full build commands β€” single-arch, cross-arch / multi-arch (`docker buildx` + QEMU), and the `sglang` pull β€” are in **[`dockerfiles/docker_instructions.md` Β§1](../../dockerfiles/docker_instructions.md#1-build)** (the authoritative build reference); per-image `ARG` version pins are in [`dockerfiles/README.md`](../../dockerfiles/README.md#version-pins). + +> **Tip:** Keep `NEMO_SKILLS_COMMIT` consistent between `Dockerfile.nemo-skills` and `pyproject.toml`. For optional containers (`megatron`, `sandbox`, `verl`), build them from the upstream [NeMo-Skills Dockerfiles](https://github.com/NVIDIA-NeMo/Skills/tree/e06c9b90/dockerfiles). + +### Step 1b: Sanity-Check Images Before Conversion + +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 trainer's 7 baked Gym component venvs are present, `tiktoken` / `openai_harmony` caches load with `--network=none`, and `tzdata` is populated. + +## Step 2: Get Images onto the Cluster + +Slurm nodes usually have no Docker, so `enroot` pulls each image from a **registry** (`docker://`, recommended) or loads it from a **saved tarball** (`dockerd://`, for sites with no registry). Tag/push and `docker save` commands for both paths are in [`dockerfiles/docker_instructions.md` Β§3](../../dockerfiles/docker_instructions.md#3-convert-to-sqsh-for-the-slurm-cluster). `sglang` can be pulled directly by `enroot` β€” no push needed unless your cluster cannot reach Docker Hub. + +## Step 3: Update Container Config + +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 `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-rl: your-registry/nvflow-nemo-rl:v0.7.0 # built locally; Gym venvs baked + nemo-skills: your-registry/nvflow-nemo-skills:v1.1.2 + vllm: your-registry/nvflow-vllm:v0.22.0 # v0.22.0 for SDG/eval + vllm-grpo: your-registry/nvflow-vllm:v0.20.0 # same repo as vllm, v0.20.0 tag 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 + +### 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`: + +```bash +# Run from a cluster login node (sbatch requires Slurm access) +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`). + +> Output filenames are derived as `-.sqsh` from the YAML key and tag (not the registry path), and any image whose file already exists is skipped β€” pass `--force` to re-download. + +**Check progress:** +```bash +tail -f outputs/logs/slurm-containers-.out +``` + +### Option B: Manual Conversion + +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 +CONTAINER_DIR= + +# Use -.sqsh so manual imports and setup_containers.sh agree. +enroot import --output $CONTAINER_DIR/nemo-skills-v1.1.2.sqsh \ + "docker://$REGISTRY/nvflow-nemo-skills:v1.1.2" # from a registry +# -- or -- +gunzip -c nvflow-nemo-skills-v1.1.2.tar.gz | docker load +enroot import --output $CONTAINER_DIR/nemo-skills-v1.1.2.sqsh \ + dockerd://nvflow-nemo-skills:v1.1.2 # from a tarball +``` + +Repeat for `vllm`, `vllm-grpo`, `nemo-gym`, and `nemo-rl`. `sglang` imports directly from its upstream registry (`docker://lmsysorg/sglang:v0.5.10.post1`). + +**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 vllm-v0.22.0.sqsh \ + "docker://nvcr.io#/nvflow-vllm:v0.22.0" + ``` +- **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-skills:v1.1.2.sqsh" nemo-skills-v1.1.2.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, then record those paths in your cluster config. diff --git a/docs/recipes/finance/README.md b/docs/recipes/finance/README.md index 15d557e..2350a4a 100644 --- a/docs/recipes/finance/README.md +++ b/docs/recipes/finance/README.md @@ -8,19 +8,19 @@ End-to-end pipeline for generating synthetic financial Q&A data from SEC filings **Two Independent SDG Approaches:** - **Template-Based SDG:** Adapts seed questions to different companies/years, maps to relevant context, generates and filters answers -- **Document-Grounded SDG:** Generates questions directly from documents with built-in verification, quality evaluation, and difficulty stratification +- **Document-Grounded SDG:** Generates questions directly from documents with built-in verification and multi-seed quality evaluation, emitting a single `final_result.jsonl` **Production-Ready Pipeline:** - **Data Generation:** Uses GPT-OSS-120B, Qwen3 (14B-235B) models for synthetic Q&A creation -- **Scale:** Processes S&P 500 companies (~100GB filings) β†’ generates 1M+ Q&A pairs +- **Scale:** Processes S&P 500 companies (~100GB filings) β†’ generates 300K+ Q&A pairs - **Training:** Full SFT pipeline on 256 GPUs (32 nodes) with Qwen3-14B - **Evaluation:** Benchmark trained models on financial reasoning tasks ## What This Recipe Produces -- **Synthetic Q&A Datasets**: 1M+ high-quality financial question-answer pairs +- **Synthetic Q&A Datasets**: 300K+ high-quality financial question-answer pairs - Template-based SDG: ~300K pairs (used in production SFT) - - Document-grounded SDG: ~800K pairs (SFT integration in progress) + - Document-grounded SDG: additional pairs (experimental; SFT integration in progress) - **Fine-tuned Models**: Financial reasoning models trained via supervised fine-tuning (SFT) - **RL-trained Models**: Models further improved via GRPO reinforcement learning with LLM-as-judge rewards - **Evaluation Results**: Model performance on financial benchmarks (SFT and GRPO checkpoints) @@ -78,14 +78,14 @@ End-to-end pipeline for generating synthetic financial Q&A data from SEC filings β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”‚ β€’ Generate questions β”‚ β”‚ β€’ Preprocess filings β”‚ β”‚ β€’ Map to context β”‚ β”‚ β€’ Generate verified Q&A β”‚ -β”‚ β€’ Generate answers β”‚ β”‚ β€’ GenSelect answers β”‚ -β”‚ β€’ GenSelect answers β”‚ β”‚ β€’ Evaluate quality β”‚ -β”‚ β€’ Filter quality β”‚ β”‚ β€’ Aggregate results β”‚ -β”‚ β”‚ β”‚ β€’ Estimate difficulty β”‚ -β”‚ β”‚ β”‚ β€’ Prepare training data β”‚ +β”‚ β€’ Generate answers β”‚ β”‚ β€’ Generate answers β”‚ +β”‚ β€’ GenSelect answers β”‚ β”‚ β€’ GenSelect answers β”‚ +β”‚ β€’ Filter quality β”‚ β”‚ β€’ Evaluate quality β”‚ +β”‚ β”‚ β”‚ β€’ Aggregate results β”‚ +β”‚ β”‚ β”‚ β€’ Post-process β†’ final β”‚ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ -β”‚ Output: ~300K Q&A β”‚ β”‚ Output: ~800K Q&A β”‚ -β”‚ [Used in SFT] β”‚ β”‚ Stratified by difficulty β”‚ +β”‚ Output: ~300K Q&A β”‚ β”‚ Output: (experimental) β”‚ +β”‚ [Used in SFT] β”‚ β”‚ Single final_result.jsonl β”‚ β”‚ β”‚ β”‚ [Work in progress] β”‚ β”‚ β”‚ β”‚ β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ @@ -139,16 +139,6 @@ End-to-end pipeline for generating synthetic financial Q&A data from SEC filings ## Getting Started -### πŸŽ₯ Video Tutorials - -> πŸ“Ή **Coming Soon:** Video walkthroughs of the complete pipeline -> - [ ] Quick Start Demo -> - [ ] Download SEC Filings -> - [ ] Template-Based SDG Explained -> - [ ] Document-Grounded SDG Explained -> - [ ] Model Training & Evaluation -> - [ ] Production Deployment Guide - ### πŸš€ First Time Users **[Quick Start Guide](quick-start.md)** - Run complete demo with 7 companies @@ -172,7 +162,7 @@ Detailed technical specifications for each stage: - **[Template-Based SDG Stages](stages/template-based-sdg.md)** - 6 stages - **[Document-Grounded SDG Stages](stages/document-grounded-sdg.md)** - 7 stages - **[SFT Stages](stages/sft.md)** - 6 stages -- **[Eval Stages](stages/eval.md)** - 9 stages +- **[Eval Stages](stages/eval.md)** - 7 stages - **[GRPO Stages](stages/grpo.md)** - 10 stages ## Quick Command Reference diff --git a/docs/recipes/finance/quick-start.md b/docs/recipes/finance/quick-start.md index 0e85fb5..ffbb71d 100644 --- a/docs/recipes/finance/quick-start.md +++ b/docs/recipes/finance/quick-start.md @@ -66,7 +66,7 @@ Evaluate three baseline models on finance benchmarks to understand pre-fine-tuni uv run nflow list-stages --config nvflow/recipes/finance/workflows/eval/demo.yaml ``` -1. `prepare_data` β€” Prepare benchmark datasets (SecQUE, FinanceBench) into `nvflow/recipes/finance/datasets/` +1. `prepare_data` β€” Prepare benchmark datasets (SecQUE, FinanceBench) into `outputs/finance/eval-datasets/` 2. `qwen3-4b` β€” Evaluate Qwen3-4B on SecQUE and FinanceBench 3. `gemma-3-4b-it` β€” Evaluate Gemma 3 4B IT on SecQUE and FinanceBench 4. `gpt-oss-20b` β€” Evaluate GPT-OSS 20B on SecQUE and FinanceBench @@ -80,12 +80,12 @@ uv run nflow run prepare_data --config nvflow/recipes/finance/workflows/eval/dem Verify the data is ready (expect ~565 SecQUE and ~150 FinanceBench examples): ```bash -wc -l nvflow/recipes/finance/datasets/secque/eval.jsonl nvflow/recipes/finance/datasets/financebench/eval.jsonl +wc -l outputs/finance/eval-datasets/secque/eval.jsonl outputs/finance/eval-datasets/financebench/eval.jsonl ``` **prepare_data output:** ``` -nvflow/recipes/finance/datasets/ # shared across workflows +outputs/finance/eval-datasets/ # shared across workflows β”œβ”€β”€ secque/ β”‚ └── eval.jsonl β”œβ”€β”€ financebench/ @@ -156,6 +156,17 @@ Qwen3-4B and GPT-OSS 20B leverage reasoning (thinking mode and Harmony format re Download 10-K and 10-Q filings for 7 demo companies from SEC EDGAR. The download utility is built into nvflow and uses the `edgartools` library to fetch filings and extract sections. +> **Required first:** SEC EDGAR rejects requests that don't identify the caller, and the config ships with placeholders. Edit the `demo` stage in `nvflow/recipes/finance/workflows/download_sec_filings.yaml` before running β€” there is no command-line override: +> +> ```yaml +> stages: +> demo: +> sec_identity_email: your.email@company.com +> sec_identity_company: YourCompany +> ``` +> +> See the [SEC Fair Access Policy](https://www.sec.gov/os/accessing-edgar-data). + **Preview stages:** ```bash uv run nflow list-stages --config nvflow/recipes/finance/workflows/download_sec_filings.yaml @@ -206,6 +217,19 @@ outputs/finance/demo/workflow-2-download-sec/ Generate financial Q&A pairs using the template-based SDG workflow. +> **Required first:** `create_seed_data` reaches both SEC EDGAR and HuggingFace, so before running: +> +> 1. Set your SEC identity in `nvflow/recipes/finance/workflows/sdg/template-based-sdg.yaml` (inherited by the demo config, and shipped with placeholders): +> +> ```yaml +> stages: +> create_seed_data: +> sec_identity_email: your.email@company.com +> sec_identity_company: YourCompany +> ``` +> +> 2. Temporarily clear `HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE` and `TRANSFORMERS_OFFLINE` in your cluster config, since the seed dataset is pulled from HuggingFace. Re-enable them afterwards. See [Offline runtime](troubleshooting.md#offline-runtime). + **Preview stages:** ```bash uv run nflow list-stages --config nvflow/recipes/finance/workflows/sdg/template-based-sdg-demo.yaml @@ -277,7 +301,7 @@ uv run nflow list-stages --config nvflow/recipes/finance/workflows/sft/qwen3_4b. **Pre-check:** If you skipped Step 1 (baseline eval), ensure benchmark datasets exist: ```bash -wc -l nvflow/recipes/finance/datasets/secque/eval.jsonl nvflow/recipes/finance/datasets/financebench/eval.jsonl +wc -l outputs/finance/eval-datasets/secque/eval.jsonl outputs/finance/eval-datasets/financebench/eval.jsonl # Expected: 565 secque + 150 financebench ``` @@ -398,7 +422,7 @@ Stage 7 (`collect_rollouts`) includes automatic sub-jobs: **Pre-check:** If you skipped Step 1 (baseline eval), ensure benchmark datasets exist: ```bash -wc -l nvflow/recipes/finance/datasets/secque/eval.jsonl nvflow/recipes/finance/datasets/financebench/eval.jsonl +wc -l outputs/finance/eval-datasets/secque/eval.jsonl outputs/finance/eval-datasets/financebench/eval.jsonl # Expected: 565 secque + 150 financebench ``` @@ -449,7 +473,7 @@ uv run nflow run collect_rollouts --config nvflow/recipes/finance/workflows/grpo # Post-rollout train/val split (CPU) uv run nflow run train_validation_split --config nvflow/recipes/finance/workflows/grpo/qwen3_4b.yaml -e finance_sec_search -# Training (Megatron, 64 GPUs β€” uses separate config for YaRN + CP=4) +# Training (Megatron, 16 GPUs β€” uses separate config for YaRN + CP=8) uv run nflow run training --config nvflow/recipes/finance/workflows/grpo/qwen3_4b_finsec.yaml -e finance_sec_search ``` @@ -468,7 +492,7 @@ squeue --me # Rollout logs (one per seed per environment) tail -f outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-5-collect-rollouts/*/logs/*.log # Training logs -tail -f outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-7-training/*/grpo-qwen3-4b-*/training-logs/ray-*-job.log +tail -f outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-8-training/*/grpo-qwen3-4b-*/training-logs/ray-*-job.log ``` **Verify rollouts (both environments):** @@ -494,10 +518,10 @@ ls outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-8-training/equivalence_llm ls outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-8-training/finance_sec_search/grpo-qwen3-4b-*/checkpoints/ ``` -**Verify evaluation:** +**Verify evaluation** (results are per-environment, matching the training checkpoints): ```bash -cat outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-8-eval/step-20/eval-results/secque/metrics.json -cat outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-8-eval/step-20/eval-results/financebench/metrics.json +cat outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-9-eval/finance_sec_search/step-20/eval-results/secque/metrics.json +cat outputs/finance/demo/workflow-5-grpo/qwen3_4b/step-9-eval/finance_sec_search/step-20/eval-results/financebench/metrics.json ``` **Output:** @@ -544,10 +568,16 @@ outputs/finance/demo/workflow-5-grpo/ β”‚ β”‚ β”œβ”€β”€ checkpoints/ β”‚ β”‚ └── training-logs/ β”‚ └── step-9-eval/ -β”‚ └── step-20/ -β”‚ └── eval-results/ -β”‚ β”œβ”€β”€ secque/metrics.json -β”‚ └── financebench/metrics.json +β”‚ β”œβ”€β”€ equivalence_llm_judge/ # Per-env results +β”‚ β”‚ └── step-20/ +β”‚ β”‚ └── eval-results/ +β”‚ β”‚ β”œβ”€β”€ secque/metrics.json +β”‚ β”‚ └── financebench/metrics.json +β”‚ └── finance_sec_search/ +β”‚ └── step-20/ +β”‚ └── eval-results/ +β”‚ β”œβ”€β”€ secque/metrics.json +β”‚ └── financebench/metrics.json ``` > **Per-environment training:** Each environment produces a separate model checkpoint. To train a single combined model on both environments, omit `-e` in the training command. @@ -578,50 +608,8 @@ outputs/finance/demo/workflow-5-grpo/ ## Troubleshooting -
-Download fails with "SEC rate limit" - -SEC EDGAR has rate limits. The downloader includes automatic throttling, but if you hit limits: -- Wait 10 minutes and retry -- Ensure `sec_identity_email` is valid in cluster config - -
- -
-"File not found: sec_metadata.parquet" +See the comprehensive **[Finance Recipe Troubleshooting](troubleshooting.md)** guide for issues across all workflows (SEC rate limits, missing `sec_metadata.parquet`, jobs not starting, eval metrics `N/A`, `Address already in use`, offline-runtime errors, resuming interrupted runs, and more). -Download stage may not have completed. Check logs: -```bash -ls outputs/finance/demo/workflow-2-download-sec/download-logs/ -``` - -
- -
-SFT job not starting - -Check SLURM queue and partition availability: -```bash -squeue --me -sinfo -p interactive -``` - -
- -
-Eval metrics show "N/A" - -Ensure the `checkpoint_path` in your SFT/GRPO config's `stages.eval` section matches your actual training output directory. - -
- -
- -vLLM server crashes with "Address already in use" - -Simply re-run the failed stage. The pipeline will retry only the chunks that did not complete. - -
--- [Workflow Documentation](workflows/) | [Stage Reference](stages/) | [Main README](README.md) diff --git a/docs/recipes/finance/stages/document-grounded-sdg.md b/docs/recipes/finance/stages/document-grounded-sdg.md index 48d86d3..9c9a521 100644 --- a/docs/recipes/finance/stages/document-grounded-sdg.md +++ b/docs/recipes/finance/stages/document-grounded-sdg.md @@ -5,18 +5,18 @@ Technical reference for all 7 stages in the document-grounded-sdg workflow. ## Quick Navigation - [dg_sdg_preprocess](#dg_sdg_preprocess) -- [generate_verified_qa](#generate_verified_qa) -- [genselect_answers](#genselect_answers) +- [generate_verified_questions](#generate_verified_questions) +- [generate_answers](#generate_answers) +- [gym_genselect_answers](#gym_genselect_answers) - [evaluate_answers](#evaluate_answers) - [aggregate_answers](#aggregate_answers) -- [difficulty_estimation](#difficulty_estimation) - [dgsdg_post_process](#dgsdg_post_process) --- ## dg_sdg_preprocess -**File:** `nvflow/recipes/finance/stages/sdg/dg_sdg_preprocess.py` +**File:** `nvflow/generic_stage/sdg/document_grounded/dg_sdg_preprocess.py` **Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="dg_sdg_preprocess"` ### Purpose @@ -34,12 +34,12 @@ Converts raw SEC 10-K and 10-Q HTML filings into structured JSONL data for quest | `input_dir` | path | Raw SEC filings directory (10-K and 10-Q HTML files) | Required | | `output_dir` | path | Preprocessed data output directory | Required | | `distribution_dir` | path | Directory with distribution CSVs (SecQue benchmark) | Required | +| `preprocess_module` | str | Dotted module path to domain CLI that chunks + samples | Required | | `max_tokens` | int | Maximum tokens per chunk | 2000 | | `overlap_tokens` | int | Overlap tokens between chunks for context coverage | 100 | | `total_samples` | int | Total samples to generate following distribution | 150000 | | `max_skip_count` | int | Stop sampling after this many skips (non-repeatable) | 20000 | | `seed` | int | Random seed for reproducibility | 42 | -| `preprocess_kwargs` | dict | Additional CPU job settings (partition, etc.) | `{}` | ### Expected Input Structure @@ -85,7 +85,8 @@ ${output_dir}/ dg_sdg_preprocess: input_dir: ${filings_dir}/data output_dir: ${base_data_dir}/step-0-preprocess - distribution_dir: /workspace/nvflow/recipes/finance/workflows/sdg/dg_sdg_distribution + distribution_dir: nvflow/recipes/finance/workflows/sdg/dg_sdg_distribution + preprocess_module: nvflow.recipes.finance.utils.sdg.dg_sdg_data_preprocess max_tokens: 3000 overlap_tokens: 500 total_samples: 150000 @@ -108,67 +109,61 @@ dg_sdg_preprocess: --- -## generate_verified_qa +## generate_verified_questions -**File:** `nvflow/recipes/finance/stages/sdg/document_grounded_question_answer_generation_pipeline.py` -**Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="generate_verified_qa"` +**File:** `nvflow/generic_stage/sdg/document_grounded/generate_verified_questions.py` +**Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="generate_verified_questions"` ### Purpose -Combined stage that generates questions from SEC filing documents, verifies their quality, and generates answers. Executes 6 internal sub-steps. +Q-side of the DG-SDG pipeline. Generates questions from SEC filing documents and verifies their quality. Executes 4 internal sub-steps. ### Internal Sub-Steps -1. **Preprocess Documents** (CPU): Preprocess sampled data for question generation -2. **Generate Questions** (GPU): Create questions from documents -3. **Preprocess Questions** (CPU): Prepare for verification -4. **Verify Questions** (GPU): Verify quality with 5 random seeds -5. **Preprocess Verified** (CPU): Filter by threshold, prepare for answers -6. **Generate Answers** (GPU): Generate answers with 5 random seeds +1. **Q-prep** (CPU): Run the recipe-supplied `question_prep_script` to attach `context` strings to each chunk +2. **Q-gen** (GPU): Generate questions from documents +3. **Q-verify-prep** (CPU): Expand each generated question into N verification trials +4. **Q-verify** (GPU): Per-question Yes/No vote with multiple random seeds ### Inputs | Parameter | Type | Description | |-----------|------|-------------| | `input_folder` | path | Preprocessed JSONL data directory from `dg_sdg_preprocess` (`${base_data_dir}/step-0-preprocess/jsonl/`) | -| `output_dir` | path | Base output directory for all sub-steps | -| `question_preprocess_kwargs` | dict | CPU job settings for preprocessing | +| `output_dir` | path | Q-pipeline output directory (e.g. `${base_data_dir}/step-1-questions`) | +| `question_prep_script` | path | Domain wrapper that injects `context_builder` into `lib.sdg.document_grounded.preprocess.construct_question_generate_input` | +| `gym_path` / `gym_config_paths` / `gym_agent_name` | various | NeMo-Gym defaults; per-substep `question_generation_*` / `question_verify_*` overrides allowed | | `question_generation_kwargs` | dict | GPU settings for question generation | -| `question_verify_kwargs` | dict | GPU settings for verification (5 seeds) | -| `answer_preprocess_kwargs` | dict | CPU settings, includes `threshold` | -| `answer_generation_kwargs` | dict | GPU settings for answer generation (5 seeds) | +| `question_verify_kwargs` | dict | GPU settings for verification (typically 5 seeds) | ### Outputs ``` ${output_dir}/ -β”œβ”€β”€ question_pipeline/ -β”‚ β”œβ”€β”€ generate_input.jsonl # Preprocessed documents -β”‚ β”œβ”€β”€ generated/ # Generated questions -β”‚ β”‚ β”œβ”€β”€ seed_0.jsonl -β”‚ β”‚ └── ... -β”‚ β”œβ”€β”€ verify_input.jsonl # Questions to verify -β”‚ └── verified/ # Verified questions -β”‚ β”œβ”€β”€ seed_0.jsonl -β”‚ └── ... -└── answer_pipeline/ - β”œβ”€β”€ answer_input.jsonl # Verified questions - └── generated/ # Generated answers ← Output - β”œβ”€β”€ seed_0.jsonl - └── ... +β”œβ”€β”€ generate_input.jsonl # step 1 output (q-prep) +β”œβ”€β”€ generated/ # step 2 output (Q-gen rollouts) +β”œβ”€β”€ verify_input.jsonl # step 3 output (q-verify-prep) +└── verified/ # step 4 output (Q-verify rollouts) + # ← consumed by generate_answers ``` ### Configuration Example ```yaml -generate_verified_qa: +generate_verified_questions: input_folder: ${base_data_dir}/step-0-preprocess/jsonl - output_dir: ${base_data_dir}/step-1-qa-pipeline + output_dir: ${base_data_dir}/step-1-questions + dependencies: [dg_sdg_preprocess] + + question_prep_script: nvflow/recipes/finance/utils/sdg/sec_question_prep.py + gym_path: *gym_path + gym_config_paths: *gym_config_paths_format_verification + gym_agent_name: *gym_agent_format_verification question_generation_kwargs: args: model: /models/gpt-oss-120b - server_gpus: 8 + num_gpus: 8 num_chunks: 5 num_random_seeds: 1 ctx_args: >- @@ -178,42 +173,102 @@ generate_verified_qa: question_verify_kwargs: args: model: /models/Qwen3-235B - server_gpus: 8 + num_gpus: 8 num_chunks: 5 num_random_seeds: 5 +``` + +### Resources + +- **Runtime:** ~2-4 hours +- **GPUs:** 40 for question generation, 200 for question verification +- **Models:** GPT-OSS-120B (questions), Qwen3-235B (verification) + +--- + +## generate_answers + +**File:** `nvflow/generic_stage/sdg/document_grounded/generate_answers.py` +**Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="generate_answers"` + +### Purpose + +A-side of the DG-SDG pipeline. Filters questions by verification pass-rate, then generates N candidate answers per surviving question. Executes 2 internal sub-steps. + +### Internal Sub-Steps + +1. **A-prep** (CPU): `construct_answer_generate_input` keeps only questions whose Q-verify pass-rate β‰₯ `threshold` +2. **A-gen** (GPU): Generate answers (typically 5 seeds for downstream genselect) + +### Inputs + +| Parameter | Type | Description | +|-----------|------|-------------| +| `input_dir` | path | Verified-questions directory from `generate_verified_questions` (`${base_data_dir}/step-1-questions/verified`) | +| `output_dir` | path | A-pipeline output directory (e.g. `${base_data_dir}/step-2-answers`) | +| `gym_path` / `gym_config_paths` / `gym_agent_name` | various | NeMo-Gym defaults; per-substep `answer_generation_*` overrides allowed | +| `answer_preprocess_kwargs` | dict | CPU settings, includes `threshold` (Q-verify pass-rate cutoff) | +| `answer_generation_kwargs` | dict | GPU settings for answer generation | + +### Outputs + +``` +${output_dir}/ +β”œβ”€β”€ answer_input.jsonl # step 1 output (a-prep) +└── generated/ # step 2 output (A-gen rollouts; consumed by gym_genselect_answers) + β”œβ”€β”€ output-rs0.jsonl + └── ... +``` + +### Configuration Example + +```yaml +generate_answers: + input_dir: ${base_data_dir}/step-1-questions/verified + output_dir: ${base_data_dir}/step-2-answers + dependencies: [generate_verified_questions] + + gym_path: *gym_path + gym_config_paths: *gym_config_paths_format_verification + gym_agent_name: *gym_agent_format_verification + + answer_preprocess_kwargs: + threshold: 1 answer_generation_kwargs: args: model: /models/gpt-oss-120b - server_gpus: 8 + num_gpus: 8 num_chunks: 5 num_random_seeds: 5 + ctx_args: >- + ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml ``` ### Resources -- **Total Runtime:** ~4-8 hours for full pipeline -- **GPUs:** 40 for question generation, 200 for question verification and answer generation -- **Models:** GPT-OSS-120B (questions, answers), Qwen3-235B (verification) +- **Runtime:** ~2-4 hours +- **GPUs:** 200 (5 seeds, 5 chunks each) +- **Model:** GPT-OSS-120B --- -## genselect_answers +## gym_genselect_answers -**File:** `nvflow/recipes/finance/stages/sdg/genselect_answers.py` -**Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="genselect_answers"` +**File:** `nvflow/generic_stage/sdg/document_grounded/gym_genselect_answers.py` +**Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="gym_genselect_answers"` ### Purpose -Select best answer from multiple candidates (same as template-based, but for document-grounded data). +Select the best answer from the multiple candidates produced by `generate_answers` (DG-SDG-specific best-of-N picker that runs through NeMo-Gym). ### Inputs | Parameter | Type | Description | |-----------|------|-------------| -| `input_dir` | path | Answer candidates from generate_verified_qa | +| `input_dir` | path | Answer candidates from `generate_answers` | | `output_file` | path | Selected answers output file | -| `prompt_config` | path | GenSelect prompt | +| `prompt_template` | path | GenSelect prompt | ### Outputs @@ -222,19 +277,19 @@ JSONL file with selected best answers. ### Configuration Example ```yaml - genselect_answers: - input_dir: ${base_data_dir}/step-1-qa-pipeline/answer_pipeline/generated - output_file: ${base_data_dir}/step-2-genselect/selected_answers.jsonl - prompt_config: nvflow/recipes/finance/prompts/genselect_answers.yaml - inline_args: "++inference.tokens_to_generate=16384" - dependencies: [generate_verified_qa] - stage_kwargs: - model: /models/Qwen3-235B-A22B-Instruct-2507 - server_type: vllm - server_gpus: 8 + gym_genselect_answers: + input_dir: ${base_data_dir}/step-2-answers/generated + output_file: ${base_data_dir}/step-3-genselect/selected_answers.jsonl + prompt_template: nvflow/recipes/finance/prompts/genselect_answers.yaml + dependencies: [generate_answers] + + policy_vllm: + model_path: /models/Qwen3-235B-A22B-Instruct-2507 + num_gpus: 8 server_nodes: 1 - num_chunks: 15 - partition: batch + num_chunks: 5 + inference_params: + max_output_tokens: 16384 ``` ### Resources @@ -247,39 +302,38 @@ JSONL file with selected best answers. ## evaluate_answers -**File:** `nvflow/recipes/finance/stages/sdg/evaluate_answers.py` +**File:** `nvflow/generic_stage/sdg/document_grounded/evaluate_answers.py` **Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="evaluate_answers"` ### Purpose -Evaluate answer quality using a large model judge. Runs 5 random seeds for robustness. +Evaluate answer quality using a large model judge. Runs 5 random seeds for robustness. Each seed's judge response ends with a JSON verdict tag `{"answerable": "YES/NO", "correct": "YES/NO"}` (parsed downstream by `aggregate_answers`). ### Inputs | Parameter | Type | Description | |-----------|------|-------------| -| `input_file` | path | Selected answers from genselect_answers | +| `input_file` | path | Selected answers from `gym_genselect_answers` | | `output_dir` | path | Directory for evaluation results | -| `prompt_config` | path | Evaluation prompt | +| `prompt_template` | path | Evaluation prompt | ### Outputs ``` ${output_dir}/ -β”œβ”€β”€ seed_0.jsonl -β”œβ”€β”€ seed_1.jsonl -β”œβ”€β”€ seed_2.jsonl -β”œβ”€β”€ seed_3.jsonl -└── seed_4.jsonl +β”œβ”€β”€ output-rs0.jsonl +β”œβ”€β”€ output-rs1.jsonl +β”œβ”€β”€ output-rs2.jsonl +β”œβ”€β”€ output-rs3.jsonl +└── output-rs4.jsonl ``` -Each file contains evaluation scores: +Each record carries the judge's raw `evaluate_generation`, whose last line is the JSON verdict tag parsed by `aggregate_answers`: ```json { - "question": "...", + "problem": "...", "generation": "...", - "evaluate_generation": "Score: 4.5/5\nReasoning: ...", - "evaluation_score": 4.5 + "evaluate_generation": "...reasoning...\n{\"answerable\": \"YES\", \"correct\": \"YES\"}" } ``` @@ -287,19 +341,21 @@ Each file contains evaluation scores: ```yaml evaluate_answers: - input_file: ${base_data_dir}/step-2-genselect/selected_answers.jsonl - output_dir: ${base_data_dir}/step-3-evaluate - prompt_config: nvflow/recipes/finance/prompts/evaluate_answers.yaml - inline_args: "++generation_key=evaluate_generation ++inference.top_p=0.9 ++inference.temperature=0.8" - dependencies: [genselect_answers] - stage_kwargs: - model: /models/Qwen3-235B-A22B-Instruct-2507 - server_type: vllm - server_gpus: 8 + input_file: ${base_data_dir}/step-3-genselect/selected_answers.jsonl + output_dir: ${base_data_dir}/step-4-evaluate + prompt_template: nvflow/recipes/finance/prompts/evaluate_answers.yaml + generation_key: evaluate_generation + dependencies: [gym_genselect_answers] + + policy_vllm: + model_path: /models/Qwen3-235B-A22B-Instruct-2507 + num_gpus: 8 server_nodes: 1 - num_chunks: 5 - num_random_seeds: 5 - partition: batch + num_chunks: 1 + num_random_seeds: 5 + inference_params: + top_p: 0.9 + temperature: 0.8 ``` ### Resources @@ -312,30 +368,29 @@ Each file contains evaluation scores: ## aggregate_answers -**File:** `nvflow/recipes/finance/stages/sdg/aggregate_answers.py` +**File:** `nvflow/generic_stage/sdg/document_grounded/aggregate_answers.py` **Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="aggregate_answers"` ### Purpose -Aggregate evaluation results from 5 random seeds into final scores. +Aggregate the 5 evaluate seeds: keep a question only if **all** seeds voted `correct=YES` with a consistent `answerable`, and attach the consensus `answerable`. ### Inputs | Parameter | Type | Description | |-----------|------|-------------| -| `input_dir` | path | Evaluation results from evaluate_answers | +| `input_dir` | path | Evaluation results from `evaluate_answers` | | `output_file` | path | Aggregated results output | ### Outputs -JSONL file with aggregated scores: +A single JSONL file of surviving records. The per-seed `evaluate_generation` / `correct` are dropped and a consensus `answerable` is added: ```json { - "question": "...", + "problem": "...", "generation": "...", - "evaluation_scores": [4.5, 4.8, 4.3, 4.7, 4.6], - "mean_score": 4.58, - "std_score": 0.18 + "reference_answer": "...", + "answerable": "YES" } ``` @@ -346,134 +401,32 @@ JSONL file with aggregated scores: --- -## difficulty_estimation - -**File:** `nvflow/recipes/finance/stages/sdg/difficulty_estimation.py` -**Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="difficulty_estimation"` - -### Purpose - -Estimate question difficulty by testing if a small model can answer correctly. Questions the small model fails are considered harder. - -### Two-Step Process - -1. **Small Model Answering**: Qwen3-4B attempts to answer (5 seeds) -2. **Large Model Judging**: GPT-OSS-120B judges if small model succeeded - -### Inputs - -| Parameter | Type | Description | -|-----------|------|-------------| -| `input_file` | path | Aggregated answers | -| `output_file` | path | Answers with difficulty scores | -| `work_dir` | path | Working directory for intermediate files | -| `num_random_seeds` | int | Random seeds for small model (default: 5) | -| `answer_model_kwargs` | dict | Settings for small model (Qwen3-4B) | -| `judge_model_kwargs` | dict | Settings for judge model (GPT-OSS-120B) | - -### Outputs - -JSONL file with difficulty scores: -```json -{ - "question": "...", - "generation": "...", - "difficulty_score": 0, # 0 = hard (small model failed) - "small_model_correct": false, - "small_model_attempts": 5, - "small_model_successes": 0 -} -``` - -**Difficulty Score:** -- `0`: Hard (small model failed all attempts) -- `1-4`: Medium (small model succeeded on some attempts) -- `5`: Easy (small model succeeded on all attempts) - -### Configuration Example - -```yaml - difficulty_estimation: - input_file: ${base_data_dir}/step-4-aggregate/aggregated_answers.jsonl - output_file: ${base_data_dir}/step-5-difficulty-data/answers_with_difficulty.jsonl - work_dir: ${base_data_dir}/step-5-difficulty - num_random_seeds: 5 - dependencies: [aggregate_answers] - - - # Small model for answering (Qwen3-4B) - answer_model_kwargs: - args: - model: /models/qwen34b - server_type: vllm - server_gpus: 8 - server_nodes: 1 - num_chunks: 5 - partition: batch - ctx_args: >- - ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml - ++inference.temperature=0.7 - - answer_prompt_config: nvflow/recipes/finance/prompts/secque_template.yaml - - # Large model for judging (GPT-OSS120) - judge_model_kwargs: - args: - model: /models/gpt-oss-120b - server_type: vllm - server_gpus: 8 - server_nodes: 1 - num_chunks: 10 - partition: batch - ctx_args: >- - ++inference.temperature=0.1 - - judge_prompt_config: nvflow/recipes/finance/prompts/judge_difficulty.yaml -``` - -### Resources - -- **GPUs:** 200 (small model), 400 (judge model) -- **Runtime:** 4-8 hours - ---- - ## dgsdg_post_process -**File:** `nvflow/recipes/finance/stages/sdg/document_grounded_data.py` +**File:** `nvflow/generic_stage/sdg/document_grounded/dgsdg_post_process.py` **Registry:** `recipe="finance"`, `workflow="document_grounded_sdg"`, `stage="dgsdg_post_process"` ### Purpose -Clean data and create difficulty-stratified training datasets. +Clean and rename fields, then emit a single `final_result.jsonl` consumed by downstream SFT / GRPO workflows. Records are not split into subsets; the per-stage trim (see `_schemas.py::STAGE_KEEP["dgsdg_post_process"]`) plus the recipe's `domain_keep_fields` defines the final allowlist of fields kept in `final_result.jsonl`. ### Inputs | Parameter | Type | Description | |-----------|------|-------------| -| `input_file` | path | Answers with difficulty from difficulty_estimation | -| `output_dir` | path | Output directory for final datasets | +| `input_file` | path | Aggregated answers from `aggregate_answers` (e.g. `${base_data_dir}/step-5-aggregate/aggregated_answers.jsonl`) | +| `output_dir` | path | Output directory for `final_result.jsonl` | +| `postprocess_script` | path | Domain CLI wrapper around `nvflow.lib.sdg.document_grounded.postprocess.dgsdg_post_process` (e.g. `recipes/finance/utils/sdg/sec_postprocess.py`) | | `seed` | int | Random seed for reproducibility (default: 42) | +| `domain_keep_fields` | list[str] | Recipe-specific fields appended to the generic allowlist before per-stage trim | ### Outputs ``` ${output_dir}/ -β”œβ”€β”€ full_data.jsonl # All cleaned records -β”œβ”€β”€ final_result.jsonl # difficulty_score in [1,2,3,4], filtered -└── hard_rl_data.jsonl # difficulty_score = 0 (hardest) +└── final_result.jsonl # Single cleaned + renamed dataset consumed by SFT / GRPO ``` -**final_result.jsonl** - Medium difficulty training data: -- Medium difficulty questions -- Filtered by filing type and quality -- Ready for training - -**hard_rl_data.jsonl** - Hard difficulty training data: -- Hardest questions (small model failed) -- High-quality answers -- Suitable for advanced training or challenging evaluation - ### Resources - **Compute:** CPU only @@ -485,12 +438,13 @@ ${output_dir}/ | Stage | Purpose | Compute | Runtime | |-------|---------|---------|---------| -| generate_verified_qa | Generate & verify Q&A | GPU | 6-8h | -| genselect_answers | Select best answers | GPU | 1-2h | +| dg_sdg_preprocess | Chunk + sample documents | CPU | 2-4h | +| generate_verified_questions | Generate + verify questions | GPU | 2-4h | +| generate_answers | Generate candidate answers | GPU | 2-4h | +| gym_genselect_answers | Select best answers | GPU | 1-2h | | evaluate_answers | Evaluate quality | GPU | 2-3h | | aggregate_answers | Aggregate scores | CPU | 10m | -| difficulty_estimation | Estimate difficulty | GPU | 2-3h | -| dgsdg_post_process | Create final datasets | CPU | 10m | +| dgsdg_post_process | Clean + rename β†’ final_result.jsonl | CPU | 10m | **Total:** ~10-12 hours for full production run diff --git a/docs/recipes/finance/stages/download-sec.md b/docs/recipes/finance/stages/download-sec.md index e040568..4f016fa 100644 --- a/docs/recipes/finance/stages/download-sec.md +++ b/docs/recipes/finance/stages/download-sec.md @@ -2,12 +2,13 @@ Technical reference for the download-sec workflow stage. -## Stage: sap-500 / demo +## Stage: smoke / demo / sap-500 **File:** `nvflow/recipes/finance/stages/download/download_sec_filings.py` -**Registry:** `recipe="finance"`, `workflow="download-sec"`, `stage="sap-500"` and `stage="demo"` +**Registry:** `recipe="finance"`, `workflow="download-sec"`, `stage="smoke"`, `stage="demo"` and `stage="sap-500"` -> **Note:** Both `sap-500` and `demo` stages use the same implementation but load different configuration files: +> **Note:** All three stages share one implementation and differ only in the ticker config they load: +> - `smoke`: 2 companies, 1 year β€” for pipeline smoke tests > - `demo`: 7 companies (NVDA, AAPL, GOOG, MSFT, CSCO, META, IBM) with 10-K and 10-Q forms (2020-2024) > - `sap-500`: 500+ S&P 500 companies with 10-K, 10-Q, and 8-K forms diff --git a/docs/recipes/finance/stages/eval.md b/docs/recipes/finance/stages/eval.md index 7b5ede7..31e83e4 100644 --- a/docs/recipes/finance/stages/eval.md +++ b/docs/recipes/finance/stages/eval.md @@ -94,12 +94,12 @@ stages: eval: eval_steps: [2600, 5000, 7408] checkpoint_path: ${directories.step-4-training}/model-name - format: megatron # Use "fsdp" for GRPO demo, "megatron" for GRPO production + format: megatron # Match the checkpoint's training backend: "fsdp" or "megatron" baseline_model: /hf_models/Qwen/Qwen3-14B server_type: vllm gpus: 1 inference_args: >- - ++prompt_config=/workspace/nvflow/recipes/finance/prompts/secque_template.yaml + ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml ++inference.temperature=0.6 server_args: "--max-model-len 40960 --async-scheduling --reasoning-parser qwen3" ``` @@ -110,7 +110,7 @@ stages: |-----------|------|-------------| | `eval_steps` | list[int] | Training steps to evaluate | | `checkpoint_path` | path | Base path to training checkpoints | -| `format` | str | `"megatron"` (SFT), `"fsdp"` (GRPO demo), or `"megatron"` (GRPO production) | +| `format` | str | Must match the checkpoint's training backend: `"megatron"` (SFT, finance_sec_search GRPO, production) or `"fsdp"` (equivalence_llm_judge GRPO demo); `"hf"` for HF checkpoints | | `baseline_model` | path | HF model path for baseline comparison | | `server_type` | str | Inference server: `"vllm"`, `"openai"` | | `gpus` | int | GPUs for model server | diff --git a/docs/recipes/finance/stages/finance-agent-eval.md b/docs/recipes/finance/stages/finance-agent-eval.md index 7c0d2fd..cd088a1 100644 --- a/docs/recipes/finance/stages/finance-agent-eval.md +++ b/docs/recipes/finance/stages/finance-agent-eval.md @@ -2,14 +2,12 @@ > **Status:** finance_agent evaluation is currently disabled in `eval/base.yaml` pending further validation. The configuration below is preserved for re-enablement. -Technical reference for the finance-agent evaluation stages (vals-ai/finance-agent benchmark). - -> **Note:** Finance agent evaluation is now integrated into the main eval workflow. The `finance_agent` benchmark is defined in `workflows/eval/base.yaml` and runs alongside SEC-QUE and FinanceBench. See [Eval Workflow](../workflows/05-eval.md) for usage. +Technical reference for the finance-agent evaluation stages (vals-ai/finance-agent benchmark). The stage config is defined in `workflows/eval/base.yaml` but is **currently commented out** (see status above); the reference below applies once it is re-enabled. See [Eval Workflow](../workflows/05-eval.md) for the active benchmarks (SEC-QUE, FinanceBench). ## Quick Navigation - [prepare_data](#prepare_data) -- [agent-gpt-oss-120b](#agent-gpt-oss-120b) +- [Agent eval configuration](#agent-eval-configuration) - [Common Agent Parameters](#common-agent-parameters) --- @@ -21,7 +19,7 @@ Technical reference for the finance-agent evaluation stages (vals-ai/finance-age ### Purpose -The shared `prepare_data` stage now downloads **all** benchmark datasets including `finance_agent`. The `finance_agent` dataset is configured in `workflows/eval/base.yaml` under `benchmarks`. +The shared `prepare_data` stage downloads the **enabled** benchmark datasets (`secque`, `financebench`). `finance_agent` is currently excluded from `dataset_names` in `workflows/eval/base.yaml`; re-add it there when the benchmark is re-enabled. ### Finance Agent Dataset @@ -52,13 +50,13 @@ ${output_dir}/ --- -## agent-gpt-oss-120b - -**Registry:** `recipe="finance"`, `workflow="eval"`, `stage="agent-gpt-oss-120b"` +## Agent eval configuration ### Purpose -Evaluate GPT-OSS-120B as a **multi-turn agent** on the finance-agent benchmark. Uses GENERATION_MODULE from the dataset (`agent_gen`) to run the agent loop with tool calls (Tavily web search, SEC EDGAR, HTML parsing). +Evaluate a model as a **multi-turn agent** on the finance-agent benchmark, using the dataset's GENERATION_MODULE (`agent_gen`) to run the agent loop with tool calls (Tavily web search, SEC EDGAR, HTML parsing). + +Eval stages are derived from the `models:` keys in `eval/*.yaml`, so there is no dedicated agent stage to enable β€” you add a model entry. The block below is a worked example using GPT-OSS-120B; it is not shipped in any config. ### Key Differences from Standard Eval @@ -71,12 +69,12 @@ Evaluate GPT-OSS-120B as a **multi-turn agent** on the finance-agent benchmark. | max_turns | N/A | 50 | | max_concurrent_requests | Parallel | 1 (sequential per question) | -### Configuration (from eval/base.yaml benchmarks section) +### Example model entry ```yaml -agent-gpt-oss-120b: +agent-gpt-oss-120b: # example name; choose your own benchmarks: [finance_agent] - datasets_dir: /workspace/nvflow/recipes/finance/datasets + datasets_dir: /workspace/outputs/finance/eval-datasets judge: *judge_finance_strict installation_command: "pip install -q model-library==0.1.8 func-timeout backoff tavily compute-eval @ git+..." extra_args: >- @@ -101,7 +99,7 @@ agent-gpt-oss-120b: ### Resources - **GPUs:** 8 (120B model) -- **Judge:** GPT-5.1 via OpenAI API (external) +- **Judge:** `gpt-5-mini` via OpenAI API (external) - **Tools:** Tavily API (web search), compute-eval for tool execution - **Runtime:** Longer than single-turn (multi-turn + tool calls) @@ -114,7 +112,7 @@ agent-gpt-oss-120b: | Parameter | Description | |-----------|-------------| | `installation_command` | Pip install model-library, func-timeout, tavily, compute-eval | -| `judge` | `judge_finance_strict` (GPT-5.1, sec_judge_strict.yaml) | +| `judge` | `judge_finance_strict` (`gpt-5-mini`, sec_judge_strict.yaml) | | `extra_args.max_turns` | Max agent turns per question (default: 50) | | `extra_args.max_concurrent_requests` | 1 (sequential to avoid API rate limits) | | `rollouts.extra_args.prompt_format` | `openai` (OpenAI function-calling format) | @@ -122,7 +120,7 @@ agent-gpt-oss-120b: ### Judge (judge_finance_strict) Strict finance-domain judge matching vals-ai/finance-agent's judge_new.py: -- **Model:** GPT-5.1 +- **Model:** `gpt-5-mini` - **Prompt:** `sec_judge_strict.yaml` (domain tolerance rules, few-shot examples) - **Temperature:** 0.0 - **Skip extraction:** Yes (judgement only) @@ -156,7 +154,7 @@ models: gpus: 2 nodes: 1 inference_args: >- - ++prompt_config=/workspace/nvflow/recipes/finance/prompts/secque_template.yaml + ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml ++inference.tokens_to_generate=32768 ++inference.temperature=0.0 server_args: "--max-model-len 65536 --async-scheduling" @@ -174,7 +172,7 @@ ls outputs/finance/sap-500/workflow-1-baseline-eval/baselines/gpt-oss-120b/eval- cat outputs/finance/sap-500/workflow-1-baseline-eval/baselines/gpt-oss-120b/eval-results/finance_agent/metrics.json | jq . # Check prepared dataset -ls /workspace/nvflow/recipes/finance/datasets/finance_agent/ +ls /workspace/outputs/finance/eval-datasets/finance_agent/ ``` --- @@ -199,4 +197,4 @@ ls /workspace/nvflow/recipes/finance/datasets/finance_agent/ --- -See [Finance Agent Benchmark](../workflows/06-finance-agent-eval.md) for an overview, or [Eval Workflow](../workflows/05-eval.md) for full usage examples and configuration. +See [Eval Workflow](../workflows/05-eval.md) for full usage examples and configuration. diff --git a/docs/recipes/finance/stages/grpo.md b/docs/recipes/finance/stages/grpo.md index f22bd34..a425274 100644 --- a/docs/recipes/finance/stages/grpo.md +++ b/docs/recipes/finance/stages/grpo.md @@ -1,18 +1,54 @@ # GRPO Stages Reference -Technical reference for all 10 stages in the GRPO RL training workflow (9 active + 1 optional). +Technical reference for the GRPO RL training workflow: 10 active stages plus `compute_rewards`, which is optional and commented out by default. + +> **Pass `-e `.** A model config's `environments` block *merges* with `grpo/base.yaml` rather than replacing it, and `base.yaml` declares three environments (`equivalence_llm_judge`, `mcqa`, `finance_sec_search`). Running a single-environment config without `-e` trains all three jointly, including `mcqa`, which is a placeholder with `raw_train_data: null` and is not runnable. ## Quick Navigation -- [data_transformation](#data_transformation) -- [apply_prompt_template](#apply_prompt_template) -- [convert_to_responses_api](#convert_to_responses_api) -- [train_validation_split](#train_validation_split) -- [prepare_data](#prepare_data) -- [collect_rollouts](#collect_rollouts) -- [compute_rewards](#compute_rewards) -- [training](#training) -- [eval](#eval) +Listed in execution order. `prefetch_cache` is optional and has no `step-N` directory. + +- [validate_questions](#validate_questions) β€” step 0 +- [data_transformation](#data_transformation) β€” step 1 +- [apply_prompt_template](#apply_prompt_template) β€” step 2 +- [convert_to_responses_api](#convert_to_responses_api) β€” step 3 +- [prepare_data](#prepare_data) β€” step 4 +- [prefetch_cache](#prefetch_cache) β€” optional +- [collect_rollouts](#collect_rollouts) β€” step 5 +- [compute_rewards](#compute_rewards) β€” step 6, optional +- [train_validation_split](#train_validation_split) β€” step 7 +- [training](#training) β€” step 8 +- [eval](#eval) β€” step 9 + +--- + +## validate_questions + +**File:** `nvflow/recipes/finance/stages/rl/validate_questions.py` +**Registry:** `recipe="finance"`, `workflow="grpo"`, `stage="validate_questions"` + +### Purpose + +Drop structurally-broken SDG questions before they enter the pipeline, per environment, in two phases: + +1. **Regex prefilter (CPU).** Drops questions that say "the company" / "the firm" with no named company or ticker anywhere in the text. Deliberately narrow β€” recall over precision. +2. **LLM classifier (GPU).** Asks a judge model (GPT-OSS-120B by default) for `VALID` / `INVALID` on each survivor. Parse failures default to `VALID`. + +The kept stream is written where `data_transformation` can read it, so a model config re-points `env.raw_train_data` at this stage's output. + +### Outputs + +``` +${step-0-validate-questions}/${env_name}/ +β”œβ”€β”€ final_result.jsonl # VALID records, consumed by data_transformation +β”œβ”€β”€ phase1_regex/ # prefiltered + dropped + stats (audit) +└── phase2_llm/ # raw generation, parsed tags, dropped, stats +``` + +### Resources + +- **Phase 1:** CPU only +- **Phase 2:** GPU, for the judge model --- @@ -154,7 +190,7 @@ Split data into training and validation sets using stratified sampling to mainta ### Purpose -Run `ng_prepare_data` to stamp each JSONL record with an `agent_ref` field that tells NeMo-Gym which agent server to route the example to during training. Auto-generates an agent config overlay YAML from the workflow's `agents` list. +Run `gym dataset collate` (formerly `ng_prepare_data`) to stamp each JSONL record with an `agent_ref` field that tells NeMo-Gym which agent server to route the example to during training. Auto-generates an agent config overlay YAML from the workflow's `agents` list. ### Inputs @@ -171,7 +207,7 @@ Run `ng_prepare_data` to stamp each JSONL record with an `agent_ref` field that ### Modes -- **`train_preparation`**: Produces `train.jsonl` + `validation.jsonl` +- **`train_preparation`**: Produces `train.jsonl`. Only a single `train` dataset is collated here; the train/validation split happens later, in [train_validation_split](#train_validation_split), on reward-filtered data. ### Agent Configuration @@ -192,13 +228,13 @@ agents: - name: train type: train license: "TBD" - jsonl_fpath: ${directories.step-3-train-validation-split}/train.jsonl + jsonl_fpath: ${directories.step-3-convert-to-responses-api}/train.jsonl ``` ### Outputs - `${output_dir}/agent_config_overlay.yaml` β€” Auto-generated agent config -- `${output_dir}/train.jsonl` + `validation.jsonl` β€” with `agent_ref` routing fields +- `${output_dir}/train.jsonl` β€” with `agent_ref` routing fields ### Resources @@ -207,6 +243,39 @@ agents: --- +## prefetch_cache + +**File:** `nvflow/recipes/finance/stages/rl/prefetch_cache.py` +**Registry:** `recipe="finance"`, `workflow="grpo"`, `stage="prefetch_cache"` + +### Purpose + +Optional CPU-only stage that populates the SEC filing metadata cache before rollout collection. Doing it here keeps SEC.gov calls out of the GPU-intensive rollout jobs and avoids races when several seeds share one cache directory. + +It runs per environment and processes only those whose config carries a `prefetch` block; the rest are skipped silently. In practice that means `finance_sec_search`. + +### Inputs + +Read from each environment's `prefetch` block: + +| Key | Description | +|-----|-------------| +| `script` | Upstream Gym prefetch script to run | +| `cache_dir` | Where the cache is written | +| `ticker_config` | Ticker set to prefetch | +| `force` | Re-fetch even if the cache is populated (default `false`) | + +### Outputs + +The cache directory declared by the environment. For `finance_sec_search` this is `cache-finance-sec-search`, i.e. `${base_output_dir}/cache/finance_sec_search`, holding `filings/`, `filings_metadata/` and `tickers.json`. + +### Resources + +- **Compute:** CPU only +- **Network:** needs SEC EDGAR access, so run it on a connected node + +--- + ## collect_rollouts **File:** `nvflow/recipes/finance/stages/rl/collect_rollouts.py` @@ -218,38 +287,57 @@ Collect model rollouts against a NeMo-Gym environment with reward scoring. Suppo ### Inputs +Top-level keys are orchestration; rollout behaviour is nested under `rollout`. + | Parameter | Type | Description | Default | |-----------|------|-------------|---------| | `output_dir` | path | Output directory | Required | -| `gym_path` | path | Path to NeMo-Gym | Required | -| `container` | string | Container name | Required | -| `input_data` | path | Prepared JSONL from prepare_data | Required | -| `agent_name` | string | Agent name (must match prepare_data) | Required | -| `model_path` | path | Model to collect rollouts from | Required | -| `nemo_gym_config_paths` | list | NeMo-Gym config paths | Required | -| `num_repeats` | int | Repeats per sample | `1` | -| `num_samples_in_parallel` | int | Concurrent requests | `4` | +| `prepare_data_dir` | path | Collated data from `prepare_data` | Required | +| `gym_path` | path | NeMo-Gym root inside the container | `/opt/Gym` | +| `gym_uv_venv_dir` | path | Baked per-component venvs reused by `ng_run` | `/opt/gym-venvs` | +| `container` | string | Rollout client + Gym env servers (CPU) | `nemo-gym` | +| `postprocess_container` | string | Merge/analyze/aggregate/filter (CPU) | `nemo-skills` | +| `vllm_container` | string | Policy and judge vLLM servers (GPU) | `vllm-grpo` | +| `environments` | dict | Environments to collect for | `${environments}` | + +**`rollout`** β€” job fan-out and per-request settings: + +| Parameter | Type | Description | Default | +|-----------|------|-------------|---------| +| `num_samples_in_parallel` | int | Concurrent requests | `64` | +| `max_num_samples` | int | Truncate to first N rows; `null` for all | `null` | | `num_chunks` | int | Split input into N parallel jobs | `1` | -| `num_random_seeds` | int | Independent runs per chunk | `1` | +| `num_random_seeds` | int | Independent runs per chunk | `8` | | `starting_seed` | int | First seed value | `0` | -| `dependent_jobs` | int | Chain N+1 Slurm jobs per chunk via `afterany` for timeout recovery | `0` | -| `responses_create_params` | dict | Pass-through params for NeMo-Gym (e.g., `max_output_tokens`) | `{}` | +| `dependent_jobs` | int | Chain N+1 jobs per (seed, chunk) for timeout resume | `0` | | `rerun_done` | bool | Force re-execution | `false` | -| `num_gpus` | int | GPUs per Slurm job | `8` | -| `tensor_parallel_size` | int | Policy vLLM TP | `2` | +| `responses_create_params` | dict | Per-request overrides, e.g. `max_output_tokens` | `{}` | + +**`rollout.policy_vllm`** β€” the policy server, shared across environments. `num_gpus`, `server_nodes`, `base_url` and `model_path` are orchestration-only; every other key becomes a `--key value` argument to `vllm serve`. + +| Parameter | Type | Description | Default | +|-----------|------|-------------|---------| +| `model_path` | path | Model to serve | Required, set in the model config | +| `num_gpus` | int | Slurm GPUs for this endpoint; `0` with `base_url` for an external server | `2` | +| `server_nodes` | int | Nodes for this vLLM; `>1` uses Ray | `1` | | `max_model_len` | int | Max sequence length | `32768` | -| `vllm_base_url` | string | External vLLM URL (optional) | None | +| `enable_auto_tool_choice` | bool | Required for tool-calling environments | `true` | +| `tool_call_parser` | string | Tool-call parser | `hermes` | + +> **Don't set `tensor_parallel_size`.** It is derived from `num_gpus` and is silently ignored here. ### Judge Configuration -| Parameter | Type | Description | -|-----------|------|-------------| -| `judge_model_path` | path | Local vLLM judge model | -| `judge_tensor_parallel_size` | int | Judge TP size | -| `judge_max_model_len` | int | Judge max sequence length | -| `judge_openai_base_url` | string | External OpenAI API URL | -| `judge_openai_model` | string | OpenAI model name | -| `judge_openai_api_key` | string | API key override (defaults to `$OPENAI_API_KEY`) | +The judge is configured **per environment**, not on the stage, because each environment decides whether it needs one: + +```yaml +environments: + finance_sec_search: + judge_vllm: + num_gpus: 0 # 0 means no local judge -- override in the model config +``` + +Set `num_gpus` above zero to stand up a local judge vLLM for that environment, and use `responses_create_params` alongside it to override the shared rollout defaults. ### Execution Model @@ -419,7 +507,7 @@ The stage validates parallelism before job submission: ### Outputs ``` -${output_dir}/grpo-{model}-{nodes}n-tp{tp}-cp{cp}-seq{seq}k/ +${output_dir}/grpo-{model}-{total_gpus}g-tp{tp}-cp{cp}-seq{seq}k/ β”œβ”€β”€ checkpoints/ β”‚ β”œβ”€β”€ step_1/ β”‚ └── step_2/ @@ -427,12 +515,14 @@ ${output_dir}/grpo-{model}-{nodes}n-tp{tp}-cp{cp}-seq{seq}k/ └── run_metadata_*.yaml # Full config for reproducibility ``` +The directory name is built from the resolved layout, so `grpo-qwen3-4b-16g-tp2-cp1-seq32k` means 16 GPUs total, TP=2, CP=1 and a 32K sequence budget. + ### Resources | Model Size | GPUs | Runtime (demo) | |------------|------|----------------| -| 4B | 16 (2 nodes) | ~20 min | -| 14B | 64 (8 nodes) | TBD | +| 4B | 16 | ~20 min | +| 30B-A3B | 64 | Longer; see `grpo/qwen3_30b_a3b.yaml` | --- @@ -454,7 +544,7 @@ Also registered for the SFT workflow, making it a shared evaluation stage across | `eval_output_dir` | path | Output directory for evaluation results | Required | | `eval_steps` | list | Checkpoint steps to evaluate | `[]` | | `checkpoint_path` | path | Path to training checkpoints | Required | -| `format` | string | Checkpoint format: `"hf"`, `"fsdp"`, `"megatron"` | `"fsdp"` (demo) / `"megatron"` (production) | +| `format` | string | Checkpoint format; match the training backend: `"hf"`, `"fsdp"` (equivalence demo), `"megatron"` (finance_sec_search demo + production) | backend-dependent | | `baseline_model` | path | Baseline model for comparison evaluation | Optional | | `server_type` | string | Inference server type | `"vllm"` | | `gpus` | int | GPUs for inference server | `1` | diff --git a/docs/recipes/finance/stages/sft.md b/docs/recipes/finance/stages/sft.md index 312d5fe..ef95ca9 100644 --- a/docs/recipes/finance/stages/sft.md +++ b/docs/recipes/finance/stages/sft.md @@ -204,7 +204,7 @@ Group training examples by total sequence length (input + output tokens) to redu | `input_file` | path | Training data from train_validation_split | Required | | `output_dir` | path | Directory for grouped/bucketed data | Required | | `tokenizer_path` | path | Tokenizer for computing lengths (optional if pre-computed) | None | -| `bucket_sizes` | list | Token length boundaries for buckets | `[16000, 32000, 64000]` | +| `bucket_sizes` | list | Token length boundaries for buckets | `[16000, 24000, 32000, 48000]` | ### Bucket Configuration @@ -255,43 +255,68 @@ Fine-tune the language model on financial Q&A data using supervised learning. ### Inputs +Training uses NeMo-RL's config schema: pick a `preset`, then patch it through `overrides`, which is passed to NeMo-RL nested and unflattened. + | Parameter | Type | Description | |-----------|------|-------------| -| `model_name_or_path` | path | Base model to fine-tune | -| `train_file` | path | Training data | -| `val_file` | path | Validation data | -| `output_dir` | path | Directory for checkpoints and logs | -| `num_train_epochs` | int | Number of training epochs (default: 3) | -| `learning_rate` | float | Learning rate (default: 2e-5) | -| `per_device_train_batch_size` | int | Batch size per GPU (default: 4) | -| `gradient_accumulation_steps` | int | Gradient accumulation (default: 8) | -| `save_steps` | int | Checkpoint save frequency (default: 500) | -| `eval_steps` | int | Evaluation frequency (default: 500) | +| `model_name` | string | Model identifier, e.g. `Qwen/Qwen3-14B` | +| `hf_checkpoint_path` | path | Base model on disk, e.g. `/hf_models/Qwen/Qwen3-14B` | +| `backend` | string | `megatron` or `dtensor` | +| `total_gpus` | int | GPUs for the job; data parallelism is derived from it | +| `dependent_jobs` | int | Extra chained jobs, for training longer than one time limit | +| `preset` | string | Base config to start from, e.g. `sft-base` | +| `overrides` | dict | Nested patch over the preset, grouped into `sft`, `checkpointing`, `policy` and `data` | + +Commonly overridden keys: + +| Key | Description | +|-----|-------------| +| `sft.max_num_epochs` | Number of epochs | +| `sft.val_period` | Validate every N steps | +| `checkpointing.save_period` | Save every N steps | +| `checkpointing.keep_top_k` | Checkpoints to retain | +| `policy.train_global_batch_size` | Global batch size | +| `policy.train_micro_batch_size` | Per-rank micro batch | +| `policy.max_total_sequence_length` | Sequence budget | +| `policy.megatron_cfg.*` | Parallelism (`tensor_model_parallel_size`, `context_parallel_size`, …) | +| `policy.megatron_cfg.optimizer.lr` | Learning rate | ### Training Configuration ```yaml -training: - learning_rate: 2e-5 - global_batch_size: 128 - max_num_epochs: 5 +stages: + training: + model_name: Qwen/Qwen3-14B + hf_checkpoint_path: /hf_models/Qwen/Qwen3-14B + backend: megatron + total_gpus: 256 + preset: "sft-base" + overrides: + sft: + max_num_epochs: 3 + policy: + train_global_batch_size: 128 + max_total_sequence_length: 49152 + megatron_cfg: + tensor_model_parallel_size: 4 + context_parallel_size: 8 + optimizer: + lr: 5e-6 ``` ### Outputs ``` -${output_dir}/ +${output_dir}/model-{model}-{total_gpus}g-tp{tp}-pp{pp}-cp{cp}-seq{seq}k/ β”œβ”€β”€ checkpoints/ -β”‚ β”œβ”€β”€ checkpoint-500/ -β”‚ β”œβ”€β”€ checkpoint-1000/ -β”‚ β”œβ”€β”€ checkpoint-1500/ -β”‚ └── final/ # ← Final model -β”œβ”€β”€ logs/ -β”‚ └── training.log -β”œβ”€β”€ runs/ # Tensorboard logs -└── training_args.json +β”‚ β”œβ”€β”€ step_10/ +β”‚ └── step_20/ +β”œβ”€β”€ training-logs/ +└── run_metadata_*.yaml ``` +Checkpoints are step-numbered; there is no `final/` directory. The `eval` stage converts a chosen step to HuggingFace format when it needs one. + ### Resources | Model Size | GPUs | Memory/GPU | Runtime | @@ -401,24 +426,20 @@ Convert Qwen3 chat-templated training data to OpenAI messages format. Parses Qwe ## Common Training Parameters +All of these live under `overrides` in the training stage. + ### Learning Rate -| Model Size | Recommended LR | -|------------|----------------| -| 7-14B | 2e-5 | -| 32B | 1e-5 | -| 70B+ | 5e-6 | +Set at `policy.megatron_cfg.optimizer.lr`. The shipped configs use `5e-6` with `min_lr: 5e-7`, cosine decay, and warmup from `1e-7`. Treat `5e-6` as the starting point rather than scaling by model size. ### Batch Size -Effective batch size = `per_device_train_batch_size` Γ— `gradient_accumulation_steps` Γ— `total_gpus` - -Recommended: 32-128 for most models +`policy.train_global_batch_size` is the global batch, and `policy.train_micro_batch_size` the per-rank micro batch; gradient accumulation is derived from the two together with the data-parallel width. The production 14B config uses `128` global and `1` micro. ### Checkpointing -- **save_steps**: 500-1000 (more frequent for smaller datasets) -- **save_total_limit**: 3-5 (keep only recent checkpoints to save space) -- **eval_steps**: Same as save_steps +- **`checkpointing.save_period`**: save every N steps β€” `100` for full training, `10` in the demo +- **`checkpointing.keep_top_k`**: checkpoints to retain +- **`sft.val_period`**: validate every N steps See [SFT Workflow](../workflows/04-sft.md) for usage examples and configuration details. diff --git a/docs/recipes/finance/troubleshooting.md b/docs/recipes/finance/troubleshooting.md index 3e53e9e..8da6307 100644 --- a/docs/recipes/finance/troubleshooting.md +++ b/docs/recipes/finance/troubleshooting.md @@ -5,11 +5,14 @@ Comprehensive troubleshooting guide for common issues across all finance recipe ## Quick Navigation - [Cluster & Infrastructure](#cluster--infrastructure) -- [Offline Runtime](#self-sufficient-runtime) +- [Offline Runtime](#offline-runtime) - [Resource Issues](#resource-issues) - [Data Issues](#data-issues) - [Training Issues](#training-issues) - [Workflow-Specific Issues](#workflow-specific-issues) +- [Resuming Interrupted Workflows](#resuming-interrupted-workflows) +- [Frequently Asked Questions](#frequently-asked-questions) +- [Getting Additional Help](#getting-additional-help) --- @@ -92,59 +95,62 @@ scontrol show config | grep SLURM_VERSION # Confirmed: SLURM 25.11.2 needs this fix, SLURM 24.x works without it ``` -**Additional notes:** -- If Ray cluster hangs during initialization, apply this fix -- The fix changes how containers are executed (uses `enroot exec` instead of `--container-name`) -- Test on your cluster - symptom is Ray cluster initialization hang +This changes how containers are launched, using `enroot exec` instead of `--container-name`. --- ## 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. +The NVFlow images (`nvflow-nemo-skills`, `nvflow-vllm` at both tags, `nvflow-nemo-gym`, `nvflow-nemo-rl`) run with **no outbound network access** at job time, including the `training` stage: `nvflow-nemo-rl` bakes the Gym venvs at build time, so nothing needs resolving over the network. `UV_OFFLINE` is nonetheless left **unset**, which preserves dev mode: mount local Gym source and `uv` resolves it. 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` +### `huggingface_hub.errors.OfflineModeIsEnabled` / `LocalEntryNotFoundError` -**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: +**Problem:** A stage fails trying to pull a model or dataset from HuggingFace Hub. -``` -bash: /opt/NeMo-RL/3rdparty/Gym-workspace/Gym/.venv/bin/activate: No such file or directory -``` +**Cause:** Air-gap mode is on (`HF_HUB_OFFLINE=1`, etc.) but the asset isn't pre-staged on disk. -**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:** +- **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`. The artifacts persist under `/workspace` and are reused by every subsequent run. -**Solution:** Remove the overlay mounts from `cluster_configs/my_cluster.yaml`. The self-sufficient image already contains everything GRPO needs: +### GRPO `training` fails: `uv` tries to resolve, or `ng_run` / `nemo_gym` not found -```yaml -mounts: - # COMMENT THESE OUT (or delete) for normal production runs: - # - :/opt/NeMo-RL - # - :/opt/NeMo-RL/3rdparty/Gym-workspace/Gym -``` +**Problem:** The `training` stage fails soon after start with `uv` trying to download packages, a hung resolution, or a missing Gym module. -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. +**Cause:** Nothing should resolve at runtime β€” `nvflow-nemo-rl` bakes one Gym venv per component. A resolve attempt means those baked venvs aren't the ones in use, which has two usual causes: a host clone bind-mounted over `/opt/nemo-rl/3rdparty/Gym-workspace/Gym`, shadowing the baked source and venvs; or the job running the stock upstream `nemo-rl` base, which ships the RL environment but leaves the Gym venvs unbuilt. -### `huggingface_hub.errors.OfflineModeIsEnabled` / `LocalEntryNotFoundError` +**Solution:** +1. Confirm `containers.nemo-rl` in your cluster config points at the image built from [`dockerfiles/Dockerfile.nemo-rl`](../../../dockerfiles/Dockerfile.nemo-rl), not the stock base. +2. Remove any Gym or NeMo-RL source mount from the `mounts:` block. +3. Only if you are deliberately running dev mode against mounted source: leave `UV_OFFLINE` unset and confirm the compute nodes can reach a pypi mirror. See [`docs/development/nemo-rl-gym.md`](../../development/nemo-rl-gym.md). -**Problem:** A stage fails trying to pull a model or dataset from HuggingFace Hub. +The Gym-only stages (`collect_rollouts`, `compute_rewards`, `prefetch_cache`, `prepare_data`) instead run on the self-contained `nvflow-nemo-gym` image (baked venvs, no build); if one of those reports `ng_run: command not found`, the image is missing its baked venvs -- re-check the nemo-gym build in [`docs/maintainers/containers.md`](../../maintainers/containers.md). -**Cause:** Air-gap mode is on (`HF_HUB_OFFLINE=1`, etc.) but the asset isn't pre-staged on disk. +### `omegaconf.errors.InterpolationKeyError: Interpolation key '' not found` after mounting a Gym branch -**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. +**Problem:** A `training` or `ng_run`-driven job fails at NeMo-Gym config-load time with, e.g.: -### `uv` errors with "package not installed" or tries to resolve from PyPI +``` +omegaconf.errors.InterpolationKeyError: Interpolation key 'tavily_api_key' not found + full_key: tavily_api_key + object_type=dict +``` + +**Cause:** The Gym source introduced a new `${}` interpolation in a resource-server YAML that the overlays under `nvflow/recipes/finance/workflows/grpo/overlays/` don't yet define. This is drift between the Gym source and the overlays, not a runtime requirement β€” the runtime treats the value as optional (an empty `tavily_api_key` disables Tavily web_search gracefully). -**Problem:** A Ray worker or stage script fails because `uv` is trying to download a package. +**Solution (clean, no upstream change):** Add a placeholder for the missing key in the relevant overlay under `nvflow/recipes/finance/workflows/grpo/overlays/`. For `tavily_api_key` specifically, that's `finance_sec_search_env.yaml`: -**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. +```yaml +# Required since upstream Gym introduced ${tavily_api_key} in finance_sec_search.yaml. +# Empty string disables tavily gracefully -- finance_sec_search uses SEC tools only. +tavily_api_key: "" +``` -**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). +Restart the job; OmegaConf will resolve the interpolation against the overlay value and the resource server will log `No tavily_api_key configured β€” web_search will be unavailable` and continue. -**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). +If this happens for a key other than `tavily_api_key`, the same recipe applies: identify which Gym resource-server YAML references the new `${}`, and add the corresponding overlay placeholder under `nvflow/recipes/finance/workflows/grpo/overlays/`. ### `tiktoken` / `openai_harmony` fails to load offline @@ -163,7 +169,7 @@ env_vars: Verify the cache exists inside the image: ```bash -docker run --rm nvflow-nemo-skills:0229040 ls /opt/tiktoken_cache +docker run --rm nvflow-nemo-skills:v1.1.2 ls /opt/tiktoken_cache # Expect: cl100k_base.tiktoken (and o200k_base.tiktoken in vllm images) ``` @@ -175,7 +181,7 @@ docker run --rm nvflow-nemo-skills:0229040 ls /opt/tiktoken_cache **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 \ +docker run --rm nvflow-nemo-skills:v1.1.2 bash -c \ 'python3 -c "import pyarrow as pa; pa.array([], type=pa.timestamp(\"ns\", tz=\"UTC\")); print(\"OK\")"' ``` diff --git a/docs/recipes/finance/workflows/02-template-based-sdg.md b/docs/recipes/finance/workflows/02-template-based-sdg.md index a41d8dc..8c6729a 100644 --- a/docs/recipes/finance/workflows/02-template-based-sdg.md +++ b/docs/recipes/finance/workflows/02-template-based-sdg.md @@ -18,6 +18,7 @@ Before running this workflow, ensure you have: - **Why needed:** Stage 0 downloads the [SecQue dataset](https://huggingface.co/datasets/nvidia/SecQue) (seed questions) from HuggingFace - **Public dataset:** No token required for public access, but token avoids rate limits - **Login alternative:** Run `huggingface-cli login` if you prefer interactive login + - **Offline clusters:** Because Stage 0 reaches the Hub, temporarily clear `HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE` and `TRANSFORMERS_OFFLINE` in your cluster config for this run, then restore them. See [Offline runtime](../troubleshooting.md#offline-runtime) - βœ… **SEC EDGAR identity configured** in workflow YAML: ```yaml diff --git a/docs/recipes/finance/workflows/03-document-grounded-sdg.md b/docs/recipes/finance/workflows/03-document-grounded-sdg.md index 3e0e621..7afd9b8 100644 --- a/docs/recipes/finance/workflows/03-document-grounded-sdg.md +++ b/docs/recipes/finance/workflows/03-document-grounded-sdg.md @@ -2,13 +2,13 @@ ## Purpose -Generate high-quality financial Q&A pairs directly from SEC filing documents with built-in verification, evaluation, and difficulty estimation. +Generate high-quality financial Q&A pairs directly from SEC filing documents with built-in question verification, multi-seed answer evaluation, and per-stage field trimming. -> **Note:** This workflow generates ~800K Q&A pairs. SFT integration is currently in progress. For production SFT pipeline, see [Template-Based SDG](02-template-based-sdg.md). +> **Note:** This workflow generates ~800K Q&A pairs in a single `final_result.jsonl`. The previous difficulty-stratified outputs (`full_data.jsonl`, `hard_rl_data.jsonl`) and the `difficulty_estimation` stage have been removed; downstream SFT / GRPO workflows read `final_result.jsonl` directly. For the production template-based pipeline, see [Template-Based SDG](02-template-based-sdg.md). ## Prerequisites -- βœ… SEC filings downloaded ([Workflow 1](01-download-sec.md)) +- SEC filings downloaded ([Workflow 1](01-download-sec.md)) - Will be preprocessed in Stage 0 (dg_sdg_preprocess) ## Key Differences from Template-Based @@ -18,55 +18,54 @@ Generate high-quality financial Q&A pairs directly from SEC filing documents wit | **Question Source** | Seed questions | Generated from documents | | **Verification** | None | Built-in verification step | | **Quality Control** | GenSelect + Filter | GenSelect + Evaluation + Aggregation | -| **Difficulty** | Not estimated | Estimated via small model testing | -| **Output** | Single dataset | Stratified by difficulty (medium/hard) | +| **Output** | Single dataset | Single `final_result.jsonl` (no stratification) | ## Pipeline Flow ``` -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 0. dg_sdg_preprocess β”‚ Preprocessing: SEC HTML β†’ Chunked JSONL -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β–Ό -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 1. generate_verified_qa β”‚ Q&A Generation: Questions + Answers -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β–Ό -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 2. genselect_answers β”‚ Selection: Best answer from candidates -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β–Ό -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 3. evaluate_answers β”‚ Evaluation: Quality scoring (5 seeds) -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β–Ό -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 4. aggregate_answers β”‚ Aggregation: Combine evaluation results -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β–Ό -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 5. difficulty_estimationβ”‚ Difficulty: Small model testing -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ - β”‚ - β–Ό -β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 6. dgsdg_post_process β”‚ Output: Stratified training datasets -β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 0. dg_sdg_preprocess β”‚ Preprocessing: SEC HTML β†’ Chunked JSONL +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 1. generate_verified_questionsβ”‚ Q-pipeline: prep + Q-gen + verify-prep + Q-verify +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 2. generate_answers β”‚ A-pipeline: a-prep (threshold filter) + A-gen +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 3. gym_genselect_answers β”‚ Selection: Best answer from candidates +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 4. evaluate_answers β”‚ Evaluation: Quality scoring (multi-seed) +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 5. aggregate_answers β”‚ Aggregation: Combine evaluation results +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ 6. dgsdg_post_process β”‚ Output: Cleaned + renamed β†’ final_result.jsonl +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ ``` ## 7 Stages (Overview) 0. **dg_sdg_preprocess**: Preprocess SEC filings (chunk HTML β†’ create JSONL data following SecQue distribution) -1. **generate_verified_qa**: Generate questions from documents, verify them, generate answers (6 internal sub-steps) -2. **genselect_answers**: Select best answer from multiple candidates -3. **evaluate_answers**: Evaluate answer quality (5 random seeds for robustness) -4. **aggregate_answers**: Aggregate evaluation results -5. **difficulty_estimation**: Estimate difficulty using small model -6. **dgsdg_post_process**: Clean and create difficulty-stratified datasets +1. **generate_verified_questions**: Generate questions from documents and verify them (4 internal sub-steps: q-prep + Q-gen + verify-prep + Q-verify) +2. **generate_answers**: Filter questions by verification pass-rate, generate N candidate answers (2 internal sub-steps: a-prep + A-gen) +3. **gym_genselect_answers**: Select best answer from multiple candidates +4. **evaluate_answers**: Evaluate answer quality (multi-seed for robustness) +5. **aggregate_answers**: Aggregate evaluation results +6. **dgsdg_post_process**: Clean + rename fields, emit single `final_result.jsonl` consumed by downstream SFT / GRPO **See [technical reference](../stages/document-grounded-sdg.md) for detailed stage documentation.** @@ -92,21 +91,21 @@ uv run nflow run-all --config nvflow/recipes/finance/workflows/sdg/document-grou # Stage 0: Preprocess SEC filings uv run nflow run dg_sdg_preprocess --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml -# Stage 1: Generate verified Q&A -uv run nflow run generate_verified_qa --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml +# Stage 1: Generate + verify questions +uv run nflow run generate_verified_questions --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml -# Stage 2: Select best answers -uv run nflow run genselect_answers --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml +# Stage 2: Generate candidate answers +uv run nflow run generate_answers --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml -# Stage 3: Evaluate answers +# Stage 3: Select best answers +uv run nflow run gym_genselect_answers --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml + +# Stage 4: Evaluate answers uv run nflow run evaluate_answers --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml -# Stage 4: Aggregate results +# Stage 5: Aggregate results uv run nflow run aggregate_answers --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml -# Stage 5: Estimate difficulty -uv run nflow run difficulty_estimation --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml - # Stage 6: Post process uv run nflow run dgsdg_post_process --config nvflow/recipes/finance/workflows/sdg/document-grounded-sdg.yaml ``` @@ -121,24 +120,22 @@ ${base_data_dir}/ β”‚ └── jsonl/ β”‚ β”œβ”€β”€ 10-k-data.jsonl # Sampled 10-K data β”‚ └── 10-q-data.jsonl # Sampled 10-Q data -β”œβ”€β”€ step-1-qa-pipeline/ -β”‚ β”œβ”€β”€ question_pipeline/ -β”‚ β”‚ β”œβ”€β”€ generated/ # Generated questions -β”‚ β”‚ └── verified/ # Verified questions -β”‚ └── answer_pipeline/ -β”‚ └── generated/ # Generated answers -β”œβ”€β”€ step-2-genselect/ +β”œβ”€β”€ step-1-questions/ +β”‚ β”œβ”€β”€ generate_input.jsonl # Q-prep output +β”‚ β”œβ”€β”€ generated/ # Generated questions +β”‚ β”œβ”€β”€ verify_input.jsonl # Q-verify-prep output +β”‚ └── verified/ # Verified questions (consumed by step-2) +β”œβ”€β”€ step-2-answers/ +β”‚ β”œβ”€β”€ answer_input.jsonl # A-prep output (threshold-filtered) +β”‚ └── generated/ # Generated answers (consumed by step-3) +β”œβ”€β”€ step-3-genselect/ β”‚ └── selected_answers.jsonl -β”œβ”€β”€ step-3-evaluate/ -β”‚ └── evaluation results (5 seeds) -β”œβ”€β”€ step-4-aggregate/ +β”œβ”€β”€ step-4-evaluate/ +β”‚ └── evaluation results (multi-seed) +β”œβ”€β”€ step-5-aggregate/ β”‚ └── aggregated_answers.jsonl -β”œβ”€β”€ step-5-difficulty/ -β”‚ └── difficulty scoring results └── step-6-post-process/ - β”œβ”€β”€ full_data.jsonl # All cleaned records - β”œβ”€β”€ final_result.jsonl # Medium difficulty (for SFT) - └── hard_rl_data.jsonl # Hard difficulty training data (difficulty_score=0) + └── final_result.jsonl # Cleaned + renamed records consumed by SFT / GRPO ``` ## Expected Results @@ -150,33 +147,32 @@ ${base_data_dir}/ | Questions Generated | ~2M+ | | Verified Questions | ~1.6M | | Final Q&A Pairs | ~800K | -| Medium Difficulty | ~100K | -| Hard Difficulty | ~400K | | Time | ~30 hours, affected by resources used | ## Output Format ### Final Training Data -**final_result.jsonl** - For supervised fine-tuning: -```json -{ - "question": "Based on the risk factors, what are Tesla's main supply chain concerns?", - "context": "...SEC filing excerpt...", - "generation": "...\n...", - "difficulty_score": 2, - "evaluation_score": 4.5 -} -``` +**final_result.jsonl** - Cleaned, renamed records consumed by downstream SFT / GRPO. Each line contains the per-stage allowlisted generic fields (see `nvflow/generic_stage/sdg/document_grounded/_schemas.py::STAGE_KEEP["dgsdg_post_process"]`) plus the recipe-declared `domain_keep_fields`. It also carries the Responses-API *original form* of the selected answer (`response` + `responses_create_params`) and an `expected_answer` mirroring `answer`, so the record is rollout-like and drop-in for SFT / GRPO. Example for the finance recipe: -**hard_rl_data.jsonl** - Hard difficulty training data: ```json { - "question": "How does NVIDIA's revenue recognition differ for bundled products?", - "context": "...complex accounting excerpt...", - "generation": "...\n...", - "difficulty_score": 0, - "evaluation_score": 4.8 + "context": "...SEC filing excerpt...", + "problem": "Based on the risk factors, what are Tesla's main supply chain concerns?", + "answer": "...", + "reasoning_content": "...", + "question_type": "Risk_Factors", + "answerable": "YES", + "question_voting_pass_rate": 1.0, + "question_voting_total": 5, + "expected_answer": "...", + "responses_create_params": { "...": "exact answer-gen request (Responses-API)" }, + "response": { "...": "original answer-gen response object (Responses-API)" }, + "company_name0": "Tesla, Inc.", + "year": "2023", + "item_section0": "Item 1A", + "file_path0": ".../10-K/...", + "file_type": "10-K" } ``` @@ -187,20 +183,14 @@ ${base_data_dir}/ BASE_DIR="outputs/finance/sap-500/workflow-3-document-grounded-sdg" # Stage outputs -ls $BASE_DIR/step-1-qa-pipeline/answer_pipeline/generated/ -ls $BASE_DIR/step-2-genselect/selected_answers.jsonl -ls $BASE_DIR/step-4-aggregate/aggregated_answers.jsonl +ls $BASE_DIR/step-2-answers/generated/ +ls $BASE_DIR/step-3-genselect/selected_answers.jsonl +ls $BASE_DIR/step-5-aggregate/aggregated_answers.jsonl -# Final datasets +# Final dataset ls $BASE_DIR/step-6-post-process/ - -# Count Q&A by difficulty -echo "Medium difficulty:" wc -l $BASE_DIR/step-6-post-process/final_result.jsonl -echo "Hard difficulty:" -wc -l $BASE_DIR/step-6-post-process/hard_rl_data.jsonl - # Inspect samples head -n 3 $BASE_DIR/step-6-post-process/final_result.jsonl | jq . ``` @@ -221,7 +211,8 @@ Converts raw SEC 10-K and 10-Q HTML filings into structured JSONL data for downs |-----------|-------------|---------| | `input_dir` | Raw SEC filings directory (10-K and 10-Q HTML files) | `${filings_dir}/data` | | `output_dir` | Preprocessed data output directory | `${base_data_dir}/step-0-preprocess` | -| `distribution_dir` | Directory with distribution CSVs (SecQue benchmark) | `/workspace/nvflow/recipes/finance/workflows/sdg/dg_sdg_distribution` | +| `distribution_dir` | Directory with distribution CSVs (SecQue benchmark) | `nvflow/recipes/finance/workflows/sdg/dg_sdg_distribution` | +| `preprocess_module` | Dotted module path to domain CLI that chunks + samples | `nvflow.recipes.finance.utils.sdg.dg_sdg_data_preprocess` | | `max_tokens` | Maximum tokens per chunk | 3000 | | `overlap_tokens` | Overlap tokens between chunks for context coverage | 500 | | `total_samples` | Total samples to generate following distribution | 150000 | @@ -251,18 +242,25 @@ ${filings_dir}/data/ This structure is created automatically by the SEC download workflow ([Workflow 1](01-download-sec.md)). -## Stage 1: generate_verified_qa Details +## Stage 1: generate_verified_questions Details -This stage performs 6 internal sub-steps: +This stage performs 4 internal sub-steps (Q-side of the pipeline): -1. **Preprocess Documents** (CPU): Prepare SEC filings for question generation -2. **Generate Questions** (GPU): Create questions from documents using GPT-OSS-120B -3. **Preprocess Questions** (CPU): Prepare for verification -4. **Verify Questions** (GPU): Verify quality using Qwen3-235B (5 seeds) -5. **Preprocess Verified** (CPU): Filter by threshold, prepare for answers -6. **Generate Answers** (GPU): Create answers using GPT-OSS-120B (5 seeds) +1. **Q-prep** (CPU): Run the recipe-supplied `question_prep_script` to attach `context` strings to each chunk +2. **Q-gen** (GPU): Generate questions from documents using GPT-OSS-120B +3. **Q-verify-prep** (CPU): Expand each generated question into N verification trials +4. **Q-verify** (GPU): Per-question Yes/No vote using Qwen3-235B (5 seeds) -See [technical reference](../stages/document-grounded-sdg.md#generate_verified_qa) for details. +See [technical reference](../stages/document-grounded-sdg.md#generate_verified_questions) for details. + +## Stage 2: generate_answers Details + +This stage performs 2 internal sub-steps (A-side of the pipeline): + +1. **A-prep** (CPU): `construct_answer_generate_input` keeps only questions whose Q-verify pass-rate β‰₯ `answer_preprocess_kwargs.threshold` +2. **A-gen** (GPU): Generate N candidate answers per surviving question using GPT-OSS-120B (5 seeds for downstream genselect) + +See [technical reference](../stages/document-grounded-sdg.md#generate_answers) for details. ## Customization @@ -279,11 +277,17 @@ num_chunks: 10 # Change from 1 β†’ 10 to run 10 jobs in parallel ```yaml stages: - generate_verified_qa: + generate_verified_questions: question_generation_kwargs: args: model: /path/to/your/model - server_gpus: 8 + num_gpus: 8 + + generate_answers: + answer_generation_kwargs: + args: + model: /path/to/your/model + num_gpus: 8 ``` ### Modify Prompts @@ -291,9 +295,9 @@ stages: Edit prompts in `nvflow/recipes/finance/prompts/`: - `document_grounded_generate_questions.yaml` - Question generation - `document_grounded_verify_questions.yaml` - Question verification -- `generate_answers.yaml` - Answer generation +- `secque_template.yaml` - Answer generation +- `genselect_answers.yaml` - GenSelect (best-of-N answer picker) - `evaluate_answers.yaml` - Answer evaluation -- `judge_difficulty.yaml` - Difficulty judging ## Common Issues @@ -313,13 +317,6 @@ ls outputs/finance/sap-500/workflow-2-download-sec/step-0-download/data/ - Lower threshold to 0.6 (3 out of 5 seeds) - Review question generation prompt -### Difficulty estimation takes too long - -**Solution:** -- Reduce `num_random_seeds` for answer generation -- Use fewer `num_chunks` for parallelization -- Use smaller judge model - ## Combining with Template-Based You can combine both SDG approaches: @@ -338,7 +335,7 @@ cat outputs/finance/sap-500/workflow-3-template-based-sdg/step-5-filter-answers/ After completing document-grounded SDG: -- **[SFT Training](04-sft.md)** - Train on stratified datasets +- **[SFT Training](04-sft.md)** - Train on `final_result.jsonl` - **[Evaluation](05-eval.md)** - Test model performance - Combine with template-based data for more diversity @@ -353,6 +350,5 @@ For comprehensive stage-by-stage documentation: |-------|-------|------| | GPT-OSS-120B | Question generation, answer generation | 120B | | Qwen3-235B-A22B | Question verification, answer selection, evaluation | 235B | -| Qwen3-4B | Difficulty estimation (small model baseline) | 4B | All models are configurable in the workflow YAML. diff --git a/docs/recipes/finance/workflows/04-sft.md b/docs/recipes/finance/workflows/04-sft.md index 3aa39c4..e871e8d 100644 --- a/docs/recipes/finance/workflows/04-sft.md +++ b/docs/recipes/finance/workflows/04-sft.md @@ -56,19 +56,19 @@ Fine-tune language models on synthetic financial Q&A data generated from SDG wor β”‚ β–Ό β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” -β”‚ 5. convert_to_messages β”‚ Conversion: Convert to OpenAI messages format +β”‚ 5. eval β”‚ Evaluation: Score checkpoints on finance benchmarks β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ ``` -> **Note:** All 6 stages run in the production `qwen3_14b.yaml` configuration. Some stages (`sequence_length_grouping`, `convert_to_messages`) may be optional for custom configurations. - **6 Stages:** 1. **data_transformation** (Step 0): Convert Q&A format to training format 2. **prepare_for_sft** (Step 1): Prepare data for SFT (formatting, filtering) 3. **train_validation_split** (Step 2): Split into train/validation sets 4. **sequence_length_grouping** (Step 3): Group by sequence length for efficiency 5. **training** (Step 4): Fine-tune the model -6. **convert_to_messages** (Step 5): Convert to message format for chat interfaces +6. **eval** (Step 5): Evaluate checkpoints on finance benchmarks + +> **Qwen3 models add a seventh stage.** `qwen3_14b.yaml` inserts `convert_to_messages` between `training` and `eval` to convert checkpoints to the OpenAI messages format. Other configs, including the `qwen3_4b.yaml` demo, run the six stages above. **See [technical reference](../stages/sft.md) for detailed stage documentation.** diff --git a/docs/recipes/finance/workflows/05-eval.md b/docs/recipes/finance/workflows/05-eval.md index 8584e2f..fc7c78c 100644 --- a/docs/recipes/finance/workflows/05-eval.md +++ b/docs/recipes/finance/workflows/05-eval.md @@ -36,6 +36,8 @@ Standalone Baselines └───────────── β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ ``` +> **Offline clusters:** `prepare_data` downloads the benchmark datasets from HuggingFace, so temporarily clear `HF_HUB_OFFLINE`, `HF_DATASETS_OFFLINE` and `TRANSFORMERS_OFFLINE` for its first run, then restore them. The datasets persist and are reused afterwards. See [Offline runtime](../troubleshooting.md#offline-runtime). + ## Configuration **Directory:** `workflows/eval/` @@ -51,7 +53,8 @@ Checkpoint evaluation is configured directly in the training configs: |------|-------------| | `sft/qwen3_4b.yaml` | `stages.eval` with `eval_steps: [10]` | | `sft/qwen3_14b.yaml` | `stages.eval` with `eval_steps: [2600, 5000, 7408]` | -| `grpo/qwen3_4b.yaml` | `stages.eval` with `eval_steps: [20]` | +| `grpo/qwen3_4b.yaml` (equivalence, FSDP) | `stages.eval` with `eval_steps: [20]` | +| `grpo/qwen3_4b_finsec.yaml` (finance_sec_search, Megatron) | `stages.eval` with `eval_steps: [20]` | ## Usage @@ -94,7 +97,8 @@ Configured in `eval/base.yaml`, shared across all evaluation contexts: - **SEC-QUE**: SEC filing comprehension (565 samples) - **FinanceBench**: Financial question answering (150 samples) -- **finance_agent**: Multi-turn agentic financial QA from [vals-ai/finance-agent](https://github.com/vals-ai/finance-agent) (50 samples) + +`finance_agent` (multi-turn agentic financial QA from [vals-ai/finance-agent](https://github.com/vals-ai/finance-agent)) is **disabled** β€” `eval/base.yaml` sets it to `null` pending validation of the multi-turn tool-calling path. See [finance-agent-eval](../stages/finance-agent-eval.md) to re-enable it. ## Eval Stage Configuration (in Training YAMLs) @@ -105,12 +109,12 @@ stages: eval: eval_steps: [1000, 3000, 5000] checkpoint_path: ${directories.step-4-training}/model-name - format: megatron # Use "fsdp" for GRPO demo, "megatron" for GRPO production + format: megatron # Match the checkpoint's training backend: "fsdp" or "megatron" baseline_model: /hf_models/Qwen/Qwen3-14B server_type: vllm gpus: 1 inference_args: >- - ++prompt_config=/workspace/nvflow/recipes/finance/prompts/secque_template.yaml + ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml ++inference.temperature=0.6 ++inference.top_p=0.95 ++inference.top_k=20 @@ -170,12 +174,12 @@ stages: eval: eval_steps: [100, 500, 1000] checkpoint_path: ${directories.step-4-training}/model-my-model-name - format: megatron # Use "fsdp" for GRPO demo checkpoints + format: megatron # Match the checkpoint's training backend: "fsdp" or "megatron" baseline_model: /hf_models/MyOrg/MyModel server_type: vllm gpus: 1 inference_args: >- - ++prompt_config=/workspace/nvflow/recipes/finance/prompts/secque_template.yaml + ++prompt_config=nvflow/recipes/finance/prompts/secque_template.yaml server_args: "--max-model-len 40960" ``` diff --git a/docs/recipes/finance/workflows/06-finance-agent-eval.md b/docs/recipes/finance/workflows/06-finance-agent-eval.md deleted file mode 100644 index 2d975c7..0000000 --- a/docs/recipes/finance/workflows/06-finance-agent-eval.md +++ /dev/null @@ -1,64 +0,0 @@ -# Finance Agent Benchmark - -## Overview - -The **finance_agent** benchmark ([vals-ai/finance-agent](https://github.com/vals-ai/finance-agent)) is now integrated into the main [evaluation workflow](05-eval.md). It is defined as a benchmark entry in `workflows/eval/base.yaml` alongside SEC-QUE and FinanceBench. - -> **Note:** The standalone `finance_agent_eval.yaml` workflow has been removed. All finance_agent evaluation now runs through the unified eval configs in `workflows/eval/`. - -## What is finance_agent? - -- **50 public questions** from vals-ai/finance-agent -- **Multi-turn**: Model can take up to 50 turns (tool calls + reasoning) -- **Tools**: Web search (Tavily), SEC EDGAR lookup, HTML parsing -- **Judge**: GPT-5 mini with strict finance-domain prompts (`sec_judge_strict.yaml`) - -## Configuration - -The finance_agent benchmark is configured in `workflows/eval/base.yaml` under the `benchmarks` section: - -```yaml -benchmarks: - finance_agent: - seeds: 5 - judge: *judge_finance_strict - installation_command: "pip install -q ..." - extra_args: >- - ++max_turns=50 - ++inference.tokens_to_generate=32000 - ++inference.temperature=0.0 - ++max_concurrent_requests=1 -``` - -Any model YAML that inherits from `base.yaml` will automatically include finance_agent in its evaluation benchmarks. - -## Usage - -Run finance_agent evaluation as part of any eval context: - -```bash -# Evaluate baselines on all benchmarks (including finance_agent) -uv run nflow run-all --config nvflow/recipes/finance/workflows/eval/baselines.yaml - -# SFT training + checkpoint eval (includes finance_agent) -uv run nflow run-all --config nvflow/recipes/finance/workflows/sft/qwen3_14b.yaml -``` - -## Output Structure - -Outputs appear under the model's eval-results directory: - -``` -outputs/finance/sap-500/workflow-1-baseline-eval/ -└── baselines/ - └── gpt-oss-120b/ - └── eval-results/ - └── finance_agent/ - β”œβ”€β”€ metrics.json # Aggregated metrics - └── output*.jsonl # Predictions per seed -``` - -## Related - -- **[Evaluation Workflow (05-eval)](05-eval.md)** – Full eval documentation, including all benchmarks -- **[Eval Stages Reference](../stages/eval.md)** – Technical stage documentation diff --git a/docs/recipes/finance/workflows/06-grpo.md b/docs/recipes/finance/workflows/06-grpo.md index 7e263d9..a04d2a0 100644 --- a/docs/recipes/finance/workflows/06-grpo.md +++ b/docs/recipes/finance/workflows/06-grpo.md @@ -96,10 +96,11 @@ Further improve fine-tuned models using Group Relative Policy Optimization (GRPO ### Model Configurations -| Config | Model | GPUs | Status | -|--------|-------|------|--------| -| `grpo/qwen3_4b.yaml` | Qwen3-4B | 16 (2 nodes) | Demo | -| `grpo/qwen3_30b_a3b.yaml` | Qwen3-30B-A3B (MoE) | 64 (8 nodes) | Production | +| Config | Model | Environment | Backend | GPUs | Status | +|--------|-------|-------------|---------|------|--------| +| `grpo/qwen3_4b.yaml` | Qwen3-4B | equivalence_llm_judge | FSDP v2 (32K) | 16 (2 nodes) | Demo | +| `grpo/qwen3_4b_finsec.yaml` | Qwen3-4B | finance_sec_search | Megatron (TP2Γ—CP8, 131K) | 16 (2 nodes) | Demo | +| `grpo/qwen3_30b_a3b.yaml` | Qwen3-30B-A3B (MoE) | β€” | Megatron | 64 (8 nodes) | Production | ## Usage @@ -220,9 +221,12 @@ outputs/finance/demo/workflow-5-grpo/ β”‚ β”œβ”€β”€ val.jsonl # Validation split β”‚ └── logs/ β”œβ”€β”€ step-8-training/ - β”‚ └── grpo-qwen3-4b-2n-tp2-cp4-seq131k/ # Demo (FSDP v2) - β”‚ β”œβ”€β”€ checkpoints/ # GRPO model checkpoints - β”‚ └── training-logs/ + β”‚ β”œβ”€β”€ equivalence_llm_judge/ + β”‚ β”‚ └── grpo-qwen3-4b-16g-tp2-cp1-seq32k/ # Demo, FSDP v2 + β”‚ └── finance_sec_search/ + β”‚ └── grpo-qwen3-4b-16g-tp2-cp8-seq128k/ # Demo, Megatron (YaRN 131K) + β”‚ β”œβ”€β”€ checkpoints/ # GRPO model checkpoints + β”‚ └── training-logs/ └── step-9-eval/ └── ... # Benchmark evaluation results ``` @@ -341,7 +345,7 @@ stages: ### Training Backends -The demo config (`qwen3_4b.yaml`) uses **FSDP v2** for the dense Qwen3-4B model. The production config (`qwen3_30b_a3b.yaml`) uses **Megatron** for the Qwen3-30B-A3B MoE model at 64 GPUs. +The demo runs two environments with different backends: `qwen3_4b.yaml` (equivalence_llm_judge) uses **FSDP v2** at 32K, while `qwen3_4b_finsec.yaml` (finance_sec_search) uses **Megatron** (TP2Γ—CP8) for YaRN context extension to 131K. The production config (`qwen3_30b_a3b.yaml`) uses **Megatron** for the Qwen3-30B-A3B MoE model at 64 GPUs. **Production (Megatron):** diff --git a/docs/recipes/multimodal/README.md b/docs/recipes/multimodal/README.md new file mode 100644 index 0000000..2e279f7 --- /dev/null +++ b/docs/recipes/multimodal/README.md @@ -0,0 +1,112 @@ +# Multimodal HopChain Recipe + +The multimodal recipe implements a HopChain-inspired synthetic data generation +pipeline for multi-hop vision-language reasoning. It follows the paper +[HopChain: Multi-Hop Data Synthesis for Generalizable Vision-Language Reasoning](https://arxiv.org/pdf/2603.17024) +and expresses the workflow as reusable NVFlow stages. + +Start with the [HopChain quick start](quick-start.md). + +## Workflows + +| Workflow | Demo config | Full config | +| --- | --- | --- | +| Image filter | `nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml` | `nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter.yaml` | +| SDG | `nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml` | `nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg.yaml` | + +The demo SDG config stops after verified-question visualization. It has no +external API dependency. The full config additionally runs the OpenAI +judge, reconciliation, Omni difficulty filtering, and SFT trace generation. + +## Configuration Contract + +Configuration is split between workflow and cluster files: + +- Each full workflow YAML defines its stages, model profiles, execution + IDs, chunking, and repository-relative input/output paths. +- Each demo YAML inherits its corresponding full workflow and overrides + only the stage selection and small-run settings. +- [`cluster_configs/my_cluster.yaml`](../../cluster-configuration.md) defines + the local Slurm account, partitions, mounts, and named container image paths. +- Optional private recipe changes go in git-ignored `private_*.yaml` overlays + next to the workflow they modify. +- The full workflow's OpenAI key is supplied as `OPENAI_API_KEY` under `env_vars` + in `cluster_configs/my_cluster.yaml`. + +Run the demo workflows in order: + +```bash +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml + +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml +``` + +Outputs are deterministic: + +```text +outputs/hopchain/image_filter/execution/demo/ +outputs/hopchain/sdg/execution/demo/ +``` + +Run the full workflows with their full configs: + +```bash +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter.yaml + +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg.yaml +``` + +Full-workflow outputs use these directories: + +```text +outputs/hopchain/image_filter/execution/full/ +outputs/hopchain/sdg/execution/full/ +``` + +## SDG Stages + +The full `hopchain_sdg` workflow runs: + +1. `prepare_filtered_image_inputs` +2. `preprocess_identify_categories` +3. `identify_categories` +4. `localize_instances` +5. `sample_instance_combinations` +6. `preprocess_generate_multihop_queries` +7. `generate_multihop_queries` +8. `verify_candidate_queries` +9. `visualize_candidate_hopchain_data` +10. `judge_candidate_queries_openai` +11. `reconcile_llm_judges` +12. `visualize_reconciled_hopchain_data` +13. `preprocess_filter_easy_candidates` +14. `filter_easy_candidates` +15. `preprocess_generate_sft_reasoning_traces` +16. `generate_sft_reasoning_traces` +17. `preprocess_filter_sft_reasoning_traces` +18. `filter_sft_reasoning_traces` + +## Inputs and Models + +The image filter recursively scans `data/images/` and writes +`outputs/hopchain/image_filter/execution/demo/image-filter/kept_images.jsonl`. +The SDG demo reads that file as its input. + +By default the containers must see checkpoints at: + +```text +/hf_models/Qwen/Qwen3.5-397B-A17B +/hf_models/facebook/sam3.1/sam3.1_multiplex.pt +/hf_models/nvidia/omni-step70 +``` + +Set host-to-container mappings in `my_cluster.yaml` and server behavior in an +ignored local YAML overlay. Keep machine-specific paths in those local files. + +Do not commit API keys or credential files. See the +[quick start](quick-start.md#run-the-full-workflow) for full-workflow credential setup +and the data-egress warning. diff --git a/docs/recipes/multimodal/quick-start.md b/docs/recipes/multimodal/quick-start.md new file mode 100644 index 0000000..b17399b --- /dev/null +++ b/docs/recipes/multimodal/quick-start.md @@ -0,0 +1,326 @@ +# HopChain Quick Start + +Run HopChain from a folder of images to verified multi-hop vision-language +questions. + +The demo has two commands: + +1. Filter the source images with Qwen. +2. Generate and verify multi-hop questions with Qwen and SAM 3.1. + +Plan for 30–60 minutes for a small demo run, plus Slurm queue time. Image +filtering typically takes 5–10 minutes, and the SDG dependency chain takes 20 +minutes or more. Runtime increases with the number of images and generated +queries. + +The demo ends after verified-question visualization. The full workflow +adds the OpenAI judge and Omni curation stages. + +## Pipeline Overview + +```text + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ 1. Your images │────▢│ 2. Image Filter │────▢│ 3. SDG β”‚ + β”‚ (a folder) β”‚ β”‚ keep complex β”‚ β”‚ categories β†’ localize β†’ β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ combine β†’ generate & β”‚ + β”‚ β”‚ verify multi-hop queries β”‚ + kept_images.jsonl β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β–Ό β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ 4. Judge + reconcile β”‚ β”‚ 5. Difficulty filter β”‚ + β”‚ (OpenAI API) β”‚ β”‚ + SFT reasoning tracesβ”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +The full path continues from verified questions through the external +judge, judge reconciliation, difficulty filtering, and SFT reasoning-trace +generation. Image filtering is a separate workflow so its output can be reused +by several SDG runs. + +## Prerequisites + +- Cluster access configured as described in [INSTALL.md](../../../INSTALL.md), + with all commands run from the `nvflow` repository root. +- A Slurm cluster config created from `cluster_configs/template-slurm.yaml`; see + [Step 1](#1-configure-models-and-cluster). +- GPUs for the model workers. The **core** path (Steps 3–4) needs: + - A **VLM server** for image scoring and question generation. The reference + config serves `Qwen/Qwen3.5-397B-A17B` with SGLang + - A **SAM 3.1 worker** for object localization. +- The **full** path additionally needs an OpenAI API key for the LLM judge and + the Omni reasoning VLM for difficulty filtering. + +> **Heads up:** The reference models are large. For a quick try, use +> [smaller models you can serve](#local-overrides). + +## 1. Configure Models and Cluster + +### SAM 3.1 checkpoint + +Request access to +[Meta's gated SAM 3.1 repository](https://huggingface.co/facebook/sam3.1), then +download the checkpoint once from a connected host: + +```bash +uv run hf auth login +uv run hf download facebook/sam3.1 sam3.1_multiplex.pt \ + --local-dir /path/to/models/hf_models/facebook/sam3.1 +``` + +After the checkpoint is downloaded, the compute jobs do not need `HF_TOKEN`. + +### Cluster configuration + +Follow [Configure Your Cluster](../../../INSTALL.md#configure-your-cluster) to +create `cluster_configs/my_cluster.yaml`. The +[Cluster Configuration Guide](../../cluster-configuration.md) documents every +available field. + +In `my_cluster.yaml`, configure the named `nemo-skills`, `sglang`, and `vllm` +container entries. Mount the checkout and your host model directory so they +are visible on every compute node. The reference configs expect these paths +inside the containers: + +```text +/hf_models/Qwen/Qwen3.5-397B-A17B +/hf_models/facebook/sam3.1/sam3.1_multiplex.pt +/hf_models/nvidia/omni-step70 # full workflow only +``` + +The workflow resolves the repository root from the shell's standard `PWD`. +Make the checkout visible to Slurm jobs at the same absolute path. On sites +that mount a workspace at `/workspace`, launch NVFlow from the checkout under +that mount, such as `/workspace/nvflow`. + +## 2. Add Images + +Copy or mount images anywhere below: + +```bash +mkdir -p data/images +# Copy or mount images below data/images/. +``` + +Subdirectories are scanned recursively. The demo selects at most 100 images and +the SDG step uses at most 25 images that pass filtering. Prefer visually rich +scenes, documents, charts, or infographics with several distinct regions. + +Public datasets that fit the recipe well include: + +| Dataset | Why it fits HopChain | Source | +| --- | --- | --- | +| COCO 2017 validation | Everyday multi-object scenes; a practical first run | | +| Visual Genome | Dense objects and relationships | | +| InfographicVQA, DocVQA, or ChartQA | Text- and figure-rich images for OCR reasoning | Hugging Face Datasets | +| ADE20K | Complex scene-parsing images | | +| Open Images V7 | Large and diverse multi-object collection | | + +Images remain path-referenced throughout the pipeline, so keep the directory +mounted and unchanged until the run completes. + +## 3. Filter Images (~5–10 minutes) + +Validate, preview, and submit the demo: + +```bash +uv run nflow validate \ + --config nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml + +uv run nflow list-stages \ + --config nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml + +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml +``` + +`run-all` submits Slurm work and returns. After the job finishes, inspect the +deterministic demo output: + +```bash +python -m json.tool outputs/hopchain/image_filter/execution/demo/image-filter/summary.json +wc -l outputs/hopchain/image_filter/execution/demo/image-filter/kept_images.jsonl +``` + +The second command must report at least one kept image before SDG can proceed. + +The image-filter output contains: + +```text +image-filter/ +β”œβ”€β”€ image_catalog.jsonl +β”œβ”€β”€ output.jsonl +β”œβ”€β”€ final_output.jsonl +β”œβ”€β”€ kept_images.jsonl +└── summary.json +``` + +`final_output.jsonl` includes every scored image; `kept_images.jsonl` contains +only images that passed the configured quality and complexity thresholds. + +## 4. Generate Multi-Hop Questions (~20+ minutes) + +The SDG demo reads +`outputs/hopchain/image_filter/execution/demo/image-filter/kept_images.jsonl`. + +```bash +uv run nflow validate \ + --config nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml + +uv run nflow list-stages \ + --config nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml + +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml +``` + +The demo runs the local core path: + +```text +prepare images -> identify categories -> localize with SAM -> sample object +combinations -> generate questions -> verify questions -> build visualization +``` + +After the dependency chain completes: + +```bash +python -m json.tool \ + outputs/hopchain/sdg/execution/demo/step-5-verify-candidate-queries/summary.json + +wc -l \ + outputs/hopchain/sdg/execution/demo/step-5-verify-candidate-queries/final_candidates.jsonl +``` + +Review the generated HTML under +`outputs/hopchain/sdg/execution/demo/step-6-visualize-candidate-hopchain-data/`. + +The core output layout is: + +```text +sdg/execution/demo/ +β”œβ”€β”€ step-0-prepare-filtered-inputs/filtered_image_inputs.jsonl +β”œβ”€β”€ step-1-identify-categories/final_output.jsonl +β”œβ”€β”€ step-2-localize-instances/ +β”œβ”€β”€ step-3-sample-instance-combinations/instance_combinations.jsonl +β”œβ”€β”€ step-4-generate-multihop-queries/final_output.jsonl +β”œβ”€β”€ step-5-verify-candidate-queries/ +β”‚ β”œβ”€β”€ final_candidates.jsonl +β”‚ β”œβ”€β”€ rejected_candidates.jsonl +β”‚ └── summary.json +└── step-6-visualize-candidate-hopchain-data/ +``` + +## Run the Full Workflow + +The full configs use the `full` execution ID. They process the complete input +set, use full-run chunk counts, call the +OpenAI judge, run the Omni +difficulty filter, and create SFT reasoning traces. + +Before running the full workflow, add your OpenAI key to +`cluster_configs/my_cluster.yaml`, following the existing +[environment-variable instructions](../../cluster-configuration.md#environment-variables): + +```yaml +env_vars: + # ...existing cluster environment variables... + - OPENAI_API_KEY= +``` + +The OpenAI judge sends question and image content to an external service. Only +enable the full path when that data transfer is allowed. + +Then run: + +```bash +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter.yaml + +uv run nflow run-all \ + --config nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg.yaml +``` + +Full-workflow outputs live under: + +```text +outputs/hopchain/image_filter/execution/full/ +outputs/hopchain/sdg/execution/full/ +``` + +The full-workflow stage groups are: + +| Steps | Work | Needs | +| --- | --- | --- | +| 0–6 | Prepare, identify, localize, combine, generate, verify, visualize | Qwen and SAM | +| 7–9 | OpenAI judge, reconcile, and visualize reconciled data | `OPENAI_API_KEY` | +| 10 | Filter easy candidates | Omni reasoning VLM | +| 11–12 | Generate and filter SFT reasoning traces | Qwen | + +Adjust full-run chunk counts after checking the image-filter and combination +counts for your dataset. + +## Local Overrides + +Put deployment-specific recipe changes in a small `private_*.yaml` overlay next +to the workflow it modifies (`private_*.yaml` files are git-ignored repo-wide). +For example: + +```yaml +# nvflow/recipes/multimodal/workflows/image_filter/private_hopchain-image-filter.yaml +_base_: hopchain-image-filter-demo.yaml + +execution_id: my_test +model_profiles: + qwen: + server_gpus: 4 + server_nodes: 1 + server_args: >- + --model-path /hf_models/Qwen/Qwen3.5-397B-A17B + --served-model-name qwen3.5-397b-a17b + --tp 4 + --trust-remote-code +``` + +Use another small overlay based on `hopchain-sdg-demo.yaml` (in +`workflows/sdg/`) when the SDG model profile also needs to change. Keep host +paths, Slurm partitions, mounts, and container image paths in +`cluster_configs/my_cluster.yaml`. + +## Next Steps + +- Review `final_candidates.jsonl` and the candidate HTML before enabling the + external judge. +- Tune `min_complexity_score` or `allowed_quality_ratings` in a local + image-filter overlay when the kept set is too broad or too small. +- Use a local SDG overlay to calibrate `sample_count`, query count, and chunk + counts before a full run. +- Read the [multimodal HopChain guide](README.md) for the complete stage list + and configuration contract. + +## Troubleshooting + +### The config validates, but the job cannot see files + +`validate` runs in the launch shell; the stage itself runs in a container on a +compute node. Confirm that the checkout, images, outputs, and checkpoint paths +are covered by `my_cluster.yaml` mounts and appear at the paths documented +above. + +### No images were selected + +Confirm `data/images/` contains supported image files. If filtering ran +but kept zero images, inspect `final_output.jsonl` and lower +`min_complexity_score` in a local image-filter overlay. + +### A job requests the wrong partition or container + +Partitions and container image paths come from `cluster_configs/my_cluster.yaml`. +Check `partition`, `cpu_partition`, and the named container entries there. + +### The full workflow fails at the judge stage + +Confirm `OPENAI_API_KEY` is present under `env_vars` in the ignored +`cluster_configs/my_cluster.yaml`. The cluster config injects it into the +`nemo-skills` container used by the full-workflow judge. + +[Multimodal HopChain Guide](README.md) | [Main README](../../../README.md) diff --git a/docs/remote-launch.md b/docs/remote-launch.md new file mode 100644 index 0000000..b1712db --- /dev/null +++ b/docs/remote-launch.md @@ -0,0 +1,152 @@ +# Running `nflow` over an SSH tunnel + +`nflow` is only a **submission orchestrator**: it builds Slurm jobs and submits +them β€” all data, GPU work, and training run in the worker containers **on the +cluster**. When you run `nflow` somewhere that can't reach Slurm directly (a +laptop, a dev box, or an isolated/airgapped environment), it submits over an +**SSH tunnel**. + +> **On a cluster login/dev node?** You don't need this doc β€” install per the +> [README](../README.md#-installation) and run `nflow` directly. This page is for +> the **off-cluster / tunneled** case. For all client options at a glance, see +> [INSTALL.md β†’ Choose your client setup](../INSTALL.md#choose-your-client-setup). + +## Options at a glance + +```text +Where does `nflow` run? +β”‚ +β”œβ”€ On a cluster login/dev node ────────────▢ sbatch ─▢ Slurm worker jobs (no tunnel) +β”‚ install: uv sync +β”‚ +└─ Off-cluster (laptop / dev box / airgap) ──ssh_tunnel──▢ login node ─sbatch─▢ workers + provision the launcher, pick one: + A. host install β€” uv sync (client host needs internet) + B. nvflow-client image β€” no uv sync, no client internet + β”œβ”€ enroot (cluster node) + β”œβ”€ docker/podman (off-cluster machine) + └─ pyxis srun (cluster node, via Slurm) + +Worker jobs (nemo-skills Β· vllm Β· vllm-grpo Β· nemo-rl Β· nemo-gym Β· sglang) +always run on the cluster; the client only submits. +``` + +## Prerequisites + +- **Cluster side is set up** ([INSTALL.md](../INSTALL.md)): worker `.sqsh` images + and models are staged, and you have a `my_cluster.yaml`. +- **SSH key auth** to a cluster login node that can run `sbatch`: + ```bash + ssh -i @ 'hostname && command -v sbatch' + ``` + +## Step 1 β€” Configure `my_cluster.yaml` (add the tunnel) + +Put `my_cluster.yaml` where the launcher reads it β€” **container:** the mounted +`/work` dir (`NEMO_SKILLS_CONFIG_DIR=/work`); **host install:** `cluster_configs/`. +Add an `ssh_tunnel` block so the launcher reaches Slurm over SSH (no Slurm client +or Lustre needed on the client): + +```yaml +ssh_tunnel: + host: + user: + identity: # container: /opt/ssh/ (id_rsa / id_ed25519) + job_dir: +``` + +> `/work` is a **bind mount** β€” prepare `my_cluster.yaml` before starting the +> container, or edit it live afterward; it just must be complete before +> `nflow run`. It holds secrets: keep it in `/work`, never bake it into an image. + +The rest of `my_cluster.yaml` is your standard cluster config (containers, +`mounts:`, `env_vars`); `ssh_tunnel` is the only tunnel-specific addition. In +`mounts:`, keep `/hf_models` and point `/workspace` at a **writable data dir** +(outputs + HF cache) β€” **not** the repo checkout. Recipe code and checked-in +assets reach workers via the packaged snapshot (`/nemo_run/code`), so the repo is +never mounted. See [cluster-configuration.md β†’ Mounts](cluster-configuration.md#mounts). + +## Step 2 β€” Start the launcher (pick one) + +### A. Host install (`uv sync`) β€” client host has internet + +Follow the [README install](../README.md#-installation) (`git clone` + `uv sync`). +Invoke the CLI as **`uv run nflow …`**. No client internet? Use the +`nvflow-client` image (option B below) instead. + +### B. Client container β€” airgapped / no local install (invoke as **`nflow …`**) + +The `nvflow-client` image bundles the `nflow` CLI + venv (no `uv sync`, no client +internet). Start it, mounting your **SSH key** (`β†’ /opt/ssh`) and the **`/work`** +dir holding `my_cluster.yaml`: + +```bash +# --- Cluster node (enroot) β€” if the .sqsh is already staged, skip the import --- +enroot import -o nvflow-client.sqsh 'docker://#/nvflow-client:' # only from a registry ref +enroot create --name nvflow-client /path/to/nvflow-client.sqsh +ENROOT_MOUNT_HOME=n enroot start --rw \ + -m ~/.ssh:/opt/ssh -m /path/to/work:/work \ + -e NEMO_SKILLS_CONFIG_DIR=/work nvflow-client bash + +# --- Cluster node via Slurm (pyxis/srun) β€” starts from the .sqsh directly --- +srun --container-image=/path/to/nvflow-client.sqsh \ + --container-mounts=/path/to/work:/work,$HOME/.ssh:/opt/ssh \ + --container-workdir=/opt/nvflow \ + --export=ALL,NEMO_SKILLS_CONFIG_DIR=/work --pty bash + +# --- Off-cluster machine (docker/podman) --- +docker run --rm -it -v ~/.ssh:/opt/ssh:ro -v /path/to/work:/work \ + -e NEMO_SKILLS_CONFIG_DIR=/work /nvflow-client: bash +``` + +> Prefer **enroot** (cluster) or **docker/podman** (off-cluster); the `srun` form +> burns an allocation just to host the launcher. Do **not** bind-mount over +> `/opt/nvflow` (baked source/venv/`.git` that nemo-run packages via `git archive`). +> Host keys auto-accept on first connect (baked `ssh_config` reads +> `/opt/ssh/known_hosts`; a *changed* key is still rejected). Build details: +> [containers.md](maintainers/containers.md). + +## Step 3 β€” Launch and monitor over the tunnel + +```bash +nflow list-stages --recipe finance # verify: CLI loads + config resolves +nflow run -c -e # submit (detaches when queued) +``` + +The client has **no Slurm client or cluster filesystem**, so monitor on the +cluster over the same SSH: + +```bash +ssh -i @ 'squeue --me' # or: sacct -j +ssh -i @ 'ls /...' # logs/artifacts land on Lustre +``` + +`nemo experiment status ` (printed at submit) also works over the tunnel. + +## Notes + +- **Connected-node prerequisites** (benchmark datasets, SEC filings, model + downloads) need internet and the `HF_*_OFFLINE` flags **off** for that one run β€” + do them once per [INSTALL.md](../INSTALL.md), then keep the flags **on**. The + container can stage models itself: + `uv run hf download --local-dir /hf_models/` (mount the models dir). +- **Everything runs on the cluster; the client only submits.** GPU work, data + I/O, and the rollout/judge servers all execute inside Slurm jobs. Recipe code + and checked-in assets ship with each job via `/nemo_run/code` (see Step 1), so + the client needs no repo and the repo is never mounted on workers. +- **Laptop / off-cluster specifics** (validated: a client with **no repo mount** + ran the full matrix end-to-end β€” staging β†’ SDG β†’ SFT β†’ eval and **both GRPO + workflows** (`finance_sec_search` via the client, equivalence via a repo + install) β€” proving all I/O is cluster-side and checked-in assets resolve from + `/nemo_run/code`, incl. Gym `config_paths`, prefetch `ticker`, and judge + fpaths): + - `ssh_tunnel.host` must be an **FQDN reachable from the laptop** (VPN), and + `ssh_tunnel.identity` your **local** key (e.g. `~/.ssh/id_rsa`). + - `mounts:` and `job_dir` are **cluster Lustre paths**; the laptop needs none of + them locally. Resume/chunk-skip is probed over the tunnel (`LauncherFS`), so + **no local mount is required** β€” and while `ssh_tunnel` is set a local mount + is ignored anyway. (A client running **on-cluster without** `ssh_tunnel` must + run from the repo root so `resolve_host_path` can map `/workspace/outputs/...` + back to the host outputs dir for skip-detection.) + - Dev-mode source overlays (Gym / NeMo-RL) must live **on the cluster**, not the + laptop β€” they bind into the worker jobs. diff --git a/docs/trace-viewer.md b/docs/trace-viewer.md new file mode 100644 index 0000000..67bee64 --- /dev/null +++ b/docs/trace-viewer.md @@ -0,0 +1,61 @@ +# Rollout Trace Viewer + +A lightweight, dependency-free web UI to spot-check NeMo-Gym rollout traces one +record at a time. Implemented in [`scripts/view_traces.py`](../scripts/view_traces.py) +(pure Python stdlib -- no Gradio, no extra installs). + +It reads only the requested record (seek-by-line with a lazy byte-offset cache), +so it opens record 0 or record 35,000 of a multi-GB `output-rs*.jsonl` without +loading the file. + +## Run + +```bash +cd nvflow +uv run python scripts/view_traces.py [--root ] [--port 8800] +``` + +- `--root` (default: `$NVFLOW_TRACE_ROOT` if set, else the current directory) -- + directory scanned for `*.jsonl` files (the file dropdown). Heavy/non-trace dirs + (`cache/`, `logs/`, `.venv/`, ...) and input artifacts + (`*materialized_inputs*`, `*chunk_input*`) are skipped automatically. + Point it at a single workflow output dir. Do **not** point it at a parent that + also holds the SEC filing dump -- scanning tens of thousands of filings makes + the directory listing crawl. +- `--port` (default 8800), `--host` (default `127.0.0.1`). + +## View it in the browser + +The server binds `127.0.0.1`, so reach it through the SSH tunnel: + +- In **Cursor / VS Code Remote**: the port is auto-forwarded. Open the **Ports** + panel, find the port, click the globe ("Open in Browser"). If it isn't listed, + "Forward a Port" -> enter the port. (Start the server in Cursor's integrated + terminal so auto-forward triggers.) +- Manual fallback from your laptop: `ssh -L 8800:localhost:8800 ` then open + `http://localhost:8800`. + +## Using it + +- **File dropdown**: pick a rollout file. For traces choose + `…/rollout/output-rs*.jsonl` or the curated `…/rollout/analysis_rs*/{best,worst,intermediate}.jsonl`. + A `train.jsonl` has no trace (just question + difficulty) and renders as a + collapsible JSON record. +- **Navigate one record at a time**: record-number box + **Go**, **Prev/Next**, + **Random** (Random counts the file once, then is instant). +- **Trace rendering**: the exact recorded order of `input` + `response.output` -- + each step color-coded with an icon/pill (user, reasoning, tool call, tool + output, assistant), collapsed by default with a one-line preview. Click a step + to expand; **Expand all / Collapse all** at the top right. +- **JSON as a tree**: tool-call args, tool outputs, and the **Raw JSON** view + render as a colorized, collapsible tree -- click any `{}`/`[]` to fold/unfold + nested fields. +- **Verdict header**: reward badge, judge rating/text, expected answer, + question type, uuid. + +## Notes + +- Stdlib only; runs under `uv run python` (3.12) or any `python3` (3.9+). +- Responses use `Cache-Control: no-store`, so a plain refresh always shows the + latest after a server restart (restart the server to pick up code edits). +- Single-user local tool: it serves on localhost only and reads files read-only. diff --git a/nvflow/core/__init__.py b/nvflow/core/__init__.py index a3e88f8..20059e5 100644 --- a/nvflow/core/__init__.py +++ b/nvflow/core/__init__.py @@ -14,12 +14,30 @@ # """Core infrastructure for workflow orchestration.""" +from typing import TYPE_CHECKING, Any + from nvflow.core import console from nvflow.core.base_stage import BaseStage from nvflow.core.stage_registry import StageRegistry -from nvflow.core.workflow_runner import WorkflowRunner + +if TYPE_CHECKING: + from nvflow.core.workflow_runner import WorkflowRunner __all__ = ["BaseStage", "StageRegistry", "WorkflowRunner", "console"] + +def __getattr__(name: str) -> Any: + # WorkflowRunner pulls in omegaconf, which is absent from minimal worker + # containers (e.g. the SAM localization image). Those workers import only + # leaf helper modules under nvflow.recipes, and recipe auto-discovery + # touches this package -- so importing WorkflowRunner eagerly here would + # crash them with ModuleNotFoundError. Resolve it lazily instead. + if name == "WorkflowRunner": + from nvflow.core.workflow_runner import WorkflowRunner + + return WorkflowRunner + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + # Note: nemo-skills functions are imported directly in stage files when needed: # from nemo_skills.pipeline.cli import generate, run_cmd, wrap_arguments diff --git a/nvflow/core/workflow_runner.py b/nvflow/core/workflow_runner.py index cc6843a..c425f50 100644 --- a/nvflow/core/workflow_runner.py +++ b/nvflow/core/workflow_runner.py @@ -275,8 +275,9 @@ def run( self._run_stage(stage_name, environment=environment, stages_to_run=stages_to_run) completed_stages.append(stage_name) - header("βœ… Workflow Complete!") - success(f"Completed {len(completed_stages)} stage(s): {', '.join(completed_stages)}") + header("βœ… Workflow Submitted") + success(f"Submitted {len(completed_stages)} stage(s): {', '.join(completed_stages)}") + detail("Note", "Stages run as Slurm jobs -- track them with squeue") def _preflight_pipeline_health( self, diff --git a/nvflow/generic_stage/__init__.py b/nvflow/generic_stage/__init__.py new file mode 100644 index 0000000..efbef28 --- /dev/null +++ b/nvflow/generic_stage/__init__.py @@ -0,0 +1,15 @@ +# 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. +# +"""Shared stage implementations reusable across recipes.""" diff --git a/nvflow/generic_stage/sdg/__init__.py b/nvflow/generic_stage/sdg/__init__.py new file mode 100644 index 0000000..5e74ab1 --- /dev/null +++ b/nvflow/generic_stage/sdg/__init__.py @@ -0,0 +1,15 @@ +# 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. +# +"""Shared SDG stage implementations.""" diff --git a/nvflow/generic_stage/sdg/document_grounded/__init__.py b/nvflow/generic_stage/sdg/document_grounded/__init__.py new file mode 100644 index 0000000..0ac2c98 --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/__init__.py @@ -0,0 +1,49 @@ +# 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. +# +"""Shared DG-SDG stages and per-recipe registration helper.""" + +from nvflow.core import StageRegistry + +from .aggregate_answers import AggregateAnswersStage +from .dg_sdg_preprocess import DGSDGPreprocessStage +from .dgsdg_post_process import DGSDGPostProcessStage +from .evaluate_answers import EvaluateAnswersStage +from .generate_answers import GenerateAnswersStage +from .generate_verified_questions import GenerateVerifiedQuestionsStage +from .gym_genselect_answers import GymGenselectAnswersStage + +WORKFLOW = "document_grounded_sdg" +SHARED_STAGES: list[tuple[type, str]] = [ + (AggregateAnswersStage, "aggregate_answers"), + (EvaluateAnswersStage, "evaluate_answers"), + (GymGenselectAnswersStage, "gym_genselect_answers"), + (GenerateVerifiedQuestionsStage, "generate_verified_questions"), + (GenerateAnswersStage, "generate_answers"), + (DGSDGPreprocessStage, "dg_sdg_preprocess"), + (DGSDGPostProcessStage, "dgsdg_post_process"), +] + + +def register_for_recipe(recipe: str) -> None: + """Register all shared DG-SDG stages for a concrete recipe name.""" + for stage_class, stage_name in SHARED_STAGES: + if StageRegistry.has(recipe=recipe, workflow=WORKFLOW, stage=stage_name): + raise ValueError( + f"register_for_recipe({recipe!r}) would re-register " + f"{recipe}.{WORKFLOW}.{stage_name}. " + "This usually means old per-recipe shim modules are still imported " + "or the helper was called twice." + ) + StageRegistry.register(recipe=recipe, workflow=WORKFLOW, stage=stage_name)(stage_class) diff --git a/nvflow/generic_stage/sdg/document_grounded/_helpers.py b/nvflow/generic_stage/sdg/document_grounded/_helpers.py new file mode 100644 index 0000000..47e6565 --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/_helpers.py @@ -0,0 +1,365 @@ +# 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. +# +"""Shared helpers for generic DG-SDG stages.""" + +import json +import shlex +from typing import Any + +from ._schemas import ALWAYS_DROP, STAGE_KEEP + + +def clean_stale_experiments(cluster: str, expnames: list[str]) -> None: + """Remove ``/experiments//`` for each name in *expnames*. + + ``rollout()`` (reused unmodified from RL) does not clean stale nemo-run + experiment dirs, but SDG needs it: nemo-run caches the generated bash + scripts per experiment, so a stale dir makes (a) code edits silently + no-op (cached scripts re-used; SKILL.md Gotcha #8) and (b) ``run_after`` + resolve to a stale FINISHED experiment, skipping the Slurm dependency + (Gotcha #1). Replicated here on the SDG side so RL code stays untouched. + Idempotent; safe because nemo-run regenerates scripts on next launch and + we run before any new job is submitted. + """ + import shutil + from pathlib import Path + + import nemo_skills.pipeline.utils as pipeline_utils + + cluster_config = pipeline_utils.get_cluster_config(cluster) + job_dir = cluster_config.get("job_dir") + if not job_dir: + return + root = Path(job_dir) / "experiments" + if not root.is_dir(): + return + for expname in expnames: + target = root / expname + if target.is_dir(): + shutil.rmtree(target, ignore_errors=True) + + +ENRICH_MODULE = "nvflow.lib.sdg.document_grounded.enrich_rollouts" +ENRICH_MODULE_EVALUATE = "nvflow.lib.sdg.document_grounded.enrich_rollouts_evaluate" +ANALYZE_MODULE = "nvflow.lib.sdg.document_grounded.analyze_rollouts" + + +def submit_gym_generation( + *, + cluster: str, + rollout_expname: str, + run_after: list[str] | None, + input_file: str, + output_dir: str, + prompt_template: str, + gym_path: str, + gym_config_paths: list[str], + gym_agent_name: str, + container: str, + installation_command: str | None, + model_path: str, + num_gpus: int, + server_nodes: int = 1, + num_chunks: int = 1, + num_random_seeds: int = 1, + inference_params: dict[str, Any] | None = None, + vllm_extra: dict[str, Any] | None = None, + extra_record_fields: dict[str, Any] | None = None, + extra_record_field_mappers: dict[str, str] | None = None, + enrich_module: str = ENRICH_MODULE, + rerun_done: bool = False, + gym_uv_venv_dir: str = "", +) -> None: + """Render SDG JSONL to Responses API, then collect rollouts via ``rollout()``. + + Up to two jobs are submitted: + + 1. ``{rollout_expname}-render`` (CPU): ``responses_api render_and_convert`` + turns the flat SDG input into Responses-API rows (per-row prompt under + ``responses_create_params.input`` + per-row ``verifier`` from + ``extra_record_fields``). ``inference_params`` are NOT rendered in -- + they are applied by ``rollout()`` as global ``responses_create_params`` + overrides, keeping ``responses_create_params.input`` stable so the + content-hash join in ``enrich`` matches input<->output rows. + SKIPPED when the render output already exists (unless ``rerun_done``): + re-rendering on resume is wasteful and races a resumed merge's enrich + (see the guard below). When skipped, ``rollout()`` inherits the render's + own ``run_after`` so downstream ordering is preserved. + 2. ``rollout()`` (GPU): chunk + ng_collect_rollouts + per-seed merge, then + the merge job runs ``enrich`` (restore SDG fields + extract generation) + and ``analyze`` (sync ``rollout/output-rs*.jsonl`` up to ``output_dir/``). + + The caller is responsible for any per-stage trim / postprocess, submitted + as a separate ``run_cmd`` under the *stage* expname with + ``run_after=[rollout_expname]`` (so downstream ``run_after=[stage_expname]`` + waits for trim -> rollout). + """ + from nemo_skills.pipeline.cli import run_cmd, wrap_arguments + + from nvflow.core import console + from nvflow.lib.rl.helpers import resolve_host_path + from nvflow.lib.rl.rollout import rollout + + rapi_file = f"{output_dir}/.responses_api_input.jsonl" + render_expname = f"{rollout_expname}-render" + + render_cmd_parts = [ + "python -m nvflow.lib.sdg.document_grounded.responses_api render_and_convert", + f"--input_file {shlex.quote(input_file)}", + f"--output_file {shlex.quote(rapi_file)}", + f"--prompt_template {shlex.quote(prompt_template)}", + ] + if extra_record_fields: + payload = json.dumps(extra_record_fields) + render_cmd_parts.append(f"--extra_record_fields {shlex.quote(payload)}") + if extra_record_field_mappers: + payload = json.dumps(extra_record_field_mappers) + render_cmd_parts.append(f"--extra_record_field_mappers {shlex.quote(payload)}") + render_cmd = " ".join(render_cmd_parts) + + # Skip re-rendering when the Responses-API input already exists. The render + # is a deterministic 1:1 transform of *input_file*, so recomputing it on a + # resume is pure waste (100s of GB rewrite). It is also unsafe: the per-seed + # merge's enrich() reads THIS exact file, and when a seed's chunks are all + # `.done` the merge loses its (transitive, via chunk jobs) dependency on the + # render -- it then runs immediately and can race a concurrent render rewrite, + # reading a half-written file (enrich alignment-check failure). Skipping the + # render keeps the input stable for any resumed merge. Mirrors the + # skip-if-exists guards on the q-prep / q-verify-prep steps; `rerun_done` + # forces a fresh render, kept in lock-step with the rollout rerun. + # NOTE: execute() runs on the orchestrator node, so resolve the container + # path to its host path before checking existence. + rapi_host = resolve_host_path(rapi_file) + rapi_exists = rapi_host.exists() and rapi_host.stat().st_size > 0 + if rapi_exists and not rerun_done: + console.success("Render skipped (reusing existing Responses-API input)") + console.detail("Responses-API input", rapi_file) + rollout_run_after = run_after + else: + run_cmd( + ctx=wrap_arguments(render_cmd), + cluster=cluster, + expname=render_expname, + log_dir=f"{output_dir}/render-logs", + run_after=run_after, + ) + rollout_run_after = [render_expname] + + cfg = build_rollout_config( + input_file=rapi_file, + output_dir=output_dir, + gym_path=gym_path, + gym_config_paths=gym_config_paths, + gym_agent_name=gym_agent_name, + container=container, + installation_command=installation_command, + model_path=model_path, + num_gpus=num_gpus, + server_nodes=server_nodes, + num_chunks=num_chunks, + num_random_seeds=num_random_seeds, + inference_params=inference_params, + vllm_extra=vllm_extra, + rerun_done=rerun_done, + gym_uv_venv_dir=gym_uv_venv_dir, + ) + rollout( + config=cfg, + cluster=cluster, + expname=rollout_expname, + run_after=rollout_run_after, + enrich_module=enrich_module, + analyze_module=ANALYZE_MODULE, + ) + + +def parse_stage_kwargs(stage_kwargs: dict[str, Any]) -> dict[str, Any]: + """Extract normalized fields from a legacy ``args`` / ``ctx_args`` block. + + Returns a dict with ``model_path``, ``num_gpus``, ``server_nodes``, + ``num_chunks``, ``num_random_seeds``, ``prompt_template``, + ``generation_key``, ``inference_params`` and ``vllm_extra`` (any remaining + ``args`` keys that are vLLM serve flags). Used by the generate_* shims to + feed both the render step (prompt_template) and :func:`build_rollout_config`. + """ + args = stage_kwargs.get("args", {}).copy() + ctx_args = stage_kwargs.get("ctx_args", "") + + model_path = args.pop("model", "") + num_gpus = args.pop("server_gpus", args.pop("num_gpus", 8)) + server_nodes = args.pop("server_nodes", 1) + num_chunks = args.pop("num_chunks", 1) + num_random_seeds = args.pop("num_random_seeds", 1) + args.pop("server_type", None) + args.pop("skip_filled", None) + + prompt_template = "" + generation_key = "generation" + inference_params: dict[str, Any] = {} + for part in ctx_args.split(): + if part.startswith("++prompt_config="): + prompt_template = part.split("=", 1)[1] + elif part.startswith("++inference."): + key = part.split("=")[0].replace("++inference.", "") + val = part.split("=", 1)[1] + try: + inference_params[key] = float(val) + except ValueError: + inference_params[key] = val + elif part.startswith("++generation_key="): + generation_key = part.split("=", 1)[1] + + vllm_extra = {k: v for k, v in args.items() if k != "generation_key"} + + return { + "model_path": model_path, + "num_gpus": num_gpus, + "server_nodes": server_nodes, + "num_chunks": num_chunks, + "num_random_seeds": num_random_seeds, + "prompt_template": prompt_template, + "generation_key": generation_key, + "inference_params": inference_params, + "vllm_extra": vllm_extra, + } + + +def build_rollout_config( + *, + input_file: str, + output_dir: str, + gym_path: str, + gym_config_paths: list[str], + gym_agent_name: str, + container: str, + installation_command: str | None, + model_path: str, + num_gpus: int, + server_nodes: int = 1, + num_chunks: int = 1, + num_random_seeds: int = 1, + inference_params: dict[str, Any] | None = None, + vllm_extra: dict[str, Any] | None = None, + rerun_done: bool = False, + env_key: str = "sdg_format_verification", + gym_uv_venv_dir: str = "", +) -> dict[str, Any]: + """Translate SDG generation params into a config for ``rollout()``. + + ``rollout()`` is reused unmodified (the adapter lives entirely on the SDG + side). Notes: + + - ``input_file`` MUST already be in Responses API format (per-row + ``responses_create_params.input`` + per-row ``verifier``), produced by + ``responses_api.render_and_convert``. The per-row prompt and verifier + live in the data, NOT here. + - ``inference_params`` (temperature, top_p, max_output_tokens, ...) become + global ``responses_create_params`` overrides applied by ng_collect. + - No ``judge_vllm`` is set -> ``determine_judge_mode`` returns + ``policy_as_judge`` (no judge server). + - ``environments`` carries the SDG overlay; ``build_config_paths_str`` + prepends the vLLM model config automatically. + """ + # Some knobs are rollout-level (consumed by ``rollout()``), not vLLM serve + # flags, but they arrive mixed into ``vllm_extra`` from a stage's ``args`` / + # ``policy_vllm`` block. Intercept them here so they reach the ``rollout`` + # config instead of leaking into ``policy_vllm`` -> ``build_vllm_server_args`` + # as invalid CLI flags. + # - num_samples_in_parallel: concurrent requests per server (default 4). + # - dependent_jobs: chained resume jobs per chunk so a rollout that doesn't + # finish inside the Slurm walltime continues in the next chained job + # (default 0). Essential for big/slow models where one 4h job can't + # finish (long-tail generations) -- the chained job resumes the few + # remaining samples and exits early once done. + rollout_level_keys = ("num_samples_in_parallel", "dependent_jobs") + extra = dict(vllm_extra or {}) + rollout_level = {k: extra.pop(k) for k in rollout_level_keys if k in extra} + + policy_vllm: dict[str, Any] = { + "model_path": model_path, + "num_gpus": num_gpus, + "server_nodes": server_nodes, + } + policy_vllm.update(extra) + + rollout_cfg: dict[str, Any] = { + "input_data": input_file, + "policy_vllm": policy_vllm, + "responses_create_params": inference_params or {}, + "num_chunks": num_chunks, + "num_random_seeds": num_random_seeds, + "rerun_done": rerun_done, + } + rollout_cfg.update(rollout_level) + + return { + "output_dir": output_dir, + "gym_path": gym_path, + "gym_uv_venv_dir": gym_uv_venv_dir, + "container": container, + "installation_command": installation_command, + "rollout": rollout_cfg, + "environments": { + env_key: { + "agent_name": gym_agent_name, + "config_paths": list(gym_config_paths), + } + }, + } + + +def build_trim_cmd( + *, + stage_name: str, + paths: list[str], + domain_keep_fields: list[str] | None, + extra_keep_fields: list[str] | None = None, +) -> str: + """Build the shell command that trims this stage's output JSONL files. + + The returned string invokes ``nvflow.generic_stage.sdg.document_grounded._trim_cli`` + with the keep-list ``(STAGE_KEEP[stage_name] | domain_keep_fields - + ALWAYS_DROP) | extra_keep_fields`` and the given ``paths`` (files, + directories, or globs -- the CLI expands them). + + ``extra_keep_fields`` is unioned *after* the ``ALWAYS_DROP`` subtraction, so + it is the only way to retain a field that is otherwise in ``ALWAYS_DROP`` + (e.g. ``responses_create_params`` on the final ``dgsdg_post_process`` + output, where the Responses-API original form must survive). Use sparingly. + + The command is meant to be either: + - appended to ``postprocess_cmd`` for Gym-driven stages + (``sdg_generate``-based: question gen/verify, answer gen, genselect, + evaluate), so it runs inside the merge job and the producing stage's + advertised expname does not need to change; or + - chained via ``&&`` to the stage's main CPU command for non-Gym stages + (aggregate, difficulty aggregate, post-process). + + Either way the trim is guaranteed to finish before any downstream stage's + ``run_after`` clears, with zero extra Slurm overhead. + """ + if stage_name not in STAGE_KEEP: + raise KeyError( + f"build_trim_cmd: stage {stage_name!r} is not in STAGE_KEEP. " + f"Known stages: {sorted(STAGE_KEEP)}" + ) + domain = set(domain_keep_fields or []) + keep = ((STAGE_KEEP[stage_name] | domain) - ALWAYS_DROP) | set(extra_keep_fields or []) + keep_args = " ".join(sorted(keep)) + paths_arg = " ".join(shlex.quote(p) for p in paths) + return ( + "python -m nvflow.generic_stage.sdg.document_grounded._trim_cli " + f"--paths {paths_arg} --keep_fields {keep_args}" + ) diff --git a/nvflow/generic_stage/sdg/document_grounded/_schemas.py b/nvflow/generic_stage/sdg/document_grounded/_schemas.py new file mode 100644 index 0000000..6631ce5 --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/_schemas.py @@ -0,0 +1,181 @@ +# 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. +# +"""Per-stage JSONL field allowlists for DG-SDG. + +Each generic DG-SDG stage projects its output JSONL to ``STAGE_KEEP[stage] | +domain_keep_fields`` (set union) at the stage boundary, before the next stage +reads it. Goal: drop stale fields that would silently contaminate downstream +stages -- most importantly the NeMo-Gym rollout metadata and the ``generation`` +/ ``reasoning_content`` keys that get overwritten by every Gym call. + +Domain-specific fields (e.g. ``company_name``, ``file_path0``) are supplied +per recipe via the workflow YAML key ``domain_keep_fields`` and unioned with +``STAGE_KEEP[stage]`` at trim time. Generic stage code never hardcodes them. +""" + +# Cross-stage scratch / noise that we *never* want to survive a stage boundary. +# These are always dropped on top of (i.e. removed from) the per-stage KEEP +# allowlist so that even if a future contributor adds one to STAGE_KEEP by +# mistake, the trim still filters it out. +ALWAYS_DROP: frozenset[str] = frozenset( + { + # NeMo-Gym rollout passthrough metadata (added by responses_api on every + # Gym call; never read downstream). + "_ng_task_index", + "_ng_rollout_index", + "agent_ref", + "reward", + "match_details", + "verifier", + # Generation-time bookkeeping added by responses_api / Gym workers. + "serialized_output", + "num_generated_tokens", + "finish_reason", + "generation_start_time", + "generation_end_time", + "generation_time", + "responses_create_params", + } +) + + +# Per-stage allowlist of *generic* fields (i.e. fields that the lib code +# produces or that downstream lib code needs). Domain-specific fields come +# from the workflow YAML's ``domain_keep_fields`` and are unioned at trim time. +# +# Stage 0 (``dg_sdg_preprocess``) is intentionally absent: it manufactures the +# initial JSONL from raw documents, so there is no upstream record to project +# from. The Stage 1 trim acts as the safety net if the recipe writes junk. +STAGE_KEEP: dict[str, frozenset[str]] = { + # Q-side output (``verified/output-rs*.jsonl``): keep the Yes/No + # ``generation`` because the A-prep step votes on it; drop the Q-verify + # CoT (``reasoning_content``) -- nobody downstream reads it. + "generate_verified_questions": frozenset( + { + "context", + "problem", + "question_type", + "generation", + } + ), + # A-side output (``generated/output-rs*.jsonl``): keep the answer text + # (``generation``) and the answer CoT (``reasoning_content``); both get + # snapshotted into ``reference_*`` by genselect.postprocess in Stage 3. + # + # ``answer_response`` / ``answer_responses_create_params`` carry the *full* + # Responses-API original form of each candidate answer (the exact request + + # response object the A-gen model produced). They are the literal + # ``response`` / ``responses_create_params`` snapshotted under a non- + # ALWAYS_DROP alias by ``enrich_rollouts`` so the trim keeps them. + # genselect collapses the per-seed ``answer_response`` into + # ``answer_responses_list`` and selects one into ``reference_response`` for + # the final post-process output (Responses-API ``final_result.jsonl``). + "generate_answers": frozenset( + { + "context", + "problem", + "question_type", + "question_voting_pass_rate", + "question_voting_total", + "generation", + "reasoning_content", + "answer_response", + "answer_responses_create_params", + } + ), + # GenSelect-picked output (``selected_answers.jsonl``): ``reference_*`` + # carry the selected answer through evaluate/aggregate/difficulty; + # ``generation`` carries the same selected answer as the prompt input for + # evaluate. Genselect scaffolding (solutions/generations_list/answer_N/...) + # is dropped because it has served its purpose. + "gym_genselect_answers": frozenset( + { + "context", + "problem", + "question_type", + "question_voting_pass_rate", + "question_voting_total", + "reference_answer", + "reference_reasoning", + "reference_response", + "reference_responses_create_params", + "generation", + "genselect_answers_metadata", + } + ), + # Multi-seed eval rollouts: keep ``evaluate_generation`` for aggregate to + # parse; drop ``reasoning_content`` which by now is the evaluate-judge CoT + # (not the answer reasoning) and would otherwise silently overwrite the + # real answer CoT carried in ``reference_reasoning``. + "evaluate_answers": frozenset( + { + "context", + "problem", + "question_type", + "question_voting_pass_rate", + "question_voting_total", + "reference_answer", + "reference_reasoning", + "reference_response", + "reference_responses_create_params", + "generation", + "evaluate_generation", + } + ), + # Aggregated answers: per-seed ``evaluate_generation`` and ``correct`` are + # dropped; the consensus ``answerable`` survives. + "aggregate_answers": frozenset( + { + "context", + "problem", + "question_type", + "question_voting_pass_rate", + "question_voting_total", + "reference_answer", + "reference_reasoning", + "reference_response", + "reference_responses_create_params", + "generation", + "answerable", + } + ), + # Final training data (``final_result.jsonl``): post-process has already + # renamed ``reference_reasoning -> reasoning_content`` and + # ``reference_answer -> answer``, so the allowlist uses the post-rename + # names. ``genselect_answers_metadata`` is intentionally dropped from the + # final output -- it was useful for debugging mid-pipeline but is noise + # for SFT / RL. + # ``response`` + ``responses_create_params`` are the Responses-API original form + # post-process restores (renamed from ``reference_response`` / + # ``reference_responses_create_params``); ``expected_answer`` mirrors + # ``answer``. Note ``responses_create_params`` is in ALWAYS_DROP, so the + # post-process stage re-adds it via ``build_trim_cmd(extra_keep_fields=...)`` + # -- listing it here is documentation; the trim would otherwise strip it. + "dgsdg_post_process": frozenset( + { + "context", + "problem", + "answer", + "reasoning_content", + "question_type", + "answerable", + "question_voting_pass_rate", + "question_voting_total", + "expected_answer", + "response", + "responses_create_params", + } + ), +} diff --git a/nvflow/generic_stage/sdg/document_grounded/_trim_cli.py b/nvflow/generic_stage/sdg/document_grounded/_trim_cli.py new file mode 100644 index 0000000..91db185 --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/_trim_cli.py @@ -0,0 +1,121 @@ +# 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. +# +"""Trim DG-SDG JSONL files in place to a per-stage allowlist. + +Invoked at every DG-SDG stage boundary (either as part of the producing job's +``postprocess_cmd`` for Gym stages, or chained with ``&&`` to the CPU command +for non-Gym stages). Drops every JSON key not present in ``--keep_fields``, +including the cross-stage scratch listed in ``_schemas.ALWAYS_DROP``. + +The trim is in-place via a ``.trim_tmp`` rename, so partial failures +don't leave a half-written file at the canonical path. + +Usage:: + + python -m nvflow.generic_stage.sdg.document_grounded._trim_cli \\ + --paths /abs/path/to/file.jsonl /abs/path/to/dir \\ + --keep_fields context problem generation +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def trim_file(path: Path, keep: set[str]) -> tuple[int, int]: + """Rewrite ``path`` in place keeping only top-level keys in ``keep``. + + Returns ``(records_processed, field_instances_dropped)``. + """ + tmp = path.with_suffix(path.suffix + ".trim_tmp") + rec_count = 0 + drop_count = 0 + with path.open() as fin, tmp.open("w") as fout: + for line in fin: + stripped = line.strip() + if not stripped: + continue + record = json.loads(stripped) + slim = {k: v for k, v in record.items() if k in keep} + drop_count += len(record) - len(slim) + fout.write(json.dumps(slim) + "\n") + rec_count += 1 + tmp.replace(path) + return rec_count, drop_count + + +def _resolve_paths(args_paths: list[str]) -> list[Path]: + """Expand globs and directories into a flat list of JSONL files.""" + matched: list[Path] = [] + for raw in args_paths: + candidate = Path(raw) + if "*" in raw or "?" in raw: + matched.extend(sorted(candidate.parent.glob(candidate.name))) + elif candidate.is_dir(): + # Unlike shell globs, pathlib's glob("*.jsonl") also matches + # dotfiles (e.g. ``.responses_api_input.jsonl``, the internal + # render/join cache written by responses_api.render_and_convert). + # That file is never a stage *output* -- trimming it strips + # ``responses_create_params`` (ALWAYS_DROP), which enrich_rollouts' + # join key is computed from, silently poisoning the cache for any + # future re-merge. Exclude dotfiles to match intended shell-glob + # semantics and keep internal caches out of stage-boundary trims. + matched.extend( + sorted(p for p in candidate.glob("*.jsonl") if not p.name.startswith(".")) + ) + else: + matched.append(candidate) + return matched + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--paths", + nargs="+", + required=True, + help="JSONL files, directories (globbed as *.jsonl), or glob patterns.", + ) + parser.add_argument( + "--keep_fields", + nargs="+", + required=True, + help="Top-level JSON keys to keep. Everything else is dropped.", + ) + args = parser.parse_args(argv) + + keep = set(args.keep_fields) + files = _resolve_paths(args.paths) + if not files: + print( + f"[trim] no files matched from {args.paths!r}; nothing to do", + file=sys.stderr, + ) + return 0 + + for f in files: + if not f.exists(): + print(f"[trim] {f}: missing, skipping", file=sys.stderr) + continue + n, d = trim_file(f, keep) + print(f"[trim] {f}: {n} records, dropped {d} field-instances") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/nvflow/recipes/finance/stages/sdg/aggregate_answers.py b/nvflow/generic_stage/sdg/document_grounded/aggregate_answers.py similarity index 66% rename from nvflow/recipes/finance/stages/sdg/aggregate_answers.py rename to nvflow/generic_stage/sdg/document_grounded/aggregate_answers.py index 7f39179..d759e44 100644 --- a/nvflow/recipes/finance/stages/sdg/aggregate_answers.py +++ b/nvflow/generic_stage/sdg/document_grounded/aggregate_answers.py @@ -17,27 +17,13 @@ from pathlib import Path from typing import Any -from nvflow.core import BaseStage, StageRegistry, console +from nvflow.core import BaseStage, console +from ._helpers import build_trim_cmd -@StageRegistry.register( - recipe="finance", - workflow="document_grounded_sdg", - stage="aggregate_answers", -) -class AggregateAnswersStage(BaseStage): - """Aggregate multi-seed evaluation results. - - This stage processes output-rs*.jsonl files in streaming mode: - - Reads all seed files line-by-line in parallel (no intermediate files) - - Parses evaluate_generation inline - - Only keeps records where ALL seeds have correct=YES - - Only keeps records where ALL seeds have consistent answerable (all YES or all NO) - - Adds a final 'answerable' field based on the consistent value - This ensures high-quality data where the evaluation is confident and consistent - across multiple random samples. Uses O(1) memory regardless of file size. - """ +class AggregateAnswersStage(BaseStage): + """Aggregate multi-seed evaluation results.""" workflow = "document_grounded_sdg" @@ -61,22 +47,21 @@ def execute( console.detail("Num seeds", str(num_seeds)) console.blank() - # The evaluate_answers stage creates: {input_dir}/{input_file_stem}/output-rsN.jsonl - # Input file stem is "selected_answers" based on workflow config generation_folder = Path(input_dir) / "selected_answers" - - aggregate_module = "nvflow.recipes.finance.utils.sdg.aggregate_evaluate" - - # Aggregate results (parse + aggregate combined, no intermediate files) - full_cmd = ( - f"python3 -m {aggregate_module} " + aggregate_cmd = ( + "python -m nvflow.lib.sdg.document_grounded.aggregate " f"--input_dir {generation_folder} " f"--output_file {output_file} " f"--num_seeds {num_seeds}" ) + trim_cmd = build_trim_cmd( + stage_name="aggregate_answers", + paths=[output_file], + domain_keep_fields=config.get("domain_keep_fields"), + ) + full_cmd = f"{aggregate_cmd} && {trim_cmd}" console.status("Running aggregation (streaming, no intermediate files)") - run_cmd( ctx=wrap_arguments(full_cmd), cluster=cluster, diff --git a/nvflow/recipes/finance/stages/sdg/dg_sdg_preprocess.py b/nvflow/generic_stage/sdg/document_grounded/dg_sdg_preprocess.py similarity index 50% rename from nvflow/recipes/finance/stages/sdg/dg_sdg_preprocess.py rename to nvflow/generic_stage/sdg/document_grounded/dg_sdg_preprocess.py index c0ada3f..5e4bea6 100644 --- a/nvflow/recipes/finance/stages/sdg/dg_sdg_preprocess.py +++ b/nvflow/generic_stage/sdg/document_grounded/dg_sdg_preprocess.py @@ -12,32 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""SEC Data Preprocessing Stage for Document-Grounded SDG. - -This stage processes raw SEC filings (10-K and 10-Q HTML files) into structured JSONL data: -1. Chunk HTML files into Markdown, Clean HTML, and Original HTML -2. Generate CSV file lists from chunked files -3. Generate JSONL training data from CSVs -""" +"""Data preprocessing stage for Document-Grounded SDG.""" from typing import Any -from nvflow.core import BaseStage, StageRegistry, console +from nvflow.core import BaseStage, console +from nvflow.lib.rl.helpers import resolve_host_path -@StageRegistry.register( - recipe="finance", - workflow="document_grounded_sdg", - stage="dg_sdg_preprocess", -) class DGSDGPreprocessStage(BaseStage): - """Preprocess SEC filings for document-grounded SDG. - - This stage converts raw SEC HTML filings into structured JSONL data: - 1. Chunks HTML files by token count with overlap - 2. Generates CSV file lists for tracking chunks - 3. Creates JSONL training data with proper sampling distribution - """ + """Preprocess domain documents into structured JSONL data.""" workflow = "document_grounded_sdg" @@ -48,34 +32,70 @@ def execute( expname: str, run_after: list[str] | None = None, ) -> None: - """Execute the SEC data preprocessing pipeline.""" + """Execute the data preprocessing pipeline.""" from nemo_skills.pipeline.cli import run_cmd, wrap_arguments input_dir = config["input_dir"] output_dir = config["output_dir"] distribution_dir = config["distribution_dir"] - # Chunking settings max_tokens = config.get("max_tokens", 2000) overlap_tokens = config.get("overlap_tokens", 100) - - # Sampling settings total_samples = config.get("total_samples", 150000) max_skip_count = config.get("max_skip_count", 20000) seed = config.get("seed", 42) - - console.status("SEC Data Preprocessing") + preprocess_module = config["preprocess_module"] + rerun_done = config.get("rerun_done", False) + + # Domain-agnostic passthrough: arbitrary extra CLI args forwarded verbatim + # to the preprocess_module. Lets domain recipes pass module-specific flags + # (e.g. the SEC recipe's --forms) without this generic stage knowing about + # them. Bool True -> bare flag; other values -> "--key value" (quoted). + extra_args = config.get("extra_args") or {} + extra_parts: list[str] = [] + for key, value in extra_args.items(): + if isinstance(value, bool): + if value: + extra_parts.append(f"--{key}") + elif isinstance(value, list | tuple): + extra_parts.append(f"--{key} '{' '.join(str(v) for v in value)}'") + elif isinstance(value, str): + extra_parts.append(f"--{key} '{value}'") + else: + extra_parts.append(f"--{key} {value}") + extra_args_str = " ".join(extra_parts) + + console.status("Document data preprocessing") console.detail("Input dir", input_dir) console.detail("Output dir", output_dir) console.detail("Distribution dir", distribution_dir) + console.detail("Preprocess module", preprocess_module) console.detail("Max tokens", str(max_tokens)) console.detail("Overlap tokens", str(overlap_tokens)) console.detail("Total samples", str(total_samples)) console.detail("Max skip count", str(max_skip_count)) console.detail("Seed", str(seed)) + if extra_args_str: + console.detail("Extra args", extra_args_str) console.blank() - preprocess_module = "nvflow.recipes.finance.utils.sdg.dg_sdg_data_preprocess" + # Reuse previously materialized sampling output by default. + # Set rerun_done=true to force a full regenerate. + # + # ``execute()`` runs on the orchestrator/login node, so ``output_dir`` + # (a container path like ``/workspace/...``) must be resolved to its + # host path before the existence check -- otherwise it never matches and + # sampling re-runs on every launch. + forms_arg = str((extra_args or {}).get("forms", "10-K 10-Q")) + forms = [f for f in forms_arg.split() if f] + host_jsonl_dir = resolve_host_path(f"{output_dir}/jsonl") + if forms and not rerun_done: + expected_outputs = [host_jsonl_dir / f"{form.lower()}-data.jsonl" for form in forms] + all_present = all(p.exists() and p.stat().st_size > 0 for p in expected_outputs) + if all_present: + console.success("Data preprocessing skipped (reusing existing sampled output)") + console.detail("Output directory", output_dir) + return full_cmd = ( f"python3 -m {preprocess_module} " @@ -88,8 +108,8 @@ def execute( f"--max_skip_count {max_skip_count} " f"--seed {seed}" ) - - console.status("Running SEC data preprocessing") + if extra_args_str: + full_cmd += f" {extra_args_str}" run_cmd( ctx=wrap_arguments(full_cmd), @@ -98,12 +118,12 @@ def execute( run_after=run_after, ) - console.success("SEC data preprocessing job submitted") + console.success("Data preprocessing job submitted") console.detail("Output directory", output_dir) def validate_config(self, config: dict[str, Any]) -> None: """Validate stage configuration.""" - required = ["input_dir", "output_dir", "distribution_dir"] + required = ["input_dir", "output_dir", "distribution_dir", "preprocess_module"] for field in required: if field not in config: raise ValueError(f"Missing required field: {field}") diff --git a/nvflow/recipes/finance/stages/sdg/document_grounded_data.py b/nvflow/generic_stage/sdg/document_grounded/dgsdg_post_process.py similarity index 57% rename from nvflow/recipes/finance/stages/sdg/document_grounded_data.py rename to nvflow/generic_stage/sdg/document_grounded/dgsdg_post_process.py index 360ab0c..8c4c341 100644 --- a/nvflow/recipes/finance/stages/sdg/document_grounded_data.py +++ b/nvflow/generic_stage/sdg/document_grounded/dgsdg_post_process.py @@ -12,32 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Document grounded sdg data post processing stage.""" +"""Document grounded SDG data post processing stage.""" from typing import Any -from nvflow.core import BaseStage, StageRegistry, console +from nvflow.core import BaseStage, console +from ._helpers import build_trim_cmd -@StageRegistry.register( - recipe="finance", - workflow="document_grounded_sdg", - stage="dgsdg_post_process", -) -class DGSDGPostProcessStage(BaseStage): - """Post process document grounded sdg data by cleaning fields and creating subsets. - This stage: - 1. Removes unwanted fields (solutions, generations_list, etc.) - 2. Renames reference_reasoning -> reasoning_content, reference_answer -> answer - 3. Creates full_data.jsonl with all cleaned records - 4. Creates medium_sft_data.jsonl: - - Only records with difficulty_score in [1, 2, 3, 4] - - For 10-K filings: excludes Risk_Factors questions - - For 10-Q filings: only includes Risk_Factors questions - 5. Creates hard_rl_data.jsonl: - - Only records with difficulty_score = 0 - """ +class DGSDGPostProcessStage(BaseStage): + """Post process document grounded SDG data by cleaning fields and creating subsets.""" workflow = "document_grounded_sdg" @@ -48,24 +33,36 @@ def execute( expname: str, run_after: list[str] | None = None, ) -> None: - """Execute document grounded sdg data post processing.""" + """Execute document grounded SDG data post processing.""" from nemo_skills.pipeline.cli import run_cmd, wrap_arguments input_file = config["input_file"] output_dir = config["output_dir"] seed = config.get("seed", 42) + postprocess_script = config["postprocess_script"] - console.status("Post processing document grounded sdg data") + console.status("Post processing document grounded SDG data") console.detail("Input file", input_file) console.detail("Output dir", output_dir) console.detail("Random seed", str(seed)) + console.detail("Postprocess script", postprocess_script) console.blank() - module = "nvflow.recipes.finance.utils.sdg.dgsdg_post_process" - - cmd = ( - f"python3 -m {module} --input_file {input_file} --output_dir {output_dir} --seed {seed}" + postprocess_cmd = ( + f"python {postprocess_script} " + f"--input_file {input_file} " + f"--output_dir {output_dir} " + f"--seed {seed}" + ) + trim_cmd = build_trim_cmd( + stage_name="dgsdg_post_process", + paths=[f"{output_dir}/final_result.jsonl"], + domain_keep_fields=config.get("domain_keep_fields"), + # ``responses_create_params`` is in ALWAYS_DROP; re-add it here so the + # final Responses-API record retains the original request. + extra_keep_fields=["responses_create_params"], ) + cmd = f"{postprocess_cmd} && {trim_cmd}" run_cmd( ctx=wrap_arguments(cmd), @@ -74,12 +71,12 @@ def execute( run_after=run_after, ) - console.success("Document grounded sdg data post processing job submitted") + console.success("Document grounded SDG data post processing job submitted") console.detail("Output files will be in", output_dir) def validate_config(self, config: dict[str, Any]) -> None: """Validate stage configuration.""" - required = ["input_file", "output_dir"] + required = ["input_file", "output_dir", "postprocess_script"] for field in required: if field not in config: raise ValueError(f"Missing required field: {field}") diff --git a/nvflow/generic_stage/sdg/document_grounded/evaluate_answers.py b/nvflow/generic_stage/sdg/document_grounded/evaluate_answers.py new file mode 100644 index 0000000..14556db --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/evaluate_answers.py @@ -0,0 +1,152 @@ +# 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. +# +"""Evaluate answers for correctness and answerability.""" + +from pathlib import Path +from typing import Any + +from nvflow.core import BaseStage, console + +from ._helpers import ( + ENRICH_MODULE_EVALUATE, + build_trim_cmd, + clean_stale_experiments, + submit_gym_generation, +) + + +class EvaluateAnswersStage(BaseStage): + """Evaluate answers for correctness and answerability.""" + + workflow = "document_grounded_sdg" + + def execute( + self, + config: dict[str, Any], + cluster: str, + expname: str, + run_after: list[str] | None = None, + ) -> None: + """Execute answer evaluation and filtering.""" + from nemo_skills.pipeline.cli import run_cmd, wrap_arguments + + clean_stale_experiments(cluster, [f"{expname}-gen", f"{expname}-gen-render", expname]) + + input_file = config["input_file"] + output_dir = config.get("output_dir") + output_file = config.get("output_file") + prompt_template = config.get("prompt_template", config.get("prompt_config", "")) + generation_key = config.get("generation_key", "evaluate_generation") + inference_params = config.get("inference_params", {}) + num_random_seeds = config.get("num_random_seeds", 1) + + if generation_key != "evaluate_generation": + console.warning( + "evaluate_answers currently pins generation field to " + "'evaluate_generation' (rollout enrich hook is fixed-arg); " + f"configured generation_key='{generation_key}' is ignored." + ) + + console.status("Evaluating answers for correctness and answerability (NeMo-Gym)") + console.detail("Input file", input_file) + console.detail("Output dir", str(output_dir)) + console.detail("Prompt template", prompt_template) + console.detail("Num random seeds", str(num_random_seeds)) + console.blank() + + if output_dir: + generation_folder = Path(output_dir) / Path(input_file).stem + else: + generation_folder = Path(output_file).parent / Path(input_file).stem + + console.detail("Generation folder", str(generation_folder)) + + lib_evaluate = "python -m nvflow.lib.sdg.document_grounded.evaluate" + domain_keep_fields = config.get("domain_keep_fields") + + pv = dict(config.get("policy_vllm", {})) + model_path = pv.pop("model_path", "") + num_gpus = pv.pop("num_gpus", 8) + server_nodes = pv.pop("server_nodes", 1) + + console.status("Running LLM evaluation via NeMo-Gym") + gen_expname = f"{expname}-gen" + submit_gym_generation( + cluster=cluster, + rollout_expname=gen_expname, + run_after=run_after, + input_file=input_file, + output_dir=str(generation_folder), + prompt_template=prompt_template, + gym_path=config["gym_path"], + gym_config_paths=config.get("gym_config_paths", []), + gym_agent_name=config["gym_agent_name"], + container=config.get("container", "nemo-rl"), + installation_command=config.get("installation_command"), + gym_uv_venv_dir=config.get("gym_uv_venv_dir", ""), + model_path=model_path, + num_gpus=num_gpus, + server_nodes=server_nodes, + num_chunks=config.get("num_chunks", 1), + num_random_seeds=num_random_seeds, + inference_params=inference_params, + vllm_extra=pv, + extra_record_fields=config.get("extra_record_fields"), + extra_record_field_mappers=config.get("extra_record_field_mappers"), + enrich_module=ENRICH_MODULE_EVALUATE, + rerun_done=config.get("rerun_done", False), + ) + + # Parse/filter/trim (single-seed) or trim-only (multi-seed), run under + # the stage expname so downstream `run_after=[stage_expname]` waits. + if num_random_seeds <= 1: + generated_file = str(generation_folder / "output-rs0.jsonl") + parsed_file = str(generation_folder / "parsed.jsonl") + final_output = ( + output_file if output_file else str(generation_folder / "evaluated.jsonl") + ) + parse_cmd = ( + f"{lib_evaluate} parse --input_file {generated_file} --output_file {parsed_file}" + ) + filter_cmd = ( + f"{lib_evaluate} filter --input_file {parsed_file} --output_file {final_output}" + ) + trim_cmd = build_trim_cmd( + stage_name="evaluate_answers", + paths=[final_output], + domain_keep_fields=domain_keep_fields, + ) + postprocess_cmd = f"{parse_cmd} && {filter_cmd} && {trim_cmd}" + else: + postprocess_cmd = build_trim_cmd( + stage_name="evaluate_answers", + paths=[str(generation_folder)], + domain_keep_fields=domain_keep_fields, + ) + run_cmd( + ctx=wrap_arguments(postprocess_cmd), + cluster=cluster, + expname=expname, + log_dir=f"{generation_folder}/postprocess-logs", + run_after=[gen_expname], + ) + + console.success(f"Completed Answer Evaluation for: {input_file}") + if num_random_seeds > 1: + console.detail("Parsed outputs in", str(generation_folder)) + else: + console.detail( + "Output (correct answers only, with 'answerable' field)", str(output_file) + ) diff --git a/nvflow/generic_stage/sdg/document_grounded/generate_answers.py b/nvflow/generic_stage/sdg/document_grounded/generate_answers.py new file mode 100644 index 0000000..f0d2104 --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/generate_answers.py @@ -0,0 +1,195 @@ +# 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. +# +"""Answer generation pipeline for document-grounded SDG. + +Consumes verified-question records produced by GenerateVerifiedQuestionsStage +and emits N candidate answers per question for downstream genselect. +""" + +from typing import Any + +from nvflow.core import BaseStage, console +from nvflow.lib.rl.helpers import resolve_host_path + +from ._helpers import ( + build_trim_cmd, + clean_stale_experiments, + parse_stage_kwargs, + submit_gym_generation, +) + + +class GenerateAnswersStage(BaseStage): + """A-side of DG-SDG: a-prep (threshold filter) -> A-gen. + + Output layout under ``output_dir``:: + + answer_input.jsonl # step 1 output (questions surviving the + # verification threshold) + generated/ # step 2 output (A-gen rollouts; consumed by + # gym_genselect_answers) + """ + + workflow = "document_grounded_sdg" + + def execute( + self, + config: dict[str, Any], + cluster: str, + expname: str, + run_after: list[str] | None = None, + ) -> None: + from nemo_skills.pipeline.cli import run_cmd, wrap_arguments + + clean_stale_experiments( + cluster, + [ + f"{expname}-step1-a-prep", + f"{expname}-step2-a-gen", + f"{expname}-step2-a-gen-render", + expname, + ], + ) + + input_dir = config["input_dir"] + output_dir = config["output_dir"] + + gym_path = config["gym_path"] + gym_uv_venv_dir = config.get("gym_uv_venv_dir", "") + gym_config_paths_default = config.get("gym_config_paths", []) + gym_agent_name_default = config.get("gym_agent_name") + gym_container = config.get("container", "nemo-rl") + installation_command = config.get("installation_command") + extra_record_fields_default = config.get("extra_record_fields") + extra_record_field_mappers_default = config.get("extra_record_field_mappers") + + def _substep(prefix: str) -> dict[str, Any]: + agent = config.get(f"{prefix}_gym_agent_name", gym_agent_name_default) + if not agent: + raise ValueError( + f"generate_answers: '{prefix}_gym_agent_name' " + "(or stage-level 'gym_agent_name') is required." + ) + return { + "gym_config_paths": config.get( + f"{prefix}_gym_config_paths", gym_config_paths_default + ), + "gym_agent_name": agent, + "extra_record_fields": config.get( + f"{prefix}_extra_record_fields", extra_record_fields_default + ), + "extra_record_field_mappers": config.get( + f"{prefix}_extra_record_field_mappers", + extra_record_field_mappers_default, + ), + } + + a_gen_overrides = _substep("answer_generation") + + answer_preprocess_kwargs = config.get("answer_preprocess_kwargs", {}) + answer_generation_kwargs = config.get("answer_generation_kwargs", {}) + + a_generate_input_file = f"{output_dir}/answer_input.jsonl" + a_generate_output_dir = f"{output_dir}/generated" + + lib_preprocess = "python -m nvflow.lib.sdg.document_grounded.preprocess" + + # execute() runs on the orchestrator node: resolve the container path to + # its host path before checking existence (see _helpers.host_path). + step1_expname = f"{expname}-step1-a-prep" + rerun_a_prep = config.get("answer_prep_rerun_done", False) + a_prep_host = resolve_host_path(a_generate_input_file) + a_prep_exists = a_prep_host.exists() and a_prep_host.stat().st_size > 0 + a_prep_submitted = False + console.status("Step 1/2: Preparing data for answer generation") + console.detail("Output file", a_generate_input_file) + if a_prep_exists and not rerun_a_prep: + console.success("Step 1 skipped (reusing existing answer_input.jsonl)") + else: + console.detail("Input dir", input_dir) + threshold = answer_preprocess_kwargs.get("threshold", 0.5) + sbatch_kwargs = answer_preprocess_kwargs.get("sbatch_kwargs", "") + cmd = ( + f"{lib_preprocess} construct_answer_generate_input " + f"--input_dir {input_dir} " + f"--output_file {a_generate_input_file} " + f"--threshold {threshold}" + ) + run_cmd( + ctx=wrap_arguments(cmd), + cluster=cluster, + expname=step1_expname, + run_after=run_after, + sbatch_kwargs=sbatch_kwargs, + ) + a_prep_submitted = True + console.success("Step 1 job submitted") + + console.status("Step 2/2: Generating answers") + params = parse_stage_kwargs(answer_generation_kwargs) + a_gen_expname = f"{expname}-step2-a-gen" + submit_gym_generation( + cluster=cluster, + rollout_expname=a_gen_expname, + run_after=[step1_expname] if a_prep_submitted else run_after, + input_file=a_generate_input_file, + output_dir=a_generate_output_dir, + prompt_template=params["prompt_template"], + gym_path=gym_path, + gym_config_paths=a_gen_overrides["gym_config_paths"], + gym_agent_name=a_gen_overrides["gym_agent_name"], + container=gym_container, + installation_command=installation_command, + gym_uv_venv_dir=gym_uv_venv_dir, + model_path=params["model_path"], + num_gpus=params["num_gpus"], + server_nodes=params["server_nodes"], + num_chunks=params["num_chunks"], + num_random_seeds=params["num_random_seeds"], + inference_params=params["inference_params"], + vllm_extra=params["vllm_extra"], + extra_record_fields=a_gen_overrides["extra_record_fields"], + extra_record_field_mappers=a_gen_overrides["extra_record_field_mappers"], + rerun_done=answer_generation_kwargs.get("rerun_done", False), + ) + + # Per-stage trim runs under the *stage* expname (depends on a-gen) so + # downstream `run_after=[stage_expname]` waits for the trimmed output. + trim_cmd = build_trim_cmd( + stage_name="generate_answers", + paths=[a_generate_output_dir], + domain_keep_fields=config.get("domain_keep_fields"), + ) + run_cmd( + ctx=wrap_arguments(trim_cmd), + cluster=cluster, + expname=expname, + log_dir=f"{a_generate_output_dir}/trim-logs", + run_after=[a_gen_expname], + ) + console.success("Step 2 job submitted") + + console.blank() + console.success("Answer generation pipeline jobs submitted") + console.detail("Generated answers will be in", a_generate_output_dir) + + def validate_config(self, config: dict[str, Any]) -> None: + """Validate required configuration fields.""" + required = ["input_dir", "output_dir", "gym_path"] + for field in required: + if field not in config: + raise ValueError(f"Missing required field: {field}") + if "answer_generation_kwargs" not in config: + raise ValueError("Missing required field: answer_generation_kwargs") diff --git a/nvflow/generic_stage/sdg/document_grounded/generate_verified_questions.py b/nvflow/generic_stage/sdg/document_grounded/generate_verified_questions.py new file mode 100644 index 0000000..39a56a9 --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/generate_verified_questions.py @@ -0,0 +1,252 @@ +# 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. +# +"""Question generation + verification pipeline for document-grounded SDG.""" + +from typing import Any + +from nvflow.core import BaseStage, console +from nvflow.lib.rl.helpers import resolve_host_path + +from ._helpers import ( + build_trim_cmd, + clean_stale_experiments, + parse_stage_kwargs, + submit_gym_generation, +) + + +class GenerateVerifiedQuestionsStage(BaseStage): + """Q-side of DG-SDG: prep -> generate -> verify-prep -> verify. + + Output layout under ``output_dir``:: + + generate_input.jsonl # step 1 output + generated/ # step 2 output (Q-gen rollouts) + verify_input.jsonl # step 3 output + verified/ # step 4 output (Q-verify rollouts; consumed by + # the generate_answers stage) + """ + + workflow = "document_grounded_sdg" + + def execute( + self, + config: dict[str, Any], + cluster: str, + expname: str, + run_after: list[str] | None = None, + ) -> None: + from nemo_skills.pipeline.cli import run_cmd, wrap_arguments + + clean_stale_experiments( + cluster, + [ + f"{expname}-step1-q-prep", + f"{expname}-step2-q-gen", + f"{expname}-step2-q-gen-render", + f"{expname}-step3-q-verify-prep", + f"{expname}-step4-q-verify", + f"{expname}-step4-q-verify-render", + expname, + ], + ) + + input_folder = config["input_folder"] + output_dir = config["output_dir"] + question_prep_script = config["question_prep_script"] + + gym_path = config["gym_path"] + gym_uv_venv_dir = config.get("gym_uv_venv_dir", "") + gym_config_paths_default = config.get("gym_config_paths", []) + gym_agent_name_default = config.get("gym_agent_name") + gym_container = config.get("container", "nemo-rl") + installation_command = config.get("installation_command") + extra_record_fields_default = config.get("extra_record_fields") + extra_record_field_mappers_default = config.get("extra_record_field_mappers") + + def _substep(prefix: str) -> dict[str, Any]: + agent = config.get(f"{prefix}_gym_agent_name", gym_agent_name_default) + if not agent: + raise ValueError( + f"generate_verified_questions: '{prefix}_gym_agent_name' " + "(or stage-level 'gym_agent_name') is required." + ) + return { + "gym_config_paths": config.get( + f"{prefix}_gym_config_paths", gym_config_paths_default + ), + "gym_agent_name": agent, + "extra_record_fields": config.get( + f"{prefix}_extra_record_fields", extra_record_fields_default + ), + "extra_record_field_mappers": config.get( + f"{prefix}_extra_record_field_mappers", + extra_record_field_mappers_default, + ), + } + + q_gen_overrides = _substep("question_generation") + q_verify_overrides = _substep("question_verify") + + question_generation_kwargs = config.get("question_generation_kwargs", {}) + question_verify_kwargs = config.get("question_verify_kwargs", {}) + rerun_q_prep = config.get("question_prep_rerun_done", False) + rerun_q_verify_prep = config.get("question_verify_prep_rerun_done", False) + + q_generate_input_file = f"{output_dir}/generate_input.jsonl" + q_generate_output_dir = f"{output_dir}/generated" + q_verify_input_file = f"{output_dir}/verify_input.jsonl" + q_verify_output_dir = f"{output_dir}/verified" + + lib_preprocess = "python -m nvflow.lib.sdg.document_grounded.preprocess" + + step1_expname = f"{expname}-step1-q-prep" + # execute() runs on the orchestrator node: resolve the container path to + # its host path before checking existence (see _helpers.host_path). + step1_host = resolve_host_path(q_generate_input_file) + step1_exists = step1_host.exists() and step1_host.stat().st_size > 0 + step1_submitted = False + if step1_exists and not rerun_q_prep: + console.status("Step 1/4: Preparing data for question generation") + console.detail("Output file", q_generate_input_file) + console.success("Step 1 skipped (reusing existing generate_input.jsonl)") + else: + console.status("Step 1/4: Preparing data for question generation") + console.detail("Input folder", input_folder) + console.detail("Output file", q_generate_input_file) + cmd = ( + f"python {question_prep_script} " + f"--input_folder {input_folder} " + f"--output_file {q_generate_input_file}" + ) + run_cmd( + ctx=wrap_arguments(cmd), + cluster=cluster, + expname=step1_expname, + run_after=run_after, + ) + step1_submitted = True + console.success("Step 1 job submitted") + + console.status("Step 2/4: Generating questions") + q_gen_params = parse_stage_kwargs(question_generation_kwargs) + submit_gym_generation( + cluster=cluster, + rollout_expname=f"{expname}-step2-q-gen", + run_after=[step1_expname] if step1_submitted else run_after, + input_file=q_generate_input_file, + output_dir=q_generate_output_dir, + prompt_template=q_gen_params["prompt_template"], + gym_path=gym_path, + gym_config_paths=q_gen_overrides["gym_config_paths"], + gym_agent_name=q_gen_overrides["gym_agent_name"], + container=gym_container, + installation_command=installation_command, + gym_uv_venv_dir=gym_uv_venv_dir, + model_path=q_gen_params["model_path"], + num_gpus=q_gen_params["num_gpus"], + server_nodes=q_gen_params["server_nodes"], + num_chunks=q_gen_params["num_chunks"], + num_random_seeds=q_gen_params["num_random_seeds"], + inference_params=q_gen_params["inference_params"], + vllm_extra=q_gen_params["vllm_extra"], + extra_record_fields=q_gen_overrides["extra_record_fields"], + extra_record_field_mappers=q_gen_overrides["extra_record_field_mappers"], + rerun_done=question_generation_kwargs.get("rerun_done", False), + ) + console.success("Step 2 job submitted") + + step3_expname = f"{expname}-step3-q-verify-prep" + step3_host = resolve_host_path(q_verify_input_file) + step3_exists = step3_host.exists() and step3_host.stat().st_size > 0 + step3_submitted = False + if step3_exists and not rerun_q_verify_prep: + console.status("Step 3/4: Preparing data for question verification") + console.detail("Output file", q_verify_input_file) + console.success("Step 3 skipped (reusing existing verify_input.jsonl)") + else: + console.status("Step 3/4: Preparing data for question verification") + cmd = ( + f"{lib_preprocess} construct_question_verify_input " + f"--input_dir {q_generate_output_dir} " + f"--output_file {q_verify_input_file}" + ) + run_cmd( + ctx=wrap_arguments(cmd), + cluster=cluster, + expname=step3_expname, + run_after=[f"{expname}-step2-q-gen"], + ) + step3_submitted = True + console.success("Step 3 job submitted") + + console.status("Step 4/4: Verifying questions") + q_verify_params = parse_stage_kwargs(question_verify_kwargs) + q_verify_expname = f"{expname}-step4-q-verify" + submit_gym_generation( + cluster=cluster, + rollout_expname=q_verify_expname, + run_after=[step3_expname] if step3_submitted else [f"{expname}-step2-q-gen"], + input_file=q_verify_input_file, + output_dir=q_verify_output_dir, + prompt_template=q_verify_params["prompt_template"], + gym_path=gym_path, + gym_config_paths=q_verify_overrides["gym_config_paths"], + gym_agent_name=q_verify_overrides["gym_agent_name"], + container=gym_container, + installation_command=installation_command, + gym_uv_venv_dir=gym_uv_venv_dir, + model_path=q_verify_params["model_path"], + num_gpus=q_verify_params["num_gpus"], + server_nodes=q_verify_params["server_nodes"], + num_chunks=q_verify_params["num_chunks"], + num_random_seeds=q_verify_params["num_random_seeds"], + inference_params=q_verify_params["inference_params"], + vllm_extra=q_verify_params["vllm_extra"], + extra_record_fields=q_verify_overrides["extra_record_fields"], + extra_record_field_mappers=q_verify_overrides["extra_record_field_mappers"], + rerun_done=question_verify_kwargs.get("rerun_done", False), + ) + + # Stage trim runs under the stage expname (depends on q-verify) so the + # downstream stage's `run_after=[stage_expname]` waits for trimmed output. + trim_cmd = build_trim_cmd( + stage_name="generate_verified_questions", + paths=[q_verify_output_dir], + domain_keep_fields=config.get("domain_keep_fields"), + ) + run_cmd( + ctx=wrap_arguments(trim_cmd), + cluster=cluster, + expname=expname, + log_dir=f"{q_verify_output_dir}/trim-logs", + run_after=[q_verify_expname], + ) + console.success("Step 4 job submitted") + + console.blank() + console.success("Question generation + verification pipeline jobs submitted") + console.detail("Verified questions will be in", q_verify_output_dir) + + def validate_config(self, config: dict[str, Any]) -> None: + """Validate required configuration fields.""" + required = ["input_folder", "output_dir", "gym_path", "question_prep_script"] + for field in required: + if field not in config: + raise ValueError(f"Missing required field: {field}") + if "question_generation_kwargs" not in config: + raise ValueError("Missing required field: question_generation_kwargs") + if "question_verify_kwargs" not in config: + raise ValueError("Missing required field: question_verify_kwargs") diff --git a/nvflow/generic_stage/sdg/document_grounded/gym_genselect_answers.py b/nvflow/generic_stage/sdg/document_grounded/gym_genselect_answers.py new file mode 100644 index 0000000..df61c6c --- /dev/null +++ b/nvflow/generic_stage/sdg/document_grounded/gym_genselect_answers.py @@ -0,0 +1,155 @@ +# 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. +# +"""Generate and select best answers using NeMo-Gym inference.""" + +from typing import Any + +from nvflow.core import BaseStage, console +from nvflow.lib.rl.helpers import resolve_host_path + +from ._helpers import ( + build_trim_cmd, + clean_stale_experiments, + submit_gym_generation, +) + + +class GymGenselectAnswersStage(BaseStage): + """Generate and select best answers via NeMo-Gym collect_rollouts.""" + + workflow = "document_grounded_sdg" + + def execute( + self, + config: dict[str, Any], + cluster: str, + expname: str, + run_after: list[str] | None = None, + ) -> None: + """Execute genselect answer generation via rollout().""" + from nemo_skills.pipeline.cli import run_cmd, wrap_arguments + + clean_stale_experiments( + cluster, + [f"{expname}-prep", f"{expname}-gen", f"{expname}-gen-render", expname], + ) + + input_dir = config["input_dir"] + output_file = config["output_file"] + prompt_template = config["prompt_template"] + + output_dir = output_file.replace(".jsonl", "") + prepped_file = output_dir + "_prepped.jsonl" + + console.status("Generating and selecting best answers (NeMo-Gym)") + console.detail("Input dir", input_dir) + console.detail("Output file", output_file) + console.detail("Prepped file", prepped_file) + console.detail("Output dir", output_dir) + console.detail("Prompt template", prompt_template) + console.blank() + + # execute() runs on the orchestrator node: resolve the container path to + # its host path before checking existence (see _helpers.host_path). + prep_expname = f"{expname}-prep" + rerun_prep = config.get("genselect_prep_rerun_done", False) + prep_host = resolve_host_path(prepped_file) + prep_exists = prep_host.exists() and prep_host.stat().st_size > 0 + prep_submitted = False + console.status("Step 1: Preparing genselect data") + if prep_exists and not rerun_prep: + console.success("Step 1 skipped (reusing existing prepped genselect input)") + else: + run_cmd( + ctx=wrap_arguments( + f"python -m nvflow.lib.sdg.document_grounded.genselect merge " + f"--input_dir={input_dir} --output_file={prepped_file}" + ), + cluster=cluster, + expname=prep_expname, + log_dir=f"{output_dir}/prep-data-logs", + run_after=run_after, + ) + prep_submitted = True + + pv = dict(config.get("policy_vllm", {})) + model_path = pv.pop("model_path", "") or config.get("model", "") + num_gpus = pv.pop("num_gpus", 0) or config.get("server_gpus", 8) + server_nodes = pv.pop("server_nodes", 1) + + console.status("Step 2: Generating answers via NeMo-Gym") + gen_expname = f"{expname}-gen" + submit_gym_generation( + cluster=cluster, + rollout_expname=gen_expname, + run_after=[prep_expname] if prep_submitted else run_after, + input_file=prepped_file, + output_dir=output_dir, + prompt_template=prompt_template, + gym_path=config["gym_path"], + gym_config_paths=config.get("gym_config_paths", []), + gym_agent_name=config["gym_agent_name"], + container=config.get("container", "nemo-rl"), + installation_command=config.get("installation_command"), + gym_uv_venv_dir=config.get("gym_uv_venv_dir", ""), + model_path=model_path, + num_gpus=num_gpus, + server_nodes=server_nodes, + num_chunks=config.get("num_chunks", 1), + num_random_seeds=config.get("num_random_seeds", 1), + inference_params=config.get("inference_params", {}), + vllm_extra=pv, + extra_record_fields=config.get("extra_record_fields"), + extra_record_field_mappers=config.get("extra_record_field_mappers"), + rerun_done=config.get("rerun_done", False), + ) + + # Genselect postprocess (select best answer -> output_file) + trim, run + # under the stage expname so downstream `run_after=[stage_expname]` waits. + trim_cmd = build_trim_cmd( + stage_name="gym_genselect_answers", + paths=[output_file], + domain_keep_fields=config.get("domain_keep_fields"), + ) + postprocess_cmd = ( + f"cp {output_dir}/output-rs0.jsonl {output_dir}/output.jsonl && " + "python -m nvflow.lib.sdg.document_grounded.genselect postprocess " + f"--input_dir={output_dir} " + f"--output_file={output_file} && " + f"{trim_cmd}" + ) + run_cmd( + ctx=wrap_arguments(postprocess_cmd), + cluster=cluster, + expname=expname, + log_dir=f"{output_dir}/postprocess-logs", + run_after=[gen_expname], + ) + + console.success(f"Genselect answer generation submitted -> {output_file}") + + def validate_config(self, config: dict[str, Any]) -> None: + """Validate required configuration fields.""" + for field in ( + "input_dir", + "output_file", + "prompt_template", + "gym_path", + "gym_agent_name", + ): + if not config.get(field): + raise ValueError(f"'{field}' is required in genselect_answers config") + if not config.get("policy_vllm") and not config.get("model"): + raise ValueError("Either 'policy_vllm.model_path' or 'model' is required") diff --git a/nvflow/lib/cli_cmd.py b/nvflow/lib/cli_cmd.py new file mode 100644 index 0000000..bbf098f --- /dev/null +++ b/nvflow/lib/cli_cmd.py @@ -0,0 +1,165 @@ +# 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. +# +"""Shared shell command builders for stage submission. + +Stages that submit ``python3 -m [positional ...] --flag ...`` +shell commands to nemo-skills' ``run_cmd`` / ``generate`` should use +:func:`build_python_cmd` to build the command string rather than +concatenating raw f-strings. ``shlex.quote`` ensures values containing +spaces, single quotes, or shell metacharacters do not break the rendered +command -- this matters because nemo-skills interpolates the command +into a Slurm shell wrapper at submission time. + +Usage:: + + from nvflow.lib.cli_cmd import build_python_cmd + + # Flag-only invocation: + rendered = build_python_cmd( + "nvflow.recipes.finance.utils.rl.regex_prefilter_questions", + input_file=Path("/lustre/foo/in.jsonl"), + output_kept=Path("/lustre/foo/kept.jsonl"), + ) + + # With positional args (e.g. for ``argparse`` scripts that take + # ``input_files`` positionally): + rendered = build_python_cmd( + "nvflow.recipes.finance.utils.shared.dataset_transformer", + Path("/lustre/sdg/final_result.jsonl"), + output_file="/lustre/out/final.jsonl", + num_chunks=10, + ) + # β†’ "python3 -m ...dataset_transformer " + # "/lustre/sdg/final_result.jsonl " + # "--output_file /lustre/out/final.jsonl --num_chunks 10" +""" + +from __future__ import annotations + +import shlex +from pathlib import Path + +# All cluster containers in this repo provide ``python3`` (it is the +# canonical interpreter on every modern Linux base image we ship). We +# standardise on ``python3`` rather than ``python`` so ambiguity around +# the unversioned ``python`` symlink (absent in some minimal images) can +# never bite us. +_INTERPRETER = "python3" + + +def build_python_cmd( + module: str, + *positional: str | int | float | Path, + **flags: str | int | float | Path, +) -> str: + """Build a ``python3 -m [positional ...] --flag value ...`` shell command. + + Each positional and flag value is passed through :func:`shlex.quote` + so paths containing spaces, single quotes, or shell metacharacters + do not break the rendered command -- this command string is + interpolated by nemo-skills into a Slurm shell wrapper, so safe + quoting matters. + + Accepts ``str``, numeric types, or :class:`pathlib.Path` values; + non-string values are stringified via :func:`str` before quoting. + Positional args are emitted in argument order, then flags in + declaration order. This keeps the rendered command stable for + log-grepping and diffing across reruns. + + Args: + module: Fully-qualified Python module name (e.g. + ``"nvflow.recipes.finance.utils.shared.dataset_transformer"``). + *positional: Positional arguments emitted before any flags -- + useful for ``argparse``-style scripts that accept positional + inputs (e.g. one or more input file paths). + **flags: Keyword arguments rendered as ``-- `` + pairs in declaration order. Boolean flags (no value) must + be appended manually by the caller; this helper does not + support them because Python kwargs cannot express + "value-less" flags unambiguously. + + Returns: + A single-line shell command string suitable for nemo-skills' + ``run_cmd`` / ``generate`` ``ctx`` argument. + + Examples: + >>> build_python_cmd("foo.bar", input_file="/a/b.jsonl") + 'python3 -m foo.bar --input_file /a/b.jsonl' + >>> build_python_cmd("foo.bar", "/a/in.jsonl", output_file="/a/out.jsonl") + 'python3 -m foo.bar /a/in.jsonl --output_file /a/out.jsonl' + >>> build_python_cmd("foo.bar", input_file="/a path/with spaces.jsonl") + "python3 -m foo.bar --input_file '/a path/with spaces.jsonl'" + """ + parts = [_INTERPRETER, "-m", module] + for arg in positional: + parts.append(shlex.quote(str(arg))) + for flag, value in flags.items(): + parts.extend([f"--{flag}", shlex.quote(str(value))]) + return " ".join(parts) + + +def build_python_script_cmd( + script: str | Path, + *positional: str | int | float | Path, + **flags: str | int | float | Path, +) -> str: + """Build a ``python3 + + +""" + + +def build_index_html(summary: HopchainHtmlVisualizationSummary, title: str) -> str: + """Build the index page that links to all rendered HTML chunks.""" + if summary.page_files: + rows = "\n".join( + f""" + + {page.page_number} + {html.escape(page.file_name)} + {page.row_count} + {page.start_row}-{page.end_row} + + """ + for page in summary.page_files + ) + table_html = f""" + + + + + + + + + + + {rows} + +
PageFileRowsRange
+ """ + else: + table_html = '

No generated query rows were available to render.

' + + return f""" + + + + {html.escape(title)} + + + +

{html.escape(title)}

+
+
Total queries: {summary.total_queries}
+
Rendered queries: {summary.rendered_queries}
+
Rows per file: {summary.rows_per_file}
+
Pages generated: {len(summary.page_files)}
+
Queries input: {html.escape(summary.queries_input)}
+
Combinations input: {html.escape(summary.combinations_input)}
+
+ {table_html} + + +""" + + +def main() -> None: + """Entry point.""" + args = parse_args() + if args.sample_count is not None and args.sample_count <= 0: + raise SystemExit("error: --sample-count must be a positive integer") + queries_input = Path(args.queries_input) + combinations_input = Path(args.combinations_input) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + image_max_dimension = None if args.image_max_dimension == 0 else args.image_max_dimension + + queries = load_queries(queries_input) + combinations_by_id = load_combinations(combinations_input) + renderable_records = join_records(queries, combinations_by_id) + sample_seed = None + if args.sample_count is not None and args.sample_count < len(renderable_records): + random.Random(args.sample_seed).shuffle(renderable_records) + renderable_records = renderable_records[: args.sample_count] + sample_seed = args.sample_seed + + total_pages = ( + max(1, math.ceil(len(renderable_records) / args.rows_per_file)) if renderable_records else 0 + ) + page_summaries: list[HtmlPageSummary] = [] + + for page_number in range(1, total_pages + 1): + start_idx = (page_number - 1) * args.rows_per_file + end_idx = min(start_idx + args.rows_per_file, len(renderable_records)) + page_records = renderable_records[start_idx:end_idx] + page_file_name = f"hopchain_review_{page_number:04d}.html" + page_path = output_dir / page_file_name + page_path.write_text( + build_page_html( + page_records=page_records, + page_number=page_number, + total_pages=total_pages, + total_records=len(renderable_records), + rows_per_file=args.rows_per_file, + image_max_dimension=image_max_dimension, + title=args.title, + ) + ) + page_summaries.append( + HtmlPageSummary( + page_number=page_number, + file_name=page_file_name, + row_count=len(page_records), + start_row=start_idx + 1, + end_row=end_idx, + ) + ) + + summary = HopchainHtmlVisualizationSummary( + queries_input=str(queries_input), + combinations_input=str(combinations_input), + output_dir=str(output_dir), + index_file=str(output_dir / "index.html"), + total_queries=len(queries), + rendered_queries=len(renderable_records), + sample_seed=sample_seed, + rows_per_file=args.rows_per_file, + image_max_dimension=image_max_dimension, + page_files=page_summaries, + ) + + (output_dir / "index.html").write_text(build_index_html(summary, title=args.title)) + Path(args.summary).write_text(summary.model_dump_json(indent=2)) + logger.info("Wrote %s HTML page(s) to %s", len(page_summaries), output_dir) + + +if __name__ == "__main__": + main() diff --git a/nvflow/recipes/multimodal/utils/runtime_env.py b/nvflow/recipes/multimodal/utils/runtime_env.py new file mode 100644 index 0000000..10fafc0 --- /dev/null +++ b/nvflow/recipes/multimodal/utils/runtime_env.py @@ -0,0 +1,45 @@ +# 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. +# +"""Helpers for cluster-specific runtime settings.""" + +from __future__ import annotations + +from typing import Any + + +def resolve_partition( + config: dict[str, Any], + cluster: str, + *, + cpu: bool = False, + override_key: str = "partition", +) -> str: + """Resolve a stage partition from YAML or the selected cluster config.""" + if config.get(override_key): + return str(config[override_key]) + + from nemo_skills.pipeline.utils import get_cluster_config + + cluster_config = get_cluster_config( + cluster=cluster, + config_dir=config.get("cluster_config_dir"), + ) + key = "cpu_partition" if cpu else "partition" + partition = cluster_config.get(key) + if partition is None and cpu: + partition = cluster_config.get("partition") + if partition is None: + raise ValueError(f"Cluster config for {cluster!r} does not define {key!r}") + return str(partition) diff --git a/nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml b/nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml new file mode 100644 index 0000000..96973f7 --- /dev/null +++ b/nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter-demo.yaml @@ -0,0 +1,14 @@ +# Small, deterministic HopChain image-filter run used by the quick start. +# Put images under data/images, then run this file directly. + +_base_: hopchain-image-filter.yaml + +execution_id: demo + +stages: + image_filter: + image_directories: + - directory: ${project_root}/data/images + recursive: true + end_index: 100 + num_chunks: 1 diff --git a/nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter.yaml b/nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter.yaml new file mode 100644 index 0000000..eb94ee0 --- /dev/null +++ b/nvflow/recipes/multimodal/workflows/image_filter/hopchain-image-filter.yaml @@ -0,0 +1,69 @@ +# HopChain image filtering workflow +# This workflow is intentionally separate from the SDG workflow because it runs +# at a different scale and produces reusable filtered-image artifacts. + +recipe: multimodal + +workflow: + name: "hopchain_image_filter" + type: "image_filter" + description: "Filter candidate HopChain images before SDG synthesis" + +cluster: my_cluster + +model_profiles: + qwen: + model: qwen3.5-397b-a17b + server_type: sglang + server_gpus: 8 + server_nodes: 2 + server_args: >- + --model-path /hf_models/Qwen/Qwen3.5-397B-A17B + --served-model-name qwen3.5-397b-a17b + --context-length 131072 + --tp 16 + --ep-size 16 + --trust-remote-code + --reasoning-parser qwen3 + --mem-fraction-static 0.80 + --chunked-prefill-size 4096 + max_image_dimension: 2048 + tokens_to_generate: 16384 + max_concurrent_requests: 64 + time_min: 120 + use_base64_images: true + +# Run from the repository root. Machine-specific mounts, partitions, and +# container paths belong in cluster_configs/my_cluster.yaml. +project_root: ${oc.env:PWD} +base_data_dir: ${project_root}/outputs/hopchain +execution_id: full + +cluster_config_dir: ${project_root}/cluster_configs +image_filter_data_dir: ${base_data_dir}/image_filter + +directories: + image-filter: ${image_filter_data_dir}/execution/${execution_id}/image-filter + +pipeline_stages: + - image_filter + +stages: + image_filter: + image_directories: + - directory: ${project_root}/data/images + recursive: true + output_dir: ${directories.image-filter} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: nvflow/recipes/multimodal/prompts/hopchain_image_filter.txt + model_config: ${model_profiles.qwen} + min_complexity_score: 4 + allowed_quality_ratings: + - High + - Medium + temperature: 0.0 + top_p: 1.0 + tokens_to_generate: 16384 + num_chunks: 6 + dependencies: [] diff --git a/nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml b/nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml new file mode 100644 index 0000000..a2e2257 --- /dev/null +++ b/nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg-demo.yaml @@ -0,0 +1,40 @@ +# Small, deterministic HopChain SDG run used by the quick start. +# It consumes the output from hopchain-image-filter-demo.yaml and stops before +# the optional external API judge and downstream curation stages. + +_base_: hopchain-sdg.yaml + +execution_id: demo +image_filter_execution_id: demo +localize_dataloader_num_workers: 2 +localize_num_chunks: 1 +generate_multihop_queries_num_queries: 1 +generate_multihop_queries_num_chunks: 1 + +pipeline_stages: + - prepare_filtered_image_inputs + - preprocess_identify_categories + - identify_categories + - localize_instances + - sample_instance_combinations + - preprocess_generate_multihop_queries + - generate_multihop_queries + - verify_candidate_queries + - visualize_candidate_hopchain_data + +stages: + prepare_filtered_image_inputs: + sample_count: 25 + sample_seed: 42 + + identify_categories: + num_chunks: 1 + + localize_instances: + debug_save_annotated_images: false + + sample_instance_combinations: + max_combinations_per_image: 2 + + visualize_candidate_hopchain_data: + sample_count: 32 diff --git a/nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg.yaml b/nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg.yaml new file mode 100644 index 0000000..d5ae9ef --- /dev/null +++ b/nvflow/recipes/multimodal/workflows/sdg/hopchain-sdg.yaml @@ -0,0 +1,485 @@ +# HopChain SDG workflow + +recipe: multimodal + +workflow: + name: "hopchain_sdg" + type: "sdg" + description: "HopChain-inspired multimodal synthetic data generation workflow" + +cluster: my_cluster + +model_profiles: + qwen: + model: qwen3.5-397b-a17b + server_type: sglang + server_gpus: 8 + server_nodes: 2 + server_args: >- + --model-path /hf_models/Qwen/Qwen3.5-397B-A17B + --served-model-name qwen3.5-397b-a17b + --context-length 131072 + --tp 16 + --ep-size 16 + --trust-remote-code + --reasoning-parser qwen3 + --mem-fraction-static 0.80 + --chunked-prefill-size 4096 + max_image_dimension: 2048 + tokens_to_generate: 16384 + max_concurrent_requests: 64 + time_min: 120 + use_base64_images: true + + omni: + model: omni + server_type: vllm + server_gpus: 4 + server_nodes: 1 + server_args: >- + --model /hf_models/nvidia/omni-step70 + --served-model-name omni + --trust-remote-code + --max-model-len 32768 + --allowed-local-media-path / + --gpu-memory-utilization 0.9 + --limit-mm-per-prompt.image 1 + --mamba_ssm_cache_dtype float32 + --reasoning-parser nemotron_v3 + tokens_to_generate: 30000 + max_concurrent_requests: 4 + time_min: 240 + use_base64_images: false + + sam: + model: /hf_models/facebook/sam3.1/sam3.1_multiplex.pt + container: vllm + +# Run from the repository root. Machine-specific mounts, partitions, and +# container paths belong in cluster_configs/my_cluster.yaml. +project_root: ${oc.env:PWD} +base_data_dir: ${project_root}/outputs/hopchain +execution_id: full +image_filter_execution_id: full + +cluster_config_dir: ${project_root}/cluster_configs +sdg_data_dir: ${base_data_dir}/sdg +base_output_dir: ${sdg_data_dir}/execution/${execution_id} +source_kept_images_file: ${base_data_dir}/image_filter/execution/${image_filter_execution_id}/image-filter/kept_images.jsonl +gpu_time_min: 240 +cpu_time_min: 240 +max_image_dimension: 1536 +localize_dataloader_num_workers: 16 +localize_dataloader_prefetch_factor: 2 +localize_num_chunks: 6 +sample_instance_combinations_selection_strategy: balanced_by_category_then_area_confidence +sample_instance_combinations_min_instances: 3 +sample_instance_combinations_max_instances: 8 +sample_instance_combinations_max_instances_considered_per_image: 24 +sample_instance_combinations_max_instances_per_category: 2 +sample_instance_combinations_iou_dedup_threshold: 0.9 +sample_instance_combinations_iou_dedup_candidate_pool_size: 72 +sample_instance_combinations_debug_copy_selected_images: false +sample_instance_combinations_max_combinations_per_image: 6 +sample_instance_combinations_size_strategy: weighted_random +sample_instance_combinations_size_weights: + 3: 1 + 4: 2 + 5: 3 + 6: 4 + 7: 4 + 8: 4 +# area_confidence strategy params (used when selection_strategy is +# balanced_by_category_then_area_confidence or area_confidence) +sample_instance_combinations_area_confidence_area_weight: 0.5 +sample_instance_combinations_area_confidence_confidence_weight: 0.5 +sample_instance_combinations_sampling_seed: 20260430 +# sample_instance_combinations_min_confidence_threshold: 0.6 # uncomment to filter low-confidence instances +generate_multihop_queries_num_queries: 2 +generate_multihop_queries_target_hop_count_info: "4-10 hops" +generate_multihop_queries_temperature: 0.2 +generate_multihop_queries_top_p: 0.95 +generate_multihop_queries_tokens_to_generate: 32768 +# Heuristic: ~ (input_size * max_combinations_per_image / 2 * num_queries) / 1000 * factor. +# Check the actual step-3 combination count before enabling paid judge stages. +generate_multihop_queries_num_chunks: 18 # ~ 2000 * 6 / 2 * 2 / 1000 * 1.5 = 18 +verify_candidate_queries_min_hop_count: 4 +llm_judge_prompt_file: nvflow/recipes/multimodal/prompts/hopchain_llm_judge_answer_question.txt +llm_judge_temperature: 0.0 +llm_judge_top_p: 1.0 +llm_judge_reasoning_effort: medium +llm_judge_timeout_seconds: 300 +llm_judge_max_retries: 3 +llm_judge_max_workers: 64 +llm_judge_time_min: 240 +difficulty_filter_k: 5 +# 2K images * 6 max combos/image / 2 expected kept combos * 2 queries +# * ~70% LLM-judge accepted * difficulty_filter_k ~= 42K Omni requests. Omni is roughly +# 2K requests/hour/chunk, so use 21 chunks to target about 1 hour. +difficulty_filter_num_chunks: 21 +difficulty_filter_temperature: 0.6 +difficulty_filter_top_p: 0.95 +difficulty_filter_enable_thinking: false # note: omni-step70 always reasons via vLLM --reasoning-parser; this flag sets chat_template_kwargs per-datapoint which nemo_skills inference does not read +difficulty_filter_prompt_file: nvflow/recipes/multimodal/prompts/hopchain_llm_judge_answer_question.txt +difficulty_filter_min_pass_rate: 0.0 # inclusive lower bound; 0.0 = no lower filter +difficulty_filter_max_pass_rate: 1.0 # inclusive upper bound; 1.0 = no upper filter +sft_trace_k: 3 +sft_trace_generation_prompt_file: nvflow/recipes/multimodal/prompts/hopchain_sft_answer_question.txt +sft_trace_judge_prompt_file: nvflow/recipes/multimodal/prompts/hopchain_sft_trace_judge.txt +sft_trace_generation_num_chunks: 21 +sft_trace_judge_num_chunks: 21 +sft_trace_generation_temperature: 0.6 +sft_trace_generation_top_p: 0.95 +sft_trace_generation_tokens_to_generate: 32768 +sft_trace_judge_temperature: 0.6 +sft_trace_judge_top_p: 0.95 +sft_trace_judge_tokens_to_generate: 32768 + +directories: + step-0-prepare-filtered-inputs: ${base_output_dir}/step-0-prepare-filtered-inputs + step-1-identify-categories: ${base_output_dir}/step-1-identify-categories + step-1-preprocess-identify-categories: ${base_output_dir}/step-1-identify-categories/temp + step-2-localize-instances: ${base_output_dir}/step-2-localize-instances + step-3-sample-instance-combinations: ${base_output_dir}/step-3-sample-instance-combinations + step-4-preprocess-generate-multihop-queries: ${base_output_dir}/step-4-generate-multihop-queries/temp + step-4-generate-multihop-queries: ${base_output_dir}/step-4-generate-multihop-queries + step-5-verify-candidate-queries: ${base_output_dir}/step-5-verify-candidate-queries + step-6-visualize-candidate-hopchain-data: ${base_output_dir}/step-6-visualize-candidate-hopchain-data + step-7-judge-candidate-queries-openai: ${base_output_dir}/step-7-judge-candidate-queries-openai + step-8-reconcile-llm-judges: ${base_output_dir}/step-8-reconcile-llm-judges + step-9-visualize-reconciled-hopchain-data: ${base_output_dir}/step-9-visualize-reconciled-hopchain-data + step-10-preprocess-filter-easy-candidates: ${base_output_dir}/step-10-filter-easy-candidates/temp + step-10-filter-easy-candidates: ${base_output_dir}/step-10-filter-easy-candidates + step-11-preprocess-generate-sft-reasoning-traces: ${base_output_dir}/step-11-generate-sft-reasoning-traces/temp + step-11-generate-sft-reasoning-traces: ${base_output_dir}/step-11-generate-sft-reasoning-traces + step-12-preprocess-filter-sft-reasoning-traces: ${base_output_dir}/step-12-filter-sft-reasoning-traces/temp + step-12-filter-sft-reasoning-traces: ${base_output_dir}/step-12-filter-sft-reasoning-traces + +pipeline_stages: + - prepare_filtered_image_inputs + - preprocess_identify_categories + - identify_categories + - localize_instances + - sample_instance_combinations + - preprocess_generate_multihop_queries + - generate_multihop_queries + - verify_candidate_queries + - visualize_candidate_hopchain_data + - judge_candidate_queries_openai + - reconcile_llm_judges + - visualize_reconciled_hopchain_data + - preprocess_filter_easy_candidates + - filter_easy_candidates + - preprocess_generate_sft_reasoning_traces + - generate_sft_reasoning_traces + - preprocess_filter_sft_reasoning_traces + - filter_sft_reasoning_traces + +stages: + prepare_filtered_image_inputs: + run_name: ${execution_id} + input_file: ${source_kept_images_file} + output_file: ${directories.step-0-prepare-filtered-inputs}/filtered_image_inputs.jsonl + summary_file: ${directories.step-0-prepare-filtered-inputs}/summary.json + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + # sample_count: 100 # disabled for full-shard 10k execution; uncomment for small calibration runs + # sample_count_per_domain: 40 # disabled for full-shard 10k execution; uncomment for per-domain calibration runs + # sample_seed: 42 + time_min: ${cpu_time_min} + dependencies: [] + + preprocess_identify_categories: + run_name: ${execution_id} + input_file: ${directories.step-0-prepare-filtered-inputs}/filtered_image_inputs.jsonl + output_file: ${directories.step-1-preprocess-identify-categories}/input_openai_format.jsonl + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: nvflow/recipes/multimodal/prompts/hopchain_category_identification.txt + model_config: ${model_profiles.qwen} + max_image_dimension: ${max_image_dimension} + time_min: ${cpu_time_min} + dependencies: + - prepare_filtered_image_inputs + + identify_categories: + run_name: ${execution_id} + input_file: ${directories.step-0-prepare-filtered-inputs}/filtered_image_inputs.jsonl + preprocessed_input_file: ${directories.step-1-preprocess-identify-categories}/input_openai_format.jsonl + output_dir: ${directories.step-1-identify-categories} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: nvflow/recipes/multimodal/prompts/hopchain_category_identification.txt + model_config: ${model_profiles.qwen} + max_image_dimension: ${max_image_dimension} + temperature: 0.0 + top_p: 1.0 + tokens_to_generate: 8192 + num_chunks: 2 # ~ input_size / 1000 = 2000 / 1000 = 2 + time_min: ${gpu_time_min} + dependencies: + - preprocess_identify_categories + + localize_instances: + run_name: ${execution_id} + input_file: ${directories.step-1-identify-categories}/final_output.jsonl + output_dir: ${directories.step-2-localize-instances} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + container: ${model_profiles.sam.container} + prompt_file: nvflow/recipes/multimodal/prompts/hopchain_instance_localization_sam3.txt + model: ${model_profiles.sam.model} + filter_list: + - filter_method: min_image_size + min_w: 50 + min_h: 50 + max_image_dimension: ${max_image_dimension} + num_gpus: 1 + dataloader_num_workers: ${localize_dataloader_num_workers} + dataloader_prefetch_factor: ${localize_dataloader_prefetch_factor} + num_chunks: ${localize_num_chunks} + max_localization_phrases_per_category: 3 + prompt_alias_iou_dedup_threshold: 0.9 + debug_save_annotated_images: true + debug_annotated_images_dir: ${directories.step-2-localize-instances}/annotated_images + threshold: 0.5 + time_min: ${gpu_time_min} + dependencies: + - identify_categories + + sample_instance_combinations: + run_name: ${execution_id} + input_file: ${directories.step-2-localize-instances}/final_output.jsonl + output_file: ${directories.step-3-sample-instance-combinations}/instance_combinations.jsonl + summary_file: ${directories.step-3-sample-instance-combinations}/summary.json + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + min_instances: ${sample_instance_combinations_min_instances} + max_instances: ${sample_instance_combinations_max_instances} + selection_strategy: ${sample_instance_combinations_selection_strategy} + max_instances_considered_per_image: ${sample_instance_combinations_max_instances_considered_per_image} + max_instances_per_category: ${sample_instance_combinations_max_instances_per_category} + iou_dedup_threshold: ${sample_instance_combinations_iou_dedup_threshold} + iou_dedup_candidate_pool_size: ${sample_instance_combinations_iou_dedup_candidate_pool_size} + area_confidence_area_weight: ${sample_instance_combinations_area_confidence_area_weight} + area_confidence_confidence_weight: ${sample_instance_combinations_area_confidence_confidence_weight} + debug_copy_selected_images: ${sample_instance_combinations_debug_copy_selected_images} + debug_selected_images_dir: ${directories.step-3-sample-instance-combinations}/selected_images + max_combinations_per_image: ${sample_instance_combinations_max_combinations_per_image} + combination_size_strategy: ${sample_instance_combinations_size_strategy} + combination_size_weights: ${sample_instance_combinations_size_weights} + sampling_seed: ${sample_instance_combinations_sampling_seed} + time_min: ${cpu_time_min} + dependencies: + - localize_instances + + preprocess_generate_multihop_queries: + run_name: ${execution_id} + input_file: ${directories.step-3-sample-instance-combinations}/instance_combinations.jsonl + output_file: ${directories.step-4-preprocess-generate-multihop-queries}/input_openai_format.jsonl + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: nvflow/recipes/multimodal/prompts/hopchain_query_design.txt + model_config: ${model_profiles.qwen} + num_queries: ${generate_multihop_queries_num_queries} + target_hop_count_info: ${generate_multihop_queries_target_hop_count_info} + time_min: ${cpu_time_min} + dependencies: + - sample_instance_combinations + + generate_multihop_queries: + run_name: ${execution_id} + input_file: ${directories.step-3-sample-instance-combinations}/instance_combinations.jsonl + preprocessed_input_file: ${directories.step-4-preprocess-generate-multihop-queries}/input_openai_format.jsonl + output_dir: ${directories.step-4-generate-multihop-queries} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: nvflow/recipes/multimodal/prompts/hopchain_query_design.txt + model_config: ${model_profiles.qwen} + num_queries: ${generate_multihop_queries_num_queries} + target_hop_count_info: ${generate_multihop_queries_target_hop_count_info} + temperature: ${generate_multihop_queries_temperature} + top_p: ${generate_multihop_queries_top_p} + tokens_to_generate: ${generate_multihop_queries_tokens_to_generate} + num_chunks: ${generate_multihop_queries_num_chunks} # see top-level estimate + time_min: ${gpu_time_min} + dependencies: + - preprocess_generate_multihop_queries + + verify_candidate_queries: + run_name: ${execution_id} + input_file: ${directories.step-4-generate-multihop-queries}/final_output.jsonl + output_dir: ${directories.step-5-verify-candidate-queries} + output_file: ${directories.step-5-verify-candidate-queries}/verified_queries.jsonl + summary_file: ${directories.step-5-verify-candidate-queries}/summary.json + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + min_hop_count: ${verify_candidate_queries_min_hop_count} + time_min: ${cpu_time_min} + dependencies: + - generate_multihop_queries + + visualize_candidate_hopchain_data: + run_name: ${execution_id} + queries_input_file: ${directories.step-5-verify-candidate-queries}/final_candidates.jsonl + combinations_input_file: ${directories.step-3-sample-instance-combinations}/instance_combinations.jsonl + output_dir: ${directories.step-6-visualize-candidate-hopchain-data} + summary_file: ${directories.step-6-visualize-candidate-hopchain-data}/summary.json + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + rows_per_file: 100 + sample_count: 200 + sample_seed: 42 + image_max_dimension: 1024 + title: HopChain Candidate Query Review + time_min: ${cpu_time_min} + dependencies: + - verify_candidate_queries + + judge_candidate_queries_openai: + run_name: ${execution_id} + input_file: ${directories.step-5-verify-candidate-queries}/final_candidates.jsonl + output_file: ${directories.step-7-judge-candidate-queries-openai}/judged_candidates.jsonl + output_dir: ${directories.step-7-judge-candidate-queries-openai} + summary_file: ${directories.step-7-judge-candidate-queries-openai}/summary.json + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: ${llm_judge_prompt_file} + judge_name: openai_gpt_5_5 + provider: openai + model: gpt-5.5 + api_key_name: OPENAI_API_KEY + max_image_dimension: ${max_image_dimension} + temperature: ${llm_judge_temperature} + top_p: ${llm_judge_top_p} + reasoning_effort: ${llm_judge_reasoning_effort} + timeout_seconds: ${llm_judge_timeout_seconds} + max_retries: ${llm_judge_max_retries} + max_workers: ${llm_judge_max_workers} + time_min: ${llm_judge_time_min} + dependencies: + - verify_candidate_queries + + reconcile_llm_judges: + run_name: ${execution_id} + input_file: ${directories.step-5-verify-candidate-queries}/final_candidates.jsonl + output_file: ${directories.step-8-reconcile-llm-judges}/reconciled_queries.jsonl + output_dir: ${directories.step-8-reconcile-llm-judges} + summary_file: ${directories.step-8-reconcile-llm-judges}/summary.json + judge_output_files: + - ${directories.step-7-judge-candidate-queries-openai}/judged_candidates.jsonl + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + time_min: ${cpu_time_min} + dependencies: + - judge_candidate_queries_openai + + preprocess_filter_easy_candidates: + run_name: ${execution_id} + input_file: ${directories.step-8-reconcile-llm-judges}/reconciled_queries.jsonl + output_file: ${directories.step-10-preprocess-filter-easy-candidates}/input_openai_format.jsonl + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + prompt_file: ${difficulty_filter_prompt_file} + k: ${difficulty_filter_k} + time_min: ${cpu_time_min} + dependencies: + - reconcile_llm_judges + + filter_easy_candidates: + run_name: ${execution_id} + input_file: ${directories.step-8-reconcile-llm-judges}/reconciled_queries.jsonl + preprocessed_input_file: ${directories.step-10-preprocess-filter-easy-candidates}/input_openai_format.jsonl + output_dir: ${directories.step-10-filter-easy-candidates} + summary_file: ${directories.step-10-filter-easy-candidates}/summary.json + model_config: ${model_profiles.omni} + prompt_file: ${difficulty_filter_prompt_file} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + k: ${difficulty_filter_k} + num_chunks: ${difficulty_filter_num_chunks} + temperature: ${difficulty_filter_temperature} + top_p: ${difficulty_filter_top_p} + enable_thinking: ${difficulty_filter_enable_thinking} + min_pass_rate: ${difficulty_filter_min_pass_rate} + max_pass_rate: ${difficulty_filter_max_pass_rate} + time_min: ${gpu_time_min} + dependencies: + - preprocess_filter_easy_candidates + + preprocess_generate_sft_reasoning_traces: + run_name: ${execution_id} + input_file: ${directories.step-10-filter-easy-candidates}/kept_output.jsonl + output_file: ${directories.step-11-preprocess-generate-sft-reasoning-traces}/input_openai_format.jsonl + model_config: ${model_profiles.qwen} + prompt_file: ${sft_trace_generation_prompt_file} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + max_image_dimension: ${max_image_dimension} + k: ${sft_trace_k} + time_min: ${cpu_time_min} + dependencies: + - filter_easy_candidates + + generate_sft_reasoning_traces: + run_name: ${execution_id} + preprocessed_input_file: ${directories.step-11-preprocess-generate-sft-reasoning-traces}/input_openai_format.jsonl + output_file: ${directories.step-11-generate-sft-reasoning-traces}/final_result.jsonl + incorrect_output_file: ${directories.step-11-generate-sft-reasoning-traces}/incorrect_answers.jsonl + output_dir: ${directories.step-11-generate-sft-reasoning-traces} + summary_file: ${directories.step-11-generate-sft-reasoning-traces}/summary.json + model_config: ${model_profiles.qwen} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + num_chunks: ${sft_trace_generation_num_chunks} + temperature: ${sft_trace_generation_temperature} + top_p: ${sft_trace_generation_top_p} + tokens_to_generate: ${sft_trace_generation_tokens_to_generate} + time_min: ${gpu_time_min} + dependencies: + - preprocess_generate_sft_reasoning_traces + + preprocess_filter_sft_reasoning_traces: + run_name: ${execution_id} + input_file: ${directories.step-11-generate-sft-reasoning-traces}/final_result.jsonl + output_file: ${directories.step-12-preprocess-filter-sft-reasoning-traces}/input_openai_format.jsonl + prompt_file: ${sft_trace_judge_prompt_file} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + time_min: ${cpu_time_min} + dependencies: + - generate_sft_reasoning_traces + + filter_sft_reasoning_traces: + run_name: ${execution_id} + input_file: ${directories.step-11-generate-sft-reasoning-traces}/final_result.jsonl + preprocessed_input_file: ${directories.step-12-preprocess-filter-sft-reasoning-traces}/input_openai_format.jsonl + output_dir: ${directories.step-12-filter-sft-reasoning-traces} + summary_file: ${directories.step-12-filter-sft-reasoning-traces}/summary.json + model_config: ${model_profiles.qwen} + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + num_chunks: ${sft_trace_judge_num_chunks} + temperature: ${sft_trace_judge_temperature} + top_p: ${sft_trace_judge_top_p} + tokens_to_generate: ${sft_trace_judge_tokens_to_generate} + time_min: ${gpu_time_min} + dependencies: + - preprocess_filter_sft_reasoning_traces + + visualize_reconciled_hopchain_data: + run_name: ${execution_id} + queries_input_file: ${directories.step-8-reconcile-llm-judges}/final_candidates.jsonl + combinations_input_file: ${directories.step-3-sample-instance-combinations}/instance_combinations.jsonl + output_dir: ${directories.step-9-visualize-reconciled-hopchain-data} + summary_file: ${directories.step-9-visualize-reconciled-hopchain-data}/summary.json + project_root: ${project_root} + cluster_config_dir: ${cluster_config_dir} + rows_per_file: 100 + sample_count: 1000 + sample_seed: 42 + image_max_dimension: 1024 + title: HopChain LLM Judge Reconciliation Review + time_min: ${cpu_time_min} + dependencies: + - reconcile_llm_judges diff --git a/nvflow/utils/__init__.py b/nvflow/utils/__init__.py index 1cac4e7..f97e400 100644 --- a/nvflow/utils/__init__.py +++ b/nvflow/utils/__init__.py @@ -12,8 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Utility functions and helpers.""" +"""Utility functions and helpers. + +Only lightweight, dependency-free helpers are re-exported here. In +particular, the JSONL helpers in :mod:`nvflow.utils.jsonl` import +``orjson`` at module top-level and MUST be imported directly via +``from nvflow.utils.jsonl import ...``. Re-exporting them from this +package would force every consumer of :func:`setup_logger` (including +standalone workers like :mod:`nvflow.lib.rl.create_overlay`, which run +inside container images that do NOT ship ``orjson`` such as the vLLM +server container) to pay the ``orjson`` import cost -- and crash on +``ModuleNotFoundError`` when the dep is absent. + +If you need the JSONL helpers, import them explicitly:: + + from nvflow.utils.jsonl import iter_jsonl, write_jsonl, write_stats_json +""" from nvflow.utils.logging_setup import setup_logger -__all__ = ["setup_logger"] +__all__ = [ + "setup_logger", +] diff --git a/nvflow/utils/jsonl.py b/nvflow/utils/jsonl.py new file mode 100644 index 0000000..ae9dce3 --- /dev/null +++ b/nvflow/utils/jsonl.py @@ -0,0 +1,226 @@ +# 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. +# +"""Shared JSONL read / write helpers for the recipes/finance utilities. + +The same buffered JSONL pattern (orjson decode/encode, ``WRITE_BUFFER_SIZE``, +manual flush, error handling) appeared verbatim in eight files under +``recipes/finance/``. This module is the single source of truth so a fix +or change applies everywhere. + +Three primitives: + +- :func:`iter_jsonl` -- read a JSONL file lazily, with a choice of + malformed-line policy (``skip`` / ``raise`` / ``yield_error``). Empty + lines are always silently skipped. +- :class:`write_jsonl` -- buffered writer context manager. ``write()`` + accepts either a ``dict`` (orjson-encoded) or raw ``bytes`` (written + verbatim). The bytes pass-through is REQUIRED for the + validate_questions pure-row-filter contract: ``apply_validate_filter`` + passes raw SDG bytes through to preserve key ordering and float + formatting from the source file. +- :func:`write_stats_json` -- atomic write of a stats JSON via + ``tmp + os.replace()``. Crash-safe: a partial write leaves only + ``{path}.tmp``, never a truncated ``{path}``. + +Performance constraints honoured to keep adoption byte-identical to the +old in-line code: + +- Buffer flush uses ``b"\\n".join(buffer) + b"\\n"`` (one syscall per flush). +- Buffer threshold is ``>=`` (matches existing behaviour). +- Last flush appends a trailing newline (file always ends in ``\\n``). +""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from pathlib import Path +from types import TracebackType +from typing import Any, Literal, overload + +import orjson + +DEFAULT_BUFFER_SIZE = 1000 +"""Default number of records to buffer before flushing to disk. + +Matches the historical ``WRITE_BUFFER_SIZE`` constant used in the +recipes/finance utilities. Kept centralised so future tuning needs to +happen in only one place. +""" + + +@overload +def iter_jsonl( + path: str | Path, + *, + on_error: Literal["skip", "raise"] = ..., +) -> Iterator[dict[str, Any]]: ... + + +@overload +def iter_jsonl( + path: str | Path, + *, + on_error: Literal["yield_error"], +) -> Iterator[tuple[dict[str, Any] | None, orjson.JSONDecodeError | None, bytes]]: ... + + +def iter_jsonl( + path: str | Path, + *, + on_error: Literal["skip", "raise", "yield_error"] = "skip", +) -> Iterator[Any]: + """Iterate JSONL records lazily. + + Empty lines are silently skipped in every mode (universal behaviour + today). Trailing whitespace is stripped before parsing. + + Args: + path: Path to the JSONL file (str or :class:`pathlib.Path`). + on_error: How to handle malformed JSON lines. + + - ``"skip"`` (default): silently skip the line. + - ``"raise"``: raise the underlying :class:`orjson.JSONDecodeError`. + - ``"yield_error"``: yield ``(None, exc, raw_line)`` for + malformed lines and ``(row, None, raw_line)`` for valid + ones. ``raw_line`` is the stripped source bytes -- callers + that emit audit records (e.g., + ``regex_prefilter_questions``) include a truncated decoded + copy in the dropped stream so an operator can inspect the + offending source line without re-opening the input file. + + Yields: + For ``skip`` / ``raise``: ``dict`` per valid line. + For ``yield_error``: + ``tuple[dict | None, orjson.JSONDecodeError | None, bytes]``. + """ + with open(path, "rb") as reader: + for raw in reader: + line = raw.strip() + if not line: + continue + try: + row = orjson.loads(line) + except orjson.JSONDecodeError as exc: + if on_error == "skip": + continue + if on_error == "raise": + raise + yield (None, exc, line) + continue + if on_error == "yield_error": + yield (row, None, line) + else: + yield row + + +class write_jsonl: # noqa: N801 -- callable-style API: pairs with iter_jsonl(path) function + """Buffered JSONL writer context manager. + + Named in lowercase intentionally so the call site reads as a + function-style helper paired with :func:`iter_jsonl`:: + + with write_jsonl(out_path) as out: + for row in iter_jsonl(in_path): + out.write(transform(row)) + + The lowercase naming violates :pep:`8` ``N801`` (CapWords for class + names); the rule is silenced via ``noqa`` because the readability + win at every call site outweighs the convention deviation, and the + pair ``iter_jsonl`` / ``write_jsonl`` is the established symmetry. + + Accepts ``dict`` (orjson-encoded with no options) or raw ``bytes`` + (written verbatim, used for byte-preserving pass-through). Bytes + must NOT contain a trailing newline -- the writer adds the line + terminator on flush, matching the historical + ``b"\\n".join(buffer) + b"\\n"`` pattern. + + Buffer flushes happen when ``len(buffer) >= buffer_size`` (matches + historical ``>=`` semantics) and once more on context exit if the + buffer is non-empty. The final flush appends a trailing newline so + the file ALWAYS ends in ``\\n`` -- matches historical behaviour and + ensures downstream tools that split on newlines see the last record. + """ + + def __init__(self, path: str | Path, *, buffer_size: int = DEFAULT_BUFFER_SIZE) -> None: + if buffer_size < 1: + raise ValueError(f"buffer_size must be >= 1 (got {buffer_size})") + self._path = path + self._buffer_size = buffer_size + self._buffer: list[bytes] = [] + self._fp: Any = None # opened in __enter__ + + def __enter__(self) -> write_jsonl: + # Open lazily on enter so the user can construct the writer + # outside a try/except without leaking file handles. + self._fp = open(self._path, "wb") + return self + + def write(self, row: dict[str, Any] | bytes) -> None: + """Append a record to the write buffer. + + ``dict`` rows are encoded via :func:`orjson.dumps` with no options + (no indenting, no key sort -- matches historical behaviour). + ``bytes`` rows are appended verbatim; they MUST be a single JSON + line WITHOUT a trailing newline (the writer adds line breaks on + flush via the join pattern). + """ + if isinstance(row, bytes): + encoded = row + else: + encoded = orjson.dumps(row) + self._buffer.append(encoded) + if len(self._buffer) >= self._buffer_size: + self._flush() + + def _flush(self) -> None: + if not self._buffer: + return + # Single syscall per flush -- matches historical performance. + self._fp.write(b"\n".join(self._buffer) + b"\n") + self._buffer.clear() + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + try: + self._flush() + finally: + self._fp.close() + self._fp = None + + +def write_stats_json(path: str | Path, fields: dict[str, Any]) -> None: + """Atomically write a stats JSON to ``path``. + + Encodes ``fields`` with :data:`orjson.OPT_INDENT_2` (matches the + historical pretty-printed stats files), writes to ``{path}.tmp``, + then renames to ``{path}`` via :func:`os.replace`. The tmp file + lives in the same directory as the target, so the rename is atomic + on POSIX (same-mount requirement satisfied). + + Crash semantics: a process killed mid-write leaves ``{path}.tmp`` + behind but never a truncated ``{path}``. Downstream tools that + cache ``{path}.exists()`` as "prior run completed" stay correct. + """ + target = Path(path) + tmp = target.with_name(target.name + ".tmp") + encoded = orjson.dumps(fields, option=orjson.OPT_INDENT_2) + with open(tmp, "wb") as f: + f.write(encoded) + os.replace(tmp, target) diff --git a/pyproject.toml b/pyproject.toml index eadd667..31761ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,16 +6,15 @@ authors = [ {name = "Your Team", email = "team@example.com"} ] readme = "README.md" -requires-python = ">=3.12" +requires-python = ">=3.12,<3.14" license = {text = "Apache-2.0"} dependencies = [ - # NeMo-Skills for cluster job submission (brings ~200+ dependencies) - # This is the core framework we orchestrate - list it first for clarity - # Pinned to specific commit for reproducibility (nemo-skills doesn't use version tags) - # Pinned on: 2026-05-05 | Commit: 0229040 (consistent with nemo-skills:0229040 container) - # To update: Find commit with containers in cluster_configs/example-slurm.yaml - "nemo-skills @ git+https://github.com/NVIDIA/NeMo-Skills.git@022904023ad7a83a87662a313cf72e7df5891d55", + # Core framework we orchestrate (brings ~200+ deps). Pinned to a commit + # since nemo-skills has no version tags. + # Pinned 2026-08-02 | e06c9b90 (paired with nemo-rl v0.7.0; NVFlow v1.1.2). + # Must match NEMO_SKILLS_COMMIT in dockerfiles/Dockerfile.nemo-skills. + "nemo-skills @ git+https://github.com/NVIDIA/NeMo-Skills.git@e06c9b900177be3f60d6a3f99135bb5de9af9bed", # Configuration & Workflow "omegaconf>=2.3.0", # YAML config loading with variable interpolation @@ -27,10 +26,10 @@ dependencies = [ # Data handling "jsonlines>=4.0.0", # JSONL file reading/writing - # urllib3>=2.6.3 blocked by torchx<1.27 constraint (via nemo-skills β†’ nemo-run β†’ torchx) - - # Note: Heavy dependencies (torch, transformers, etc.) come from nemo-skills - # Actual compute happens on Slurm cluster inside containers + # Declared directly ONLY so the CPU-index source below can route it (keeps + # CUDA out of the x86_64 client). Floor set for CVE remediation; otherwise + # tracks nemo-skills' torch. + "torch>=2.13.0", ] [project.optional-dependencies] @@ -47,6 +46,7 @@ dev = [ "ruff>=0.1.0", "mypy>=1.7.0", "types-PyYAML", + "types-requests", ] [project.scripts] @@ -104,14 +104,49 @@ override-dependencies = [ # Mirrors the override added in nemo-skills/pyproject.toml (PR #1433), # but [tool.uv] from a dep is ignored β€” overrides must be at the top-level project. "httpx[http2]>=0.28.1", - "urllib3>=2.6.3", # Force minimum versions of transitive deps for security/maintenance updates - "cryptography>=47.0.0", - "Pillow>=12.2.0", + "cryptography>=48.0.1", + "Pillow>=12.3.0", "Pygments>=2.20.0", - "GitPython>=3.1.49", + "GitPython>=3.1.52", # 3.1.50 still resolves vulnerable; fixes land in 3.1.51/3.1.52 (Trivy HIGH) + # Nspec CVE remediation (2026-07-09) β€” long-term fixes + # ray[default]>=2.54.0 was previously overridden (removed in 4e1b7ec4 during rebase cleanup). + # Restoring at >=2.56.0 to cover both the functional fix (entrypoint_label_selector added + # in 2.54.0) and the Nspec CVE recommendation. Without this, nemo-run->torchx pins ray 2.53.0. + "ray[default]>=2.56.0", + "urllib3>=2.7.0", + "transformers>=5.13.0", + # CVE floors ported from !188 (litellm capped to 1.84.x to avoid the 1.91.x/aiohttp-4.0.0a1 conflict) + # lxml>=6.1.0 clears High CVE-2026-41066 (info disclosure / local file read) + "starlette>=1.3.1", + "litellm>=1.84.0,<1.85", + "lxml>=6.1.0", + "gradio>=6.20.0", +] + +# NSpect/Trivy HIGH CVE floors (2026-07-21) β€” transitive-only, raise lower bound. +constraint-dependencies = [ + "pyjwt>=2.12.0", + "python-multipart>=0.0.27", + "mcp>=1.28.1", + "msgpack>=1.2.1", + "nltk>=3.10.0", + "aiohttp>=3.13.3", + "pyarrow>=23.0.1", + "pyasn1>=0.6.4", # CVE-2026-59886: exact big-integer exponentiation DoS + "soupsieve>=2.8.4", ] +# x86_64 torch from CPU wheels (arm64 PyPI torch is already CPU-only), keeping +# CUDA runtime out of the amd64 client. explicit = only torch uses this index. +[[tool.uv.index]] +name = "pytorch-cpu" +url = "https://download.pytorch.org/whl/cpu" +explicit = true + +[tool.uv.sources] +torch = [{ index = "pytorch-cpu", marker = "platform_machine == 'x86_64'" }] + [dependency-groups] dev = [ "pudb>=2025.1.3", diff --git a/scripts/_dump_rollout_fixtures.py b/scripts/_dump_rollout_fixtures.py new file mode 100644 index 0000000..cf5d818 --- /dev/null +++ b/scripts/_dump_rollout_fixtures.py @@ -0,0 +1,153 @@ +# 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. +# +"""One-shot script to dump current renderer outputs into the test fixtures dir. + +Used to capture the initial baseline for the bash-renderer snapshot tests +in tests/test_rollout.py. Re-run after any *intentional* change to the +renderers to refresh fixtures, then audit the diff before committing. + +Usage: uv run python3 scripts/_dump_rollout_fixtures.py +""" + +from __future__ import annotations + +from pathlib import Path + +from nvflow.lib.rl.rollout import ( + _build_client_cmd, + _build_merge_cmd, + build_aggregate_cmd, + build_filter_cmd, +) + +FIXTURES = Path(__file__).parent.parent / "tests" / "fixtures" / "rollout" +FIXTURES.mkdir(parents=True, exist_ok=True) + + +def _client_cmd(**overrides: object) -> str: + """Render _build_client_cmd with a known-good baseline + overrides.""" + base: dict[str, object] = { + "output_dir": "/out/rollout", + "gym_path": "/opt/Gym", + "model_path": "/hf_models/Qwen/Qwen3-30B-A3B", + "agent_name": "finance_agent", + "input_data": "/data/train.jsonl", + "output_file": "/out/rollout/rs0/chunk_0.jsonl", + "done_file": "/out/rollout/rs0/chunk_0.jsonl.done", + "config_paths": "vllm.yaml,env.yaml,overlay.yaml", + "num_parallel": 512, + "job_label": "rs0_chunk0", + "policy_vllm_url": "http://policy:8000/v1", + "judge_vllm_url": "http://judge:8001/v1", + "judge_ng_run_overrides": ( + ' "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \\\n' + ' "+judge_model.responses_api_models.vllm_model.base_url=http://judge:8001/v1" \\\n' + ), + "max_num_samples": 0, + "chunk_id": 0, + "num_chunks": 8, + "responses_create_params": {"max_output_tokens": 32768, "temperature": 1.0}, + } + base.update(overrides) + return _build_client_cmd(**base) # type: ignore[arg-type] + + +# 1. Dual-server Qwen3-style: policy URL + judge URL + multi-chunk +(FIXTURES / "client_cmd_dual_server.txt").write_text(_client_cmd()) + +# 2. Policy-only (policy_as_judge or no judge): empty judge URL + overrides +(FIXTURES / "client_cmd_policy_only.txt").write_text( + _client_cmd(judge_vllm_url="", judge_ng_run_overrides="") +) + +# 3. Single-chunk path: chunk slicing branch must NOT emit +(FIXTURES / "client_cmd_no_chunk.txt").write_text(_client_cmd(num_chunks=1)) + +# 4. With max_num_samples cap (truncated input) +(FIXTURES / "client_cmd_max_samples.txt").write_text(_client_cmd(max_num_samples=10000)) + +# 5. Empty responses_create_params: no extra +responses_create_params.* lines +(FIXTURES / "client_cmd_no_rcp.txt").write_text(_client_cmd(responses_create_params={})) + + +# --- _build_merge_cmd --- +def _merge_cmd(**overrides: object) -> str: + base: dict[str, object] = { + "gym_path": "/opt/Gym", + "merged_file": "/out/rollout/output-rs0.jsonl", + "analysis_dir": "/out/rollout/analysis_rs0", + "seed_label": "rs0", + "num_chunks": 8, + "chunk_file_pattern": "/out/rollout/rs0/chunk_$i.jsonl", + "merged_done_file": "/out/rollout/output-rs0.jsonl.done", + "analyze_module": "nvflow.recipes.finance.utils.rl.analyze_rollouts", + "enrich_module": "nvflow.recipes.finance.utils.rl.enrich_rollouts", + "input_data": "/data/train.jsonl", + } + base.update(overrides) + return _build_merge_cmd(**base) # type: ignore[arg-type] + + +(FIXTURES / "merge_cmd_8chunks.txt").write_text(_merge_cmd()) +(FIXTURES / "merge_cmd_1chunk.txt").write_text(_merge_cmd(num_chunks=1)) + + +# --- build_aggregate_cmd --- +(FIXTURES / "aggregate_cmd_default.txt").write_text( + build_aggregate_cmd( + rollout_dir="/out/rollout", + aggregate_module="nvflow.recipes.finance.utils.rl.aggregate_seeds", + ) +) +(FIXTURES / "aggregate_cmd_custom_filename.txt").write_text( + build_aggregate_cmd( + rollout_dir="/out/rollout", + aggregate_module="nvflow.recipes.finance.utils.rl.aggregate_seeds", + difficulty_filename="custom_difficulty.jsonl", + ) +) + + +# --- build_filter_cmd --- +(FIXTURES / "filter_cmd_minimal.txt").write_text( + build_filter_cmd( + output_dir="/out", + difficulty_dir="/out/rollout", + filter_module="nvflow.recipes.finance.utils.rl.filter_training_data", + train_data="/in/train.jsonl", + validation_data="", + ) +) +(FIXTURES / "filter_cmd_full.txt").write_text( + build_filter_cmd( + output_dir="/out", + difficulty_dir="/out/rollout", + filter_module="nvflow.recipes.finance.utils.rl.filter_training_data", + train_data="/in/train.jsonl", + validation_data="/in/val.jsonl", + min_reward_std=1e-6, + policy_model="/hf_models/Qwen/Qwen3-30B-A3B", + judge_model="/hf_models/openai/gpt-oss-120b", + train_filename="train.jsonl", + val_filename="validation.jsonl", + difficulty_filename="difficulty.jsonl", + report_filename="filter_report.json", + ) +) + + +print("Wrote fixtures:") +for p in sorted(FIXTURES.iterdir()): + print(f" {p.name}: {p.stat().st_size} bytes") diff --git a/scripts/_dump_verify_fixtures.py b/scripts/_dump_verify_fixtures.py new file mode 100644 index 0000000..e4361ea --- /dev/null +++ b/scripts/_dump_verify_fixtures.py @@ -0,0 +1,109 @@ +# 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. +# +"""One-shot script to dump current verify-renderer outputs into fixtures. + +Used to capture the baseline for the bash-renderer snapshot tests in +``tests/test_verify.py``. Re-run after any *intentional* change to +``_build_verify_cmd`` or ``_build_analysis_cmd``, then audit the diff +under ``tests/fixtures/verify/`` carefully before committing. + +Usage: uv run python3 scripts/_dump_verify_fixtures.py +""" + +from __future__ import annotations + +from pathlib import Path + +from nvflow.lib.rl.verify import _build_analysis_cmd, _build_verify_cmd + +FIXTURES = Path(__file__).parent.parent / "tests" / "fixtures" / "verify" +FIXTURES.mkdir(parents=True, exist_ok=True) + + +# Override blocks mirror what build_judge_ng_run_overrides() returns +# for each judge_mode -- keep them in sync so the snapshots reflect +# realistic call-site shapes rather than a stripped placeholder. + +_LOCAL_VLLM_OVERRIDES = ( + ' "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \\\n' + ' "+judge_model.responses_api_models.vllm_model.base_url=http://127.0.0.1:$JUDGE_PORT/v1" \\\n' + ' "+judge_model.responses_api_models.vllm_model.api_key=EMPTY" \\\n' + ' "+judge_model.responses_api_models.vllm_model.model=/hf_models/openai/gpt-oss-120b" \\\n' + ' "+judge_model.responses_api_models.vllm_model.return_token_id_information=false" \\\n' + ' "+judge_model.responses_api_models.vllm_model.uses_reasoning_parser=true" \\\n' + ' "+finance_env.resources_servers.finance_env.judge_model_server.name=judge_model" \\\n' +) + +_OPENAI_OVERRIDES = ( + ' "+judge_model.responses_api_models.openai_model.base_url=https://api.openai.com/v1" \\\n' + ' "+judge_model.responses_api_models.openai_model.api_key_env_var=OPENAI_API_KEY" \\\n' + ' "+judge_model.responses_api_models.openai_model.model=gpt-4o-mini" \\\n' + ' "+finance_env.resources_servers.finance_env.judge_model_server.name=judge_model" \\\n' +) + + +def _verify_cmd(**overrides: object) -> str: + """Render _build_verify_cmd with a known-good baseline + overrides.""" + base: dict[str, object] = { + "output_dir": "/out/verify", + "gym_path": "/opt/Gym", + "input_file": "/in/rollouts/output-rs0.jsonl", + "output_file": "/out/verify/rejudge/output-rs0.jsonl", + "done_file": "/out/verify/rejudge/output-rs0.jsonl.done", + "config_paths": "vllm.yaml,env.yaml,overlay.yaml", + "num_parallel": 8, + "job_label": "rejudge_rs0", + "judge_mode": "local_vllm", + "environment_name": "finance_env", + "judge_ng_run_overrides": _LOCAL_VLLM_OVERRIDES, + } + base.update(overrides) + return _build_verify_cmd(**base) # type: ignore[arg-type] + + +# 1. Local-vLLM judge: full overrides w/ uses_reasoning_parser +(FIXTURES / "verify_cmd_local_judge.txt").write_text(_verify_cmd()) + +# 2. OpenAI-API judge: shape of overrides differs (different keys) +(FIXTURES / "verify_cmd_openai_judge.txt").write_text( + _verify_cmd(judge_mode="openai", judge_ng_run_overrides=_OPENAI_OVERRIDES) +) + + +# --- _build_analysis_cmd --- +def _analysis_cmd(**overrides: object) -> str: + base: dict[str, object] = { + "rejudge_dir": "/out/verify/rejudge", + "gym_path": "/opt/Gym", + "analyze_module": "nvflow.recipes.finance.utils.rl.analyze_rollouts", + "analysis_entries": [ + ("rs0", "/out/verify/rejudge/output-rs0.jsonl"), + ("rs1", "/out/verify/rejudge/output-rs1.jsonl"), + ], + } + base.update(overrides) + return _build_analysis_cmd(**base) # type: ignore[arg-type] + + +(FIXTURES / "analysis_cmd_multi_seed.txt").write_text(_analysis_cmd()) +(FIXTURES / "analysis_cmd_single_seed.txt").write_text( + _analysis_cmd(analysis_entries=[("rs0", "/out/verify/rejudge/output-rs0.jsonl")]) +) +(FIXTURES / "analysis_cmd_empty_entries.txt").write_text(_analysis_cmd(analysis_entries=[])) + + +print("Wrote fixtures:") +for p in sorted(FIXTURES.iterdir()): + print(f" {p.name}: {p.stat().st_size} bytes") diff --git a/scripts/serve_vllm_patched.py b/scripts/serve_vllm_patched.py deleted file mode 100644 index b680e69..0000000 --- a/scripts/serve_vllm_patched.py +++ /dev/null @@ -1,311 +0,0 @@ -# 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. -# -"""Patched vLLM entrypoint for standalone ``vllm serve`` processes. - -Applies runtime workarounds **before** vLLM is imported, then builds and -runs the ``vllm.entrypoints.openai.api_server`` command directly. - -The upstream ``nemo_skills.inference.server.serve_vllm`` hardcodes -``tensor_parallel_size = num_gpus * num_nodes``, which assumes all GPUs -are for TP. This entrypoint replaces that with the formula:: - - tp = (num_gpus * num_nodes) // dp_size - -When ``--data-parallel-size`` is absent, ``dp_size`` defaults to 1 and the -formula gives the same result as the original code. This is model-agnostic -and backward-compatible. - -Tested with vLLM 0.17.1 and 0.18.1. -Other vLLM versions are safe -- patches skip gracefully if the expected -code snippets are not found, and vLLM starts normally unpatched. - -Active workarounds ------------------- -WORKAROUND(vllm-0.17-hermes) - Patches ``vllm/tool_parsers/hermes_tool_parser.py`` on disk so that - ``Hermes2ProToolParser.__init__`` caches tokenizer encode/decode results - behind a ``threading.Lock``. Fixes ``RuntimeError: Already borrowed`` - under concurrent chat-completion requests with tool calling enabled. - Remove when: vLLM ships fix from https://github.com/vllm-project/vllm/pull/35034 - -WORKAROUND(harmony-aarch64) - Pre-downloads tiktoken vocab files for gpt-oss models on aarch64. - The ``openai_harmony`` Rust binary cannot download them at runtime. - Remove when: openai_harmony ships a fixed aarch64 binary. - Tracking: https://github.com/openai/harmony/issues/71 -""" - -from __future__ import annotations - -import argparse -import os -import platform -import subprocess -import urllib.request -from importlib.util import find_spec -from pathlib import Path - -_TAG = "[serve_vllm_patched]" - -# --------------------------------------------------------------------------- -# WORKAROUND(vllm-0.17-hermes) -- hermes tool parser thread-safety -# -# Identical logic to RL/nemo_rl/models/generation/vllm/vllm_worker.py but -# applied on disk before vLLM is imported (standalone ``vllm serve`` has no -# in-process hook). -# -# Remove this entire section when vLLM ships the upstream fix. -# --------------------------------------------------------------------------- - -# --- Exact string snippets to locate and replace in hermes_tool_parser.py --- - -_OLD_IMPORT = "import json\nfrom collections.abc import Sequence" - -_NEW_IMPORT = "import json\nimport threading\nfrom collections.abc import Sequence" - -_OLD_CLASS_LINE = "class Hermes2ProToolParser(ToolParser):" - -_NEW_CLASS_LINE = ( - "class Hermes2ProToolParser(ToolParser):\n" - " _tokenizer_lock = threading.Lock()\n" - " _tokenizer_cache = {}" -) - -_OLD_INIT = ( - " self.tool_call_start_token_ids = self.model_tokenizer.encode(\n" - " self.tool_call_start_token, add_special_tokens=False\n" - " )\n" - " self.tool_call_end_token_ids = self.model_tokenizer.encode(\n" - " self.tool_call_end_token, add_special_tokens=False\n" - " )\n" - "\n" - " self.tool_call_start_token_array = [\n" - " self.model_tokenizer.decode([token_id])\n" - " for token_id in self.tool_call_start_token_ids\n" - " ]\n" - "\n" - " self.tool_call_end_token_array = [\n" - " self.model_tokenizer.decode([token_id])\n" - " for token_id in self.tool_call_end_token_ids\n" - " ]" -) - -_NEW_INIT = ( - " _tid = id(self.model_tokenizer)\n" - " if _tid in Hermes2ProToolParser._tokenizer_cache:\n" - " _cached = Hermes2ProToolParser._tokenizer_cache[_tid]\n" - " self.tool_call_start_token_ids = _cached['start_ids']\n" - " self.tool_call_end_token_ids = _cached['end_ids']\n" - " self.tool_call_start_token_array = _cached['start_array']\n" - " self.tool_call_end_token_array = _cached['end_array']\n" - " else:\n" - " with Hermes2ProToolParser._tokenizer_lock:\n" - " if _tid in Hermes2ProToolParser._tokenizer_cache:\n" - " _cached = Hermes2ProToolParser._tokenizer_cache[_tid]\n" - " self.tool_call_start_token_ids = _cached['start_ids']\n" - " self.tool_call_end_token_ids = _cached['end_ids']\n" - " self.tool_call_start_token_array = _cached['start_array']\n" - " self.tool_call_end_token_array = _cached['end_array']\n" - " else:\n" - " self.tool_call_start_token_ids = self.model_tokenizer.encode(\n" - " self.tool_call_start_token, add_special_tokens=False\n" - " )\n" - " self.tool_call_end_token_ids = self.model_tokenizer.encode(\n" - " self.tool_call_end_token, add_special_tokens=False\n" - " )\n" - " self.tool_call_start_token_array = [\n" - " self.model_tokenizer.decode([token_id])\n" - " for token_id in self.tool_call_start_token_ids\n" - " ]\n" - " self.tool_call_end_token_array = [\n" - " self.model_tokenizer.decode([token_id])\n" - " for token_id in self.tool_call_end_token_ids\n" - " ]\n" - " Hermes2ProToolParser._tokenizer_cache[_tid] = {\n" - " 'start_ids': self.tool_call_start_token_ids,\n" - " 'end_ids': self.tool_call_end_token_ids,\n" - " 'start_array': self.tool_call_start_token_array,\n" - " 'end_array': self.tool_call_end_token_array,\n" - " }" -) - - -def _patch_hermes_tool_parser() -> None: # WORKAROUND(vllm-0.17-hermes) - """Patch hermes_tool_parser.py on disk before vLLM imports it.""" - spec = find_spec("vllm") - if spec is None or not spec.submodule_search_locations: - print(f"{_TAG} vLLM not found -- skipping hermes patch.") - return - - base_dir = next(iter(spec.submodule_search_locations)) - target = os.path.join(base_dir, "tool_parsers", "hermes_tool_parser.py") - - if not os.path.exists(target): - print(f"{_TAG} {target} not found -- skipping hermes patch.") - return - - with open(target) as f: - content = f.read() - - if "_tokenizer_cache" in content: - print(f"{_TAG} Hermes patch already applied.") - return - - if _OLD_INIT not in content: - print(f"{_TAG} WARNING: Expected code snippet not found in {target}.") - print(f"{_TAG} The vLLM version may have changed -- skipping hermes patch.") - return - - content = content.replace(_OLD_IMPORT, _NEW_IMPORT, 1) - content = content.replace(_OLD_CLASS_LINE, _NEW_CLASS_LINE, 1) - content = content.replace(_OLD_INIT, _NEW_INIT, 1) - - with open(target, "w") as f: - f.write(content) - - print(f"{_TAG} Successfully patched {target} for thread-safety.") - - -# --------------------------------------------------------------------------- -# WORKAROUND(harmony-aarch64) -- tiktoken vocab download for gpt-oss on ARM -# -# Remove this entire section when openai_harmony ships a fixed aarch64 binary. -# --------------------------------------------------------------------------- - -_TIKTOKEN_FILES = { - "o200k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken", - "cl100k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken", -} - - -def _ensure_tiktoken_cache() -> None: # WORKAROUND(harmony-aarch64) - """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) - - for filename, url in _TIKTOKEN_FILES.items(): - dest = cache_dir / filename - if dest.exists() and dest.stat().st_size > 0: - continue - try: - print(f"{_TAG} Downloading {filename} for aarch64 workaround...") - urllib.request.urlretrieve(url, dest) - print(f"{_TAG} {dest.stat().st_size:,} bytes -> {dest}") - except Exception as e: - print(f"{_TAG} WARNING: Failed to download {filename}: {e}") - return - - os.environ["TIKTOKEN_ENCODINGS_BASE"] = str(cache_dir) - os.environ.setdefault("TIKTOKEN_RS_CACHE_DIR", str(cache_dir)) - print(f"{_TAG} TIKTOKEN_ENCODINGS_BASE={cache_dir}") - - -# --------------------------------------------------------------------------- -# TP / DP arithmetic -# --------------------------------------------------------------------------- - - -def _extract_int_flag(args: list[str], flag: str, default: int = 1) -> int: - """Read an integer flag value from *args* without removing it. - - Supports both ``--flag N`` and ``--flag=N`` forms. - """ - for i, a in enumerate(args): - if a == flag and i + 1 < len(args): - return int(args[i + 1]) - if a.startswith(f"{flag}="): - return int(a.split("=", 1)[1]) - return default - - -def _has_flag(args: list[str], flag: str) -> bool: - """Return True if *flag* is already present in *args*.""" - return any(a == flag or a.startswith(f"{flag}=") for a in args) - - -# --------------------------------------------------------------------------- -# Entrypoint -# --------------------------------------------------------------------------- - - -def main(): - os.environ["TOKENIZERS_PARALLELISM"] = "false" - - _patch_hermes_tool_parser() - _ensure_tiktoken_cache() - - parser = argparse.ArgumentParser( - description="Patched vLLM server entrypoint with DP-aware TP calculation", - ) - parser.add_argument("--model", required=True, help="Model path or HF name") - parser.add_argument("--num_gpus", type=int, required=True) - parser.add_argument("--num_nodes", type=int, default=1) - parser.add_argument("--port", type=int, default=5000, help="Server port") - parser.add_argument("--no_verbose", action="store_true", help="Suppress request logs") - args, unknown = parser.parse_known_args() - - dp_size = _extract_int_flag(unknown, "--data-parallel-size") - total_gpus = args.num_gpus * args.num_nodes - tp_size = total_gpus // dp_size - - # Multi-node DP: vLLM must know how many DP ranks fit on the local - # (master) node, otherwise it tries to place all ranks locally. - dp_size_local = args.num_gpus // tp_size - if dp_size > 1 and not _has_flag(unknown, "--data-parallel-size-local"): - unknown.extend(["--data-parallel-size-local", str(dp_size_local)]) - - print(f"{_TAG} Deploying model {args.model}") - print( - f"{_TAG} GPUs: {total_gpus} total" - f" (num_gpus={args.num_gpus} x num_nodes={args.num_nodes})" - f" -> TP={tp_size}, DP={dp_size}, DP_local={dp_size_local}" - ) - - cmd_list = [ - "python3", - "-m", - "vllm.entrypoints.openai.api_server", - f"--model={args.model}", - f"--served-model-name={args.model}", - "--trust-remote-code", - "--host=0.0.0.0", - f"--port={args.port}", - f"--tensor-parallel-size={tp_size}", - ] - if args.no_verbose: - cmd_list.extend(["--disable-log-requests", "--disable-log-stats"]) - cmd_list.extend(unknown) - - print(f"{_TAG} Starting OpenAI Server") - print(f"{_TAG} cmd: {' '.join(cmd_list)}") - subprocess.run(cmd_list, check=True) - - -if __name__ == "__main__": - main() diff --git a/scripts/view_traces.py b/scripts/view_traces.py new file mode 100644 index 0000000..e022196 --- /dev/null +++ b/scripts/view_traces.py @@ -0,0 +1,615 @@ +#!/usr/bin/env python3 +# 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. +# +"""Lightweight NeMo-Gym rollout-trace viewer (stdlib only, no Gradio). + +Serves a tiny web UI to spot-check rollout traces ONE record at a time. The +server reads only the requested record (seek by line with a lazy offset cache), +so it handles multi-GB ``output-rs*.jsonl`` files without loading them. + +Design inspired by the (removed) ``nemo_gym/dataset_viewer.py`` but with zero +external dependencies -- pure Python stdlib. + +Usage: + uv run python scripts/view_traces.py [--root ] [--port 8800] + +``--root`` defaults to ``$NVFLOW_TRACE_ROOT`` if set, otherwise the current +directory. Then open the forwarded http://localhost: in your +browser. In Cursor / VS Code Remote the port is auto-forwarded over SSH -- just +click the "open in browser" notification (or use the Ports panel). + +Rendering: + * Rollout records (have ``responses_create_params`` + ``response``) are + rendered as a conversation: prompt -> reasoning -> tool calls -> tool + outputs -> final answer, plus a verdict header (reward / judge / expected). + * Any other record falls back to pretty-printed JSON. +""" + +from __future__ import annotations + +import argparse +import json +import os +import random +import threading +import urllib.parse +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +MAX_CONTENT_CHARS = 20000 # per-turn display cap to keep the page responsive +MAX_FILES = 1000 +SCAN_MAXDEPTH = 8 +# Heavy / non-trace dirs to skip while scanning for *.jsonl (keeps the scan fast +# even when rooted at the whole grpo workflow dir -- e.g. the 14 GB SEC cache). +PRUNE_DIRS = { + ".venv", + "venv", + ".git", + "node_modules", + "__pycache__", + ".mypy_cache", + ".ruff_cache", + ".pytest_cache", + "cache", + "logs", + "filings", + "filings_metadata", +} +# jsonl filename substrings that are pipeline input artifacts (no traces) -- hidden +# from the dropdown to reduce noise. +SKIP_FILE_SUBSTRINGS = ("materialized_inputs", "chunk_input") + +# Point this at a single workflow output dir, not at the parent of the SEC dump +# tree -- scanning tens of thousands of filings makes the directory listing crawl. +DEFAULT_ROOT = os.environ.get("NVFLOW_TRACE_ROOT") or "." + + +# --------------------------------------------------------------------------- +# Record parsing (schema-aware, dict-based -- no openai/pydantic deps) +# --------------------------------------------------------------------------- +def _content_to_text(content) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for c in content: + if isinstance(c, dict): + parts.append(c.get("text") or c.get("output") or json.dumps(c, indent=2)) + else: + parts.append(str(c)) + return "\n".join(parts) + if content is None: + return "" + return json.dumps(content, indent=2) + + +def _pretty_json_str(s): + try: + return json.dumps(json.loads(s), indent=2) + except Exception: + return s if isinstance(s, str) else json.dumps(s, indent=2) + + +def parse_item(m: dict): + """Convert one input/output item into {kind, title, content}.""" + if not m.get("type") and m.get("role"): + m = {**m, "type": "message"} + t = m.get("type") + if t == "message": + role = m.get("role", "assistant") + return {"kind": role, "title": "", "content": _content_to_text(m.get("content", ""))} + if t == "function_call": + name = m.get("name", "?") + return { + "kind": "tool_call", + "title": name, + "content": _pretty_json_str(m.get("arguments", "{}")), + } + if t == "function_call_output": + return { + "kind": "tool_output", + "title": "", + "content": _pretty_json_str(m.get("output", "")), + } + if t == "reasoning": + txt = "\n".join(s.get("text", "") for s in (m.get("summary") or []) if isinstance(s, dict)) + return { + "kind": "reasoning", + "title": "", + "content": txt or _content_to_text(m.get("content", "")), + } + return {"kind": "other", "title": t or "item", "content": json.dumps(m, indent=2)} + + +HEADER_FIELDS = ( + "question", + "problem", + "expected_answer", + "question_type", + "reward", + "judge_rating", + "judge_text", + "uuid", + "current_date", +) + + +def parse_record(rec: dict) -> dict: + """Return {schema, header, turns} for rollout records, else generic fallback.""" + header = {} + for k in HEADER_FIELDS: + if k in rec and rec[k] not in (None, ""): + header[k] = rec[k] + dp = rec.get("difficulty_profile") + if isinstance(dp, dict) and "avg_reward" in dp: + header["difficulty_profile.avg_reward"] = dp["avg_reward"] + + rcp = rec.get("responses_create_params") + resp = rec.get("response") + if isinstance(rcp, dict) and isinstance(resp, dict): + raw_inp = rcp.get("input") + if isinstance(raw_inp, str): + inp: list = [{"role": "user", "content": raw_inp}] + elif isinstance(raw_inp, list): + inp = raw_inp + else: + inp = [] + raw_out = resp.get("output") + out: list = raw_out if isinstance(raw_out, list) else [] + turns = [] + turn, step = 0, 0 + for m in inp + out: + if not isinstance(m, dict): + continue + if m.get("role") == "user": + turn += 1 + step = 0 + if m.get("type") == "function_call": + step += 1 + ti = parse_item(m) + content = ti["content"] or "" + if len(content) > MAX_CONTENT_CHARS: + ti["content"] = ( + content[:MAX_CONTENT_CHARS] + + f"\n\n... [truncated {len(content) - MAX_CONTENT_CHARS} chars]" + ) + ti["turn"], ti["step"] = turn, step + turns.append(ti) + return {"schema": "rollout", "header": header, "turns": turns} + + return {"schema": "generic", "header": header, "raw": rec} + + +# --------------------------------------------------------------------------- +# JSONL reader: read one record by index via a lazy byte-offset cache. +# --------------------------------------------------------------------------- +class JsonlReader: + """Reads a single record by index without loading the whole file. + + Maintains a lazily-grown byte-offset index so repeated/sequential access is + cheap. Thread-safe: a lock guards the shared offset cache (the HTTP server is + threaded). + """ + + def __init__(self, path: Path): + self.path = path + self.offsets: list[int] = [0] # offsets[i] = byte offset of line i + self.eof = False + self.count: int | None = None # filled lazily (Random / total) + self._lock = threading.Lock() + + def _extend_to(self, index: int): + """Grow the offset cache through line `index`. Caller must hold the lock.""" + if self.eof or index < len(self.offsets): + return + with open(self.path, "rb") as f: + f.seek(self.offsets[-1]) + i = len(self.offsets) - 1 + while i <= index: + line = f.readline() + if not line: + self.eof = True + self.count = len(self.offsets) - 1 + break + i += 1 + self.offsets.append(f.tell()) + + def get(self, index: int): + if index < 0: + return None + with self._lock: + self._extend_to(index) + if index >= len(self.offsets) - 1 and self.eof: + return None + offset = self.offsets[index] + # File I/O + parse outside the lock (independent of shared state). + try: + with open(self.path, "rb") as f: + f.seek(offset) + line = f.readline() + except OSError: + return None + if not line: + return None + try: + return json.loads(line) + except Exception: + return { + "_parse_error": True, + "raw_line": line.decode("utf-8", "replace")[:MAX_CONTENT_CHARS], + } + + def total(self) -> int: + """Record count via a single chunked byte scan (cached). Used by Random.""" + with self._lock: + if self.count is not None: + return self.count + n, last = 0, b"" + with open(self.path, "rb") as f: + while True: + chunk = f.read(1 << 20) + if not chunk: + break + n += chunk.count(b"\n") + last = chunk[-1:] + if last and last != b"\n": # final line without trailing newline + n += 1 + self.count = n + return n + + +# --------------------------------------------------------------------------- +# HTTP server +# --------------------------------------------------------------------------- +def list_jsonl_files(root: Path): + files = [] + root = root.resolve() + for dirpath, dirnames, filenames in os.walk(root): + rel = Path(dirpath).relative_to(root) + if len(rel.parts) > SCAN_MAXDEPTH: + dirnames[:] = [] + continue + dirnames[:] = [d for d in dirnames if d not in PRUNE_DIRS and not d.startswith(".")] + for fn in filenames: + if fn.endswith(".jsonl") and not any(s in fn for s in SKIP_FILE_SUBSTRINGS): + files.append(str((Path(dirpath) / fn).relative_to(root))) + if len(files) >= MAX_FILES: + return sorted(files) + return sorted(files) + + +def make_handler(root: Path): + root = root.resolve() + readers: dict[str, JsonlReader] = {} + readers_lock = threading.Lock() + + def reader_for(rel: str): + # Resolve safely within root (no path traversal). + target = (root / rel).resolve() + try: + target.relative_to(root) + except ValueError: + return None + if not target.is_file(): + return None + with readers_lock: + if rel not in readers: + readers[rel] = JsonlReader(target) + return readers[rel] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): # quiet + pass + + def _send_json(self, obj, code=200): + body = json.dumps(obj).encode("utf-8") + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Cache-Control", "no-store") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _send_html(self, text): + body = text.encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Cache-Control", "no-store") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): # noqa: N802 (BaseHTTPRequestHandler requires this name) + try: + self._route() + except Exception as e: # never leave the client hanging on a bare 500 + try: + self._send_json({"error": f"{type(e).__name__}: {e}"}, 500) + except Exception: + pass + + def _route(self): + parsed = urllib.parse.urlparse(self.path) + q = urllib.parse.parse_qs(parsed.query) + path = parsed.path + if path == "/" or path == "/index.html": + self._send_html(PAGE) + return + if path == "/api/files": + self._send_json({"root": str(root), "files": list_jsonl_files(root)}) + return + if path in ("/api/record", "/api/random"): + rel = (q.get("file") or [""])[0] + rd = reader_for(rel) + if rd is None: + self._send_json({"error": f"file not found under root: {rel}"}, 404) + return + if path == "/api/random": + total = rd.total() + index = random.randint(0, max(0, total - 1)) if total else 0 + else: + try: + index = int((q.get("index") or ["0"])[0]) + except ValueError: + index = 0 + rec = rd.get(index) + if rec is None: + self._send_json( + {"error": "no record at index", "index": index, "known_total": rd.count}, + 404, + ) + return + parsed_rec = parse_record(rec) + self._send_json( + {"index": index, "known_total": rd.count, "parsed": parsed_rec, "raw": rec} + ) + return + self._send_json({"error": "not found"}, 404) + + return Handler + + +# --------------------------------------------------------------------------- +# Frontend (single self-contained page) +# --------------------------------------------------------------------------- +PAGE = r""" +Rollout Trace Viewer + + +
+ + + + + + + + + + + +
+
+ + +""" + + +def main(): + ap = argparse.ArgumentParser( + description="Lightweight NeMo-Gym rollout-trace viewer (stdlib only)." + ) + ap.add_argument( + "--root", + default=DEFAULT_ROOT, + help=( + "Directory scanned for *.jsonl files (file dropdown). " + "Defaults to $NVFLOW_TRACE_ROOT, or the current directory." + ), + ) + ap.add_argument( + "--host", + default="127.0.0.1", + help="Bind host (default 127.0.0.1; use SSH/Cursor port-forward).", + ) + ap.add_argument("--port", type=int, default=8800, help="Bind port (default 8800).") + args = ap.parse_args() + + root = Path(args.root).expanduser() + if not root.is_dir(): + raise SystemExit(f"--root is not a directory: {root}") + root = root.resolve() + + handler = make_handler(root) + try: + httpd = ThreadingHTTPServer((args.host, args.port), handler) + except OSError as e: + raise SystemExit( + f"Could not bind {args.host}:{args.port} ({e}).\n" + "A viewer may already be running -- stop it, or pass a different --port." + ) from e + print(f"Rollout trace viewer serving {root}", flush=True) + print( + f" http://{args.host}:{args.port} (Cursor/VS Code will auto-forward this port)", + flush=True, + ) + print(" Ctrl+C to stop.", flush=True) + try: + httpd.serve_forever() + except KeyboardInterrupt: + print("\nstopping.") + httpd.shutdown() + + +if __name__ == "__main__": + main() diff --git a/scripts/wandb_consolidate.py b/scripts/wandb_consolidate.py index 3a1ae1e..259b5ff 100644 --- a/scripts/wandb_consolidate.py +++ b/scripts/wandb_consolidate.py @@ -13,37 +13,43 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Consolidate multiple WandB training runs into a single dashboard. +"""Consolidate multiple WandB training runs into one continuous dashboard. -Reads metrics from individual WandB runs (filtered by group) and writes -them to a single consolidated WandB run with continuous step numbers. +Each GRPO/SFT job (including each job of a resume chain) creates a separate WandB +run. This merges them post-hoc into a single run with continuous step numbers. -Each GRPO/SFT training job creates a separate WandB dashboard. This script -merges them post-hoc into one continuous view. +Source selection: + --run-ids ID... explicit runs (precise; best when many experiments share a group) + --group NAME all runs in a group, optionally narrowed by --since TIMESTAMP -Two modes: - fresh (default): Creates a new consolidated WandB run. - --append RUN_ID: Resumes an existing consolidated run and adds new data. - Only logs steps beyond what was previously consolidated. +Failure handling (automatic): runs are merged in chronological order and each is +truncated at its successor's resume step, so a failed/rolled-back job's stale tail +is dropped and the resumed job's data wins. Use --max-step to cap a trailing failed +run (the last run has no successor to bound it). + +Modes: + fresh (default) create a new consolidated run + --append RUN_ID extend an existing consolidated run with new steps only Usage: - # First time -- new consolidated dashboard from all runs in a group + # fresh, from explicit run-ids (--entity required if the project is owned by + # another team, e.g. nvidia; or set $WANDB_ENTITY) uv run python scripts/wandb_consolidate.py \\ - --project finance-grpo \\ - --group grpo-training-finance_sec_search \\ - --name gspo-qwen3-30b-consolidated + --project finance-grpo --entity nvidia \\ + --run-ids abc123 def456 ghi789 --name nano-consolidated - # After more training -- append new data to existing dashboard + # or by group + time window uv run python scripts/wandb_consolidate.py \\ - --project finance-grpo \\ - --group grpo-training-finance_sec_search \\ - --append + --project finance-grpo --group grpo-training-finance_sec_search \\ + --since 2026-06-28T19:00:00Z --name nano-consolidated - # Dry run -- show what would be logged + # append later as the chain grows uv run python scripts/wandb_consolidate.py \\ - --project finance-grpo \\ - --group grpo-training-finance_sec_search \\ - --dry-run + --project finance-grpo --run-ids abc123 def456 ghi789 jkl012 \\ + --append + + # preview only + ... --dry-run Requirements: pip install wandb @@ -69,8 +75,21 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--group", - required=True, - help="WandB group name to filter source runs (e.g., grpo-training-finance_sec_search).", + default=None, + help="WandB group to filter source runs. Required unless --run-ids is given.", + ) + parser.add_argument( + "--run-ids", + nargs="*", + default=None, + help="Explicit source run IDs to consolidate (precise selection; ignores " + "--group/--since). Preferred when multiple experiments share one group.", + ) + parser.add_argument( + "--since", + default=None, + help="Only include --group runs created at/after this ISO8601 UTC timestamp " + "(e.g. 2026-06-28T19:00:00Z).", ) parser.add_argument( "--name", @@ -79,11 +98,11 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--entity", - default=None, + default=os.environ.get("WANDB_ENTITY"), help=( - "WandB entity (team/user). If not set, uses the default from " - "'wandb login'. Required if your default entity differs from " - "the project owner." + "WandB entity (team/user). Defaults to $WANDB_ENTITY, else your " + "'wandb login' default. Required for --run-ids when the project is " + "owned by another entity (e.g. --entity nvidia)." ), ) parser.add_argument( @@ -103,40 +122,115 @@ def parse_args() -> argparse.Namespace: default=["ray/"], help="Skip metrics whose key starts with these prefixes (default: ray/).", ) + parser.add_argument( + "--skip-histograms", + action="store_true", + help="Do not carry over wandb.Histogram metrics (train/*/histogram, " + "validation/*/histogram). By default histograms ARE carried; only " + "artifact-backed tables/plots (e.g. full_result, *_plot_sample) are skipped.", + ) + parser.add_argument( + "--max-step", + type=int, + default=None, + help="Cap consolidation at this step (inclusive). Use to exclude a trailing " + "failed job's rolled-back steps beyond its last good checkpoint (the last " + "run has no successor to bound it automatically).", + ) parser.add_argument( "--output-dir", default=None, help=( - "Directory for WandB local run data. Defaults to " - "{training-logs}/wandb_consolidated/ if not set. " - "Prevents WandB from writing to the source tree." + "Directory for WandB local run data (default: " + "/outputs/wandb_consolidated/). Keeps WandB scratch files out " + "of the source tree." ), ) return parser.parse_args() def fetch_source_runs( - api, project: str, entity: str | None, group: str, exclude_id: str | None = None + api, + project: str, + entity: str | None, + group: str | None, + exclude_id: str | None = None, + run_ids: list[str] | None = None, + since: str | None = None, ): - """Fetch source runs in the given group, excluding the consolidated run.""" + """Fetch source runs to consolidate. + + Selection precedence: + - run_ids: fetch exactly those runs (group/since ignored). Use when several + experiments share one group. + - group: fetch the group, optionally filtered by `since` (created_at >= since). + Excludes the --append target and any prior consolidated runs (CONSOLIDATED_TAG). + + Runs are returned in CHRONOLOGICAL (created_at) order, not by max step, so that + on overlapping steps from a failed-then-resumed job the later run's data wins. + """ path = f"{entity}/{project}" if entity else project - runs = api.runs(path, filters={"group": group}) - if not runs: - print(f"No runs found in project={project}, group={group}") + + if run_ids: + candidates = [] + for rid in run_ids: + try: + candidates.append(api.run(f"{path}/{rid}")) + except Exception as e: # noqa: BLE001 + print( + f" WARN: could not fetch run id {rid}: {e}\n" + f" (looked under entity='{entity}'; pass --entity or set " + f"$WANDB_ENTITY if the project is owned by another entity, e.g. nvidia)" + ) + elif group: + candidates = list(api.runs(path, filters={"group": group})) + else: + print("Provide either --group or --run-ids.") sys.exit(1) run_list = [] - for r in runs: + for r in candidates: if exclude_id and r.id == exclude_id: continue if r.tags and CONSOLIDATED_TAG in r.tags: continue + if since and not run_ids and (getattr(r, "created_at", "") or "") < since: + continue run_list.append(r) - run_list.sort(key=lambda r: r.summary.get("_step", 0)) + if not run_list: + print( + f"No source runs matched (project={project}, group={group}, " + f"run_ids={run_ids}, since={since})." + ) + sys.exit(1) + + # Chronological: later (resume) runs supersede earlier failed ones on overlapping steps. + run_list.sort(key=lambda r: getattr(r, "created_at", "") or "") return run_list +def _to_wandb_histogram(wandb, raw: dict): + """Reconstruct a wandb.Histogram from a scan_history histogram dict. + + scan_history returns histograms as {"values": [counts...], + "packedBins": {"min": m, "size": s, "count": n}, "_type": "histogram"}. + wandb.Histogram(np_histogram=(counts, bin_edges)) needs len(bin_edges) == + len(counts) + 1, so we rebuild edges from the packed (min, size, count). + Returns None if the dict can't be reconstructed. + """ + try: + values = list(raw["values"]) + pb = raw["packedBins"] + mn, size, count = float(pb["min"]), float(pb["size"]), int(pb["count"]) + edges = [mn + i * size for i in range(count + 1)] + if len(edges) != len(values) + 1: + return None + return wandb.Histogram(np_histogram=(values, edges)) + except Exception: # noqa: BLE001 + return None + + def get_max_consolidated_step(api, project: str, entity: str | None, run_id: str) -> int: """Get the highest step already logged in the consolidated run.""" path = f"{entity}/{project}" if entity else project @@ -148,37 +242,92 @@ def get_max_consolidated_step(api, project: str, entity: str | None, run_id: str return -1 -def collect_metrics(runs, skip_prefixes: list[str], min_step: int = -1): - """Collect scalar metrics from source runs, skipping already-consolidated steps.""" - all_rows: list[tuple[int, dict[str, float]]] = [] - +def collect_metrics( + runs, + skip_prefixes: list[str], + min_step: int = -1, + max_step: int | None = None, + include_histograms: bool = True, +): + """Collect scalar (and optionally histogram) metrics from source runs. + + Runs are assumed in chronological order. + + Scalars are kept as-is. WandB histograms come back from scan_history as dicts + ({"_type": "histogram", "values": [...], "packedBins": {min,size,count}}); we + keep that raw dict and reconstruct a wandb.Histogram at log time. Artifact-backed + non-scalars (table-file sample dumps, *_plot_sample images) are always skipped. + + Resume-lineage truncation: each run is valid only up to where the *next* run + resumed (its successor's min step). Steps a job logged beyond its last + carried-forward checkpoint were rolled back, so they are dropped. This handles + failed jobs (empty valid range -> auto-dropped) and partial-success jobs + (keep the checkpointed prefix, drop the rolled-back tail) using only step data. + """ + # Pass 1: read each run's rows + its min step (runs are already chronological). + per_run: list[tuple[int | None, list[tuple[int, dict[str, object]]], object]] = [] for run in runs: print(f" Reading run: {run.name} ({run.id}), state={run.state}") - row_count = 0 - skipped = 0 + rows: list[tuple[int, dict[str, object]]] = [] for row in run.scan_history(): step = row.get("_step") if step is None: continue - if int(step) <= min_step: - skipped += 1 - continue - metrics = {} + step = int(step) + metrics: dict[str, object] = {} for k, v in row.items(): - if k.startswith("_"): - continue - if any(k.startswith(p) for p in skip_prefixes): + if k.startswith("_") or any(k.startswith(p) for p in skip_prefixes): continue if isinstance(v, int | float): metrics[k] = v + elif ( + include_histograms + and isinstance(v, dict) + and v.get("_type") == "histogram" + and v.get("values") is not None + and isinstance(v.get("packedBins"), dict) + ): + # Keep raw dict; reconstructed into wandb.Histogram at log time. + metrics[k] = v if metrics: - all_rows.append((int(step), metrics)) - row_count += 1 - msg = f" {row_count} steps collected" + rows.append((step, metrics)) + # Resume boundary = first *training* step (>0). step 0 is the per-job + # val_at_start artifact (re-logged every job), NOT the resume point, so it + # must be excluded or every run's boundary collapses to 0. + resume_step = min((s for s, _ in rows if s > 0), default=None) + per_run.append((resume_step, rows, run)) + + # Pass 2: bound each run at its successor's resume point + drop already-consolidated. + all_rows: list[tuple[int, dict[str, object]]] = [] + for i, (_, rows, run) in enumerate(per_run): + upper = None # exclusive upper bound = next run's min step + for j in range(i + 1, len(per_run)): + if per_run[j][0] is not None: + upper = per_run[j][0] + break + kept = truncated = skipped = capped = 0 + for step, metrics in rows: + if step <= min_step: + skipped += 1 + continue + if max_step is not None and step > max_step: + capped += 1 + continue + if upper is not None and step >= upper: + truncated += 1 + continue + all_rows.append((step, metrics)) + kept += 1 + msg = f" {run.name}: kept {kept} steps" + if truncated: + msg += f", dropped {truncated} rolled-back (>= successor resume @ {upper})" + if capped: + msg += f", dropped {capped} above --max-step {max_step}" if skipped: - msg += f" ({skipped} already consolidated, skipped)" + msg += f", skipped {skipped} already-consolidated" print(msg) + # Stable sort by step preserves chronological order on ties (later run wins). all_rows.sort(key=lambda x: x[0]) return all_rows @@ -190,14 +339,21 @@ def print_summary(rows: list[tuple[int, dict]], runs, min_step: int): return steps = [r[0] for r in rows] all_keys: set[str] = set() + hist_keys: set[str] = set() for _, metrics in rows: all_keys.update(metrics.keys()) + for k, v in metrics.items(): + if isinstance(v, dict) and v.get("_type") == "histogram": + hist_keys.add(k) print("\nConsolidation summary:") print(f" Source runs: {len(runs)}") if min_step >= 0: print(f" Already consolidated up to step: {min_step}") print(f" New steps: {len(steps)} (min={min(steps)}, max={max(steps)})") - print(f" Unique metrics: {len(all_keys)}") + print( + f" Unique metrics: {len(all_keys)} (scalars: {len(all_keys) - len(hist_keys)}, " + f"histograms: {len(hist_keys)})" + ) print(f" Total data points: {sum(len(m) for _, m in rows)}") @@ -215,12 +371,29 @@ def consolidate(args: argparse.Namespace): exclude_id = args.append print(f"Append mode: consolidated run {args.append}, max step = {min_step}") - print(f"Fetching source runs from project={args.project}, group={args.group}") - runs = fetch_source_runs(api, args.project, args.entity, args.group, exclude_id=exclude_id) + print( + f"Fetching source runs from project={args.project} " + f"(group={args.group}, run_ids={args.run_ids}, since={args.since})" + ) + runs = fetch_source_runs( + api, + args.project, + args.entity, + args.group, + exclude_id=exclude_id, + run_ids=args.run_ids, + since=args.since, + ) print(f"Found {len(runs)} source runs\n") print("Collecting metrics...") - rows = collect_metrics(runs, args.skip_prefixes or [], min_step=min_step) + rows = collect_metrics( + runs, + args.skip_prefixes or [], + min_step=min_step, + max_step=args.max_step, + include_histograms=not args.skip_histograms, + ) print_summary(rows, runs, min_step) if not rows: @@ -241,7 +414,7 @@ def consolidate(args: argparse.Namespace): print(f"\nAppending to existing WandB run: {args.append}") else: init_kwargs["name"] = args.name - init_kwargs["group"] = args.group + init_kwargs["group"] = args.group or args.name init_kwargs["tags"] = [CONSOLIDATED_TAG] print(f"\nCreating new WandB run: {args.name}") @@ -258,12 +431,25 @@ def consolidate(args: argparse.Namespace): print(f"URL: {run.url}") logged = 0 + hist_logged = 0 for step, metrics in rows: - run.log(metrics, step=step) - logged += len(metrics) + log_metrics = {} + for k, v in metrics.items(): + if isinstance(v, dict) and v.get("_type") == "histogram": + hv = _to_wandb_histogram(wandb, v) + if hv is not None: + log_metrics[k] = hv + hist_logged += 1 + else: + log_metrics[k] = v + run.log(log_metrics, step=step) + logged += len(log_metrics) run.finish() - print(f"\nDone. Logged {logged} data points across {len(rows)} steps.") + print( + f"\nDone. Logged {logged} data points across {len(rows)} steps " + f"({hist_logged} histogram points)." + ) print(f"Run ID: {run.id} (use with --append for future updates)") diff --git a/tests/fixtures/rollout/aggregate_cmd_custom_filename.txt b/tests/fixtures/rollout/aggregate_cmd_custom_filename.txt new file mode 100644 index 0000000..7fef1f1 --- /dev/null +++ b/tests/fixtures/rollout/aggregate_cmd_custom_filename.txt @@ -0,0 +1,7 @@ +set -e +echo "Cross-Seed Aggregation (pass@k)" +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.aggregate_seeds \ + "/out/rollout" \ + "/out/rollout/aggregate" \ + --output_filename "custom_difficulty.jsonl" +echo "Done. Results in /out/rollout/aggregate/" diff --git a/tests/fixtures/rollout/aggregate_cmd_default.txt b/tests/fixtures/rollout/aggregate_cmd_default.txt new file mode 100644 index 0000000..2e90380 --- /dev/null +++ b/tests/fixtures/rollout/aggregate_cmd_default.txt @@ -0,0 +1,7 @@ +set -e +echo "Cross-Seed Aggregation (pass@k)" +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.aggregate_seeds \ + "/out/rollout" \ + "/out/rollout/aggregate" \ + --output_filename "difficulty.jsonl" +echo "Done. Results in /out/rollout/aggregate/" diff --git a/tests/fixtures/rollout/client_cmd_dual_server.txt b/tests/fixtures/rollout/client_cmd_dual_server.txt new file mode 100644 index 0000000..9c54ac9 --- /dev/null +++ b/tests/fixtures/rollout/client_cmd_dual_server.txt @@ -0,0 +1,269 @@ +set -e + +OUTPUT_DIR=/out/rollout +GYM_PATH=/opt/Gym +UV_VENV_DIR=/opt/Gym +MODEL_PATH=/hf_models/Qwen/Qwen3-30B-A3B +AGENT_NAME=finance_agent +INPUT_DATA=/data/train.jsonl +OUTPUT_FILE=/out/rollout/rs0/chunk_0.jsonl +DONE_FILE=/out/rollout/rs0/chunk_0.jsonl.done +CONFIG_PATHS=vllm.yaml,env.yaml,overlay.yaml +NUM_PARALLEL=512 +JOB_LABEL=rs0_chunk0 +VLLM_URL="http://policy:8000/v1" +JUDGE_URL="http://judge:8001/v1" +CHUNK_ID=0 +NUM_CHUNKS=8 + +mkdir -p "$OUTPUT_DIR/logs" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +NG_RUN_PID="" + +cleanup() { + local _nvflow_exit=$? + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + # Suppress stderr via ``2>&-`` (close fd) instead of + # ``2>/dev/null`` so the cleanup trap stays quiet even if + # the container's /dev/null disappeared mid-script -- which + # is exactly what happened in the Nemotron-Nano smoke run + # where pyxis tore down the container while bash was still + # in the cleanup path, producing ``/dev/null: No such file + # or directory`` noise on top of the original failure. + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>&- && wait $NG_RUN_PID 2>&- || true + # Best-effort merge of .prev into -async. On success, finalize already + # merged and removed .prev so this block is a no-op. On failure/kill, + # this is a first attempt; the self-heal at next startup is the guarantee. + # Chain with && so .prev is NEVER deleted unless the merge succeeds. + if [ -n "$ASYNC_FILE" ] && [ -f "$ASYNC_FILE.prev" ] && [ "${PREV_MERGED:-0}" -eq 0 ]; then + echo "[Cleanup] Restoring previous results into -async for resume ..." + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.restored" \ + && { [ ! -f "$ASYNC_FILE" ] || cat "$ASYNC_FILE" >> "$ASYNC_FILE.restored"; } \ + && mv -f "$ASYNC_FILE.restored" "$ASYNC_FILE" \ + && rm -f "$ASYNC_FILE.prev" \ + || echo "[Cleanup] WARNING: merge failed β€” self-heal will recover on next start" + fi + if [ $_nvflow_exit -ne 0 ]; then + echo "[nvflow] Client exited with code $_nvflow_exit β€” cancelling job ${SLURM_JOB_ID}" + scancel "${SLURM_JOB_ID}" 2>&- || kill 0 2>&- || true + fi +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Rollout Collection [$JOB_LABEL]" +echo "============================================================" +echo "Model: $MODEL_PATH" +echo "Agent: $AGENT_NAME" +echo "Input data: $INPUT_DATA" +echo "Output file: $OUTPUT_FILE" +echo "Policy URL: $VLLM_URL" +[ -n "$JUDGE_URL" ] && echo "Judge URL: $JUDGE_URL" +echo "============================================================" +if [ -f "$DONE_FILE" ]; then + echo "Chunk already complete (.done exists) β€” skipping." + exit 0 +fi + +echo "" +echo "[Step 1/3] Waiting for vLLM servers ..." +wait_for_server "http://policy:8000/v1/models" "Policy vLLM" $$ 400 /dev/null +wait_for_server "http://judge:8001/v1/models" "Judge vLLM" $$ 400 /dev/null + +CHUNK_INPUT="" +if [ $NUM_CHUNKS -gt 1 ]; then + echo "" + echo "[Step 1a] Extracting chunk slice ..." + TOTAL_LINES=$(wc -l < "$INPUT_DATA") + EFFECTIVE=$TOTAL_LINES + MAX_SAMPLES=0 + if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then + EFFECTIVE=$MAX_SAMPLES + fi + CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) + [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" + head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" + echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" + INPUT_DATA="$CHUNK_INPUT" +fi +ASYNC_FILE="$OUTPUT_FILE-async" +( + # Case A: output exists without .done β†’ restore to -async + if [ -f "$OUTPUT_FILE" ] && [ ! -f "$DONE_FILE" ]; then + echo "[Self-heal] Output file exists without .done β€” restoring to -async ..." + mv -f "$OUTPUT_FILE" "$ASYNC_FILE" + fi + + # Case B: orphaned .prev β†’ merge into -async + if [ -f "$ASYNC_FILE.prev" ]; then + echo "[Self-heal] Found orphaned .prev β€” merging into -async ..." + PREV_LINES=$(wc -l < "$ASYNC_FILE.prev") + ASYNC_LINES=0 + [ -f "$ASYNC_FILE" ] && ASYNC_LINES=$(wc -l < "$ASYNC_FILE") + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.healed" + [ -f "$ASYNC_FILE" ] && cat "$ASYNC_FILE" >> "$ASYNC_FILE.healed" + mv -f "$ASYNC_FILE.healed" "$ASYNC_FILE" && rm -f "$ASYNC_FILE.prev" + MERGED_LINES=$(wc -l < "$ASYNC_FILE") + echo " Recovered $PREV_LINES (prev) + $ASYNC_LINES (async) = $MERGED_LINES total rows" + fi + + # Case C: clean up orphaned temp files from prior crash + rm -f "$ASYNC_FILE.healed" "$ASYNC_FILE.restored" "$ASYNC_FILE.merged" +) || echo "[Self-heal] WARNING: recovery failed β€” continuing with available data" +REMAINING_INPUT="$OUTPUT_DIR/remaining_input_chunk$CHUNK_ID.jsonl" + +if ! PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.resume_filter "$ASYNC_FILE" "$INPUT_DATA" "$REMAINING_INPUT" 0; then + echo "ERROR: resume_filter failed" >&2 + exit 1 +fi + +if [ -f "$ASYNC_FILE" ] && [ ! -s "$REMAINING_INPUT" ]; then + echo "All rows already completed in -async -- finalizing." + cp -f "$ASYNC_FILE" "$OUTPUT_FILE" + touch "$DONE_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_FILE" "$ASYNC_FILE.prev" + echo "Done [$JOB_LABEL]." + exit 0 +fi + +HEAD_SERVER_PORT=$(find_free_port) + +cd "$GYM_PATH" + +echo "" +echo "[Step 2/3] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=$VLLM_URL" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=$MODEL_PATH" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+port_range_low=1024" \ + "+port_range_high=8999" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \ + "+judge_model.responses_api_models.vllm_model.base_url=http://judge:8001/v1" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 3/3] Collecting rollouts ..." +# Back up previous partial results before `gym eval run` clears the file. +ASYNC_BACKUP="" +PREV_MERGED=0 +if [ -s "$ASYNC_FILE" ]; then + ASYNC_BACKUP="$ASYNC_FILE.prev" + cp "$ASYNC_FILE" "$ASYNC_BACKUP" +fi +# Ensure clean slate for first attempt. Prior data is safe in .prev. +# Stale materialized_inputs from a prior Slurm job would cause +# resume_from_cache to load wrong task indexes. +rm -f "$ASYNC_FILE" +MATERIALIZED="$(dirname "$ASYNC_FILE")/$(basename "$ASYNC_FILE" .jsonl-async)_materialized_inputs.jsonl" +rm -f "$MATERIALIZED" +# Retry loop: the vLLM tokenizer race condition (RuntimeError: Already +# borrowed) can crash the client on the initial request burst. Retrying +# after a short delay shifts the timing and almost always succeeds. +# resume_from_cache=true ensures retries skip completed samples. +_NVFLOW_MAX_RETRIES=3 +_NVFLOW_RETRY_DELAY=15 +_NVFLOW_EXIT=0 +for _attempt in $(seq 1 $_NVFLOW_MAX_RETRIES); do + # Guard: truncate corrupted last line from SIGKILL mid-write + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + if [ "$(tail -c 1 "$ASYNC_FILE" | xxd -p)" != "0a" ]; then + head -n -1 "$ASYNC_FILE" > "$ASYNC_FILE.truncated" \ + && mv -f "$ASYNC_FILE.truncated" "$ASYNC_FILE" \ + || rm -f "$ASYNC_FILE.truncated" + echo "[nvflow] Truncated corrupted last line from $ASYNC_FILE" + fi + fi + set +e + gym eval run --no-serve \ + ${AGENT_NAME:++agent_name=$AGENT_NAME} \ + +input_jsonl_fpath=$REMAINING_INPUT \ + +output_jsonl_fpath=$ASYNC_FILE \ + +num_repeats=1 \ + +resume_from_cache=true \ + +num_samples_in_parallel=$NUM_PARALLEL \ + +head_server.host=127.0.0.1 \ + +head_server.port=$HEAD_SERVER_PORT \ + +responses_create_params.max_output_tokens=32768 \ + +responses_create_params.temperature=1.0 + _NVFLOW_EXIT=$? + set -e + [ $_NVFLOW_EXIT -eq 0 ] && break + # F6: `gym eval run` can exit non-zero AFTER all rollouts + # have already been written to ASYNC_FILE -- the gym's + # post-collection aggregate_metrics call returned 500 in the + # Nemotron-Nano smoke run, killing the process at 99% even + # though all 1000 rollouts were on disk. Retrying in that + # state is wasteful (re-runs everything) and dangerous if the + # container FS is being torn down (we hit + # ``/usr/bin/sleep: No such file or directory`` then). If + # ASYNC_FILE has at least as many rows as REMAINING_INPUT, we + # already have what we need; declare success and let finalize + # do its job. + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + _async_rows=$(wc -l < "$ASYNC_FILE" 2>&- || echo 0) + _input_rows=$(wc -l < "$REMAINING_INPUT" 2>&- || echo 0) + if [ "$_input_rows" -gt 0 ] && [ "$_async_rows" -ge "$_input_rows" ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT but $ASYNC_FILE has $_async_rows/$_input_rows rows -- treating as complete (post-collection error in gym, rollouts intact)." + _NVFLOW_EXIT=0 + break + fi + fi + if [ $_attempt -lt $_NVFLOW_MAX_RETRIES ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT (attempt $_attempt/$_NVFLOW_MAX_RETRIES). Retrying in ${_NVFLOW_RETRY_DELAY}s ..." + sleep $_NVFLOW_RETRY_DELAY + _NVFLOW_RETRY_DELAY=$((_NVFLOW_RETRY_DELAY * 2)) + fi +done +if [ $_NVFLOW_EXIT -ne 0 ]; then + echo "[nvflow] gym eval run failed after $_NVFLOW_MAX_RETRIES attempts." + exit $_NVFLOW_EXIT +fi + +# Merge previous partial results with new results. +if [ -n "$ASYNC_BACKUP" ] && [ -f "$ASYNC_BACKUP" ]; then + cat "$ASYNC_BACKUP" "$ASYNC_FILE" > "$ASYNC_FILE.merged" + mv -f "$ASYNC_FILE.merged" "$ASYNC_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_BACKUP" +fi +cp -f "$ASYNC_FILE" "$OUTPUT_FILE" +touch "$DONE_FILE" +# Safe to clean up β€” .done exists, chunk won't be rescheduled. +rm -f "$ASYNC_FILE" "$REMAINING_INPUT" +[ -n "$CHUNK_INPUT" ] && rm -f "$CHUNK_INPUT" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/fixtures/rollout/client_cmd_max_samples.txt b/tests/fixtures/rollout/client_cmd_max_samples.txt new file mode 100644 index 0000000..83c5902 --- /dev/null +++ b/tests/fixtures/rollout/client_cmd_max_samples.txt @@ -0,0 +1,269 @@ +set -e + +OUTPUT_DIR=/out/rollout +GYM_PATH=/opt/Gym +UV_VENV_DIR=/opt/Gym +MODEL_PATH=/hf_models/Qwen/Qwen3-30B-A3B +AGENT_NAME=finance_agent +INPUT_DATA=/data/train.jsonl +OUTPUT_FILE=/out/rollout/rs0/chunk_0.jsonl +DONE_FILE=/out/rollout/rs0/chunk_0.jsonl.done +CONFIG_PATHS=vllm.yaml,env.yaml,overlay.yaml +NUM_PARALLEL=512 +JOB_LABEL=rs0_chunk0 +VLLM_URL="http://policy:8000/v1" +JUDGE_URL="http://judge:8001/v1" +CHUNK_ID=0 +NUM_CHUNKS=8 + +mkdir -p "$OUTPUT_DIR/logs" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +NG_RUN_PID="" + +cleanup() { + local _nvflow_exit=$? + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + # Suppress stderr via ``2>&-`` (close fd) instead of + # ``2>/dev/null`` so the cleanup trap stays quiet even if + # the container's /dev/null disappeared mid-script -- which + # is exactly what happened in the Nemotron-Nano smoke run + # where pyxis tore down the container while bash was still + # in the cleanup path, producing ``/dev/null: No such file + # or directory`` noise on top of the original failure. + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>&- && wait $NG_RUN_PID 2>&- || true + # Best-effort merge of .prev into -async. On success, finalize already + # merged and removed .prev so this block is a no-op. On failure/kill, + # this is a first attempt; the self-heal at next startup is the guarantee. + # Chain with && so .prev is NEVER deleted unless the merge succeeds. + if [ -n "$ASYNC_FILE" ] && [ -f "$ASYNC_FILE.prev" ] && [ "${PREV_MERGED:-0}" -eq 0 ]; then + echo "[Cleanup] Restoring previous results into -async for resume ..." + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.restored" \ + && { [ ! -f "$ASYNC_FILE" ] || cat "$ASYNC_FILE" >> "$ASYNC_FILE.restored"; } \ + && mv -f "$ASYNC_FILE.restored" "$ASYNC_FILE" \ + && rm -f "$ASYNC_FILE.prev" \ + || echo "[Cleanup] WARNING: merge failed β€” self-heal will recover on next start" + fi + if [ $_nvflow_exit -ne 0 ]; then + echo "[nvflow] Client exited with code $_nvflow_exit β€” cancelling job ${SLURM_JOB_ID}" + scancel "${SLURM_JOB_ID}" 2>&- || kill 0 2>&- || true + fi +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Rollout Collection [$JOB_LABEL]" +echo "============================================================" +echo "Model: $MODEL_PATH" +echo "Agent: $AGENT_NAME" +echo "Input data: $INPUT_DATA" +echo "Output file: $OUTPUT_FILE" +echo "Policy URL: $VLLM_URL" +[ -n "$JUDGE_URL" ] && echo "Judge URL: $JUDGE_URL" +echo "============================================================" +if [ -f "$DONE_FILE" ]; then + echo "Chunk already complete (.done exists) β€” skipping." + exit 0 +fi + +echo "" +echo "[Step 1/3] Waiting for vLLM servers ..." +wait_for_server "http://policy:8000/v1/models" "Policy vLLM" $$ 400 /dev/null +wait_for_server "http://judge:8001/v1/models" "Judge vLLM" $$ 400 /dev/null + +CHUNK_INPUT="" +if [ $NUM_CHUNKS -gt 1 ]; then + echo "" + echo "[Step 1a] Extracting chunk slice ..." + TOTAL_LINES=$(wc -l < "$INPUT_DATA") + EFFECTIVE=$TOTAL_LINES + MAX_SAMPLES=10000 + if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then + EFFECTIVE=$MAX_SAMPLES + fi + CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) + [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" + head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" + echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" + INPUT_DATA="$CHUNK_INPUT" +fi +ASYNC_FILE="$OUTPUT_FILE-async" +( + # Case A: output exists without .done β†’ restore to -async + if [ -f "$OUTPUT_FILE" ] && [ ! -f "$DONE_FILE" ]; then + echo "[Self-heal] Output file exists without .done β€” restoring to -async ..." + mv -f "$OUTPUT_FILE" "$ASYNC_FILE" + fi + + # Case B: orphaned .prev β†’ merge into -async + if [ -f "$ASYNC_FILE.prev" ]; then + echo "[Self-heal] Found orphaned .prev β€” merging into -async ..." + PREV_LINES=$(wc -l < "$ASYNC_FILE.prev") + ASYNC_LINES=0 + [ -f "$ASYNC_FILE" ] && ASYNC_LINES=$(wc -l < "$ASYNC_FILE") + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.healed" + [ -f "$ASYNC_FILE" ] && cat "$ASYNC_FILE" >> "$ASYNC_FILE.healed" + mv -f "$ASYNC_FILE.healed" "$ASYNC_FILE" && rm -f "$ASYNC_FILE.prev" + MERGED_LINES=$(wc -l < "$ASYNC_FILE") + echo " Recovered $PREV_LINES (prev) + $ASYNC_LINES (async) = $MERGED_LINES total rows" + fi + + # Case C: clean up orphaned temp files from prior crash + rm -f "$ASYNC_FILE.healed" "$ASYNC_FILE.restored" "$ASYNC_FILE.merged" +) || echo "[Self-heal] WARNING: recovery failed β€” continuing with available data" +REMAINING_INPUT="$OUTPUT_DIR/remaining_input_chunk$CHUNK_ID.jsonl" + +if ! PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.resume_filter "$ASYNC_FILE" "$INPUT_DATA" "$REMAINING_INPUT" 0; then + echo "ERROR: resume_filter failed" >&2 + exit 1 +fi + +if [ -f "$ASYNC_FILE" ] && [ ! -s "$REMAINING_INPUT" ]; then + echo "All rows already completed in -async -- finalizing." + cp -f "$ASYNC_FILE" "$OUTPUT_FILE" + touch "$DONE_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_FILE" "$ASYNC_FILE.prev" + echo "Done [$JOB_LABEL]." + exit 0 +fi + +HEAD_SERVER_PORT=$(find_free_port) + +cd "$GYM_PATH" + +echo "" +echo "[Step 2/3] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=$VLLM_URL" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=$MODEL_PATH" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+port_range_low=1024" \ + "+port_range_high=8999" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \ + "+judge_model.responses_api_models.vllm_model.base_url=http://judge:8001/v1" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 3/3] Collecting rollouts ..." +# Back up previous partial results before `gym eval run` clears the file. +ASYNC_BACKUP="" +PREV_MERGED=0 +if [ -s "$ASYNC_FILE" ]; then + ASYNC_BACKUP="$ASYNC_FILE.prev" + cp "$ASYNC_FILE" "$ASYNC_BACKUP" +fi +# Ensure clean slate for first attempt. Prior data is safe in .prev. +# Stale materialized_inputs from a prior Slurm job would cause +# resume_from_cache to load wrong task indexes. +rm -f "$ASYNC_FILE" +MATERIALIZED="$(dirname "$ASYNC_FILE")/$(basename "$ASYNC_FILE" .jsonl-async)_materialized_inputs.jsonl" +rm -f "$MATERIALIZED" +# Retry loop: the vLLM tokenizer race condition (RuntimeError: Already +# borrowed) can crash the client on the initial request burst. Retrying +# after a short delay shifts the timing and almost always succeeds. +# resume_from_cache=true ensures retries skip completed samples. +_NVFLOW_MAX_RETRIES=3 +_NVFLOW_RETRY_DELAY=15 +_NVFLOW_EXIT=0 +for _attempt in $(seq 1 $_NVFLOW_MAX_RETRIES); do + # Guard: truncate corrupted last line from SIGKILL mid-write + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + if [ "$(tail -c 1 "$ASYNC_FILE" | xxd -p)" != "0a" ]; then + head -n -1 "$ASYNC_FILE" > "$ASYNC_FILE.truncated" \ + && mv -f "$ASYNC_FILE.truncated" "$ASYNC_FILE" \ + || rm -f "$ASYNC_FILE.truncated" + echo "[nvflow] Truncated corrupted last line from $ASYNC_FILE" + fi + fi + set +e + gym eval run --no-serve \ + ${AGENT_NAME:++agent_name=$AGENT_NAME} \ + +input_jsonl_fpath=$REMAINING_INPUT \ + +output_jsonl_fpath=$ASYNC_FILE \ + +num_repeats=1 \ + +resume_from_cache=true \ + +num_samples_in_parallel=$NUM_PARALLEL \ + +head_server.host=127.0.0.1 \ + +head_server.port=$HEAD_SERVER_PORT \ + +responses_create_params.max_output_tokens=32768 \ + +responses_create_params.temperature=1.0 + _NVFLOW_EXIT=$? + set -e + [ $_NVFLOW_EXIT -eq 0 ] && break + # F6: `gym eval run` can exit non-zero AFTER all rollouts + # have already been written to ASYNC_FILE -- the gym's + # post-collection aggregate_metrics call returned 500 in the + # Nemotron-Nano smoke run, killing the process at 99% even + # though all 1000 rollouts were on disk. Retrying in that + # state is wasteful (re-runs everything) and dangerous if the + # container FS is being torn down (we hit + # ``/usr/bin/sleep: No such file or directory`` then). If + # ASYNC_FILE has at least as many rows as REMAINING_INPUT, we + # already have what we need; declare success and let finalize + # do its job. + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + _async_rows=$(wc -l < "$ASYNC_FILE" 2>&- || echo 0) + _input_rows=$(wc -l < "$REMAINING_INPUT" 2>&- || echo 0) + if [ "$_input_rows" -gt 0 ] && [ "$_async_rows" -ge "$_input_rows" ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT but $ASYNC_FILE has $_async_rows/$_input_rows rows -- treating as complete (post-collection error in gym, rollouts intact)." + _NVFLOW_EXIT=0 + break + fi + fi + if [ $_attempt -lt $_NVFLOW_MAX_RETRIES ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT (attempt $_attempt/$_NVFLOW_MAX_RETRIES). Retrying in ${_NVFLOW_RETRY_DELAY}s ..." + sleep $_NVFLOW_RETRY_DELAY + _NVFLOW_RETRY_DELAY=$((_NVFLOW_RETRY_DELAY * 2)) + fi +done +if [ $_NVFLOW_EXIT -ne 0 ]; then + echo "[nvflow] gym eval run failed after $_NVFLOW_MAX_RETRIES attempts." + exit $_NVFLOW_EXIT +fi + +# Merge previous partial results with new results. +if [ -n "$ASYNC_BACKUP" ] && [ -f "$ASYNC_BACKUP" ]; then + cat "$ASYNC_BACKUP" "$ASYNC_FILE" > "$ASYNC_FILE.merged" + mv -f "$ASYNC_FILE.merged" "$ASYNC_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_BACKUP" +fi +cp -f "$ASYNC_FILE" "$OUTPUT_FILE" +touch "$DONE_FILE" +# Safe to clean up β€” .done exists, chunk won't be rescheduled. +rm -f "$ASYNC_FILE" "$REMAINING_INPUT" +[ -n "$CHUNK_INPUT" ] && rm -f "$CHUNK_INPUT" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/fixtures/rollout/client_cmd_no_chunk.txt b/tests/fixtures/rollout/client_cmd_no_chunk.txt new file mode 100644 index 0000000..4a163bc --- /dev/null +++ b/tests/fixtures/rollout/client_cmd_no_chunk.txt @@ -0,0 +1,269 @@ +set -e + +OUTPUT_DIR=/out/rollout +GYM_PATH=/opt/Gym +UV_VENV_DIR=/opt/Gym +MODEL_PATH=/hf_models/Qwen/Qwen3-30B-A3B +AGENT_NAME=finance_agent +INPUT_DATA=/data/train.jsonl +OUTPUT_FILE=/out/rollout/rs0/chunk_0.jsonl +DONE_FILE=/out/rollout/rs0/chunk_0.jsonl.done +CONFIG_PATHS=vllm.yaml,env.yaml,overlay.yaml +NUM_PARALLEL=512 +JOB_LABEL=rs0_chunk0 +VLLM_URL="http://policy:8000/v1" +JUDGE_URL="http://judge:8001/v1" +CHUNK_ID=0 +NUM_CHUNKS=1 + +mkdir -p "$OUTPUT_DIR/logs" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +NG_RUN_PID="" + +cleanup() { + local _nvflow_exit=$? + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + # Suppress stderr via ``2>&-`` (close fd) instead of + # ``2>/dev/null`` so the cleanup trap stays quiet even if + # the container's /dev/null disappeared mid-script -- which + # is exactly what happened in the Nemotron-Nano smoke run + # where pyxis tore down the container while bash was still + # in the cleanup path, producing ``/dev/null: No such file + # or directory`` noise on top of the original failure. + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>&- && wait $NG_RUN_PID 2>&- || true + # Best-effort merge of .prev into -async. On success, finalize already + # merged and removed .prev so this block is a no-op. On failure/kill, + # this is a first attempt; the self-heal at next startup is the guarantee. + # Chain with && so .prev is NEVER deleted unless the merge succeeds. + if [ -n "$ASYNC_FILE" ] && [ -f "$ASYNC_FILE.prev" ] && [ "${PREV_MERGED:-0}" -eq 0 ]; then + echo "[Cleanup] Restoring previous results into -async for resume ..." + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.restored" \ + && { [ ! -f "$ASYNC_FILE" ] || cat "$ASYNC_FILE" >> "$ASYNC_FILE.restored"; } \ + && mv -f "$ASYNC_FILE.restored" "$ASYNC_FILE" \ + && rm -f "$ASYNC_FILE.prev" \ + || echo "[Cleanup] WARNING: merge failed β€” self-heal will recover on next start" + fi + if [ $_nvflow_exit -ne 0 ]; then + echo "[nvflow] Client exited with code $_nvflow_exit β€” cancelling job ${SLURM_JOB_ID}" + scancel "${SLURM_JOB_ID}" 2>&- || kill 0 2>&- || true + fi +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Rollout Collection [$JOB_LABEL]" +echo "============================================================" +echo "Model: $MODEL_PATH" +echo "Agent: $AGENT_NAME" +echo "Input data: $INPUT_DATA" +echo "Output file: $OUTPUT_FILE" +echo "Policy URL: $VLLM_URL" +[ -n "$JUDGE_URL" ] && echo "Judge URL: $JUDGE_URL" +echo "============================================================" +if [ -f "$DONE_FILE" ]; then + echo "Chunk already complete (.done exists) β€” skipping." + exit 0 +fi + +echo "" +echo "[Step 1/3] Waiting for vLLM servers ..." +wait_for_server "http://policy:8000/v1/models" "Policy vLLM" $$ 400 /dev/null +wait_for_server "http://judge:8001/v1/models" "Judge vLLM" $$ 400 /dev/null + +CHUNK_INPUT="" +if [ $NUM_CHUNKS -gt 1 ]; then + echo "" + echo "[Step 1a] Extracting chunk slice ..." + TOTAL_LINES=$(wc -l < "$INPUT_DATA") + EFFECTIVE=$TOTAL_LINES + MAX_SAMPLES=0 + if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then + EFFECTIVE=$MAX_SAMPLES + fi + CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) + [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" + head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" + echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" + INPUT_DATA="$CHUNK_INPUT" +fi +ASYNC_FILE="$OUTPUT_FILE-async" +( + # Case A: output exists without .done β†’ restore to -async + if [ -f "$OUTPUT_FILE" ] && [ ! -f "$DONE_FILE" ]; then + echo "[Self-heal] Output file exists without .done β€” restoring to -async ..." + mv -f "$OUTPUT_FILE" "$ASYNC_FILE" + fi + + # Case B: orphaned .prev β†’ merge into -async + if [ -f "$ASYNC_FILE.prev" ]; then + echo "[Self-heal] Found orphaned .prev β€” merging into -async ..." + PREV_LINES=$(wc -l < "$ASYNC_FILE.prev") + ASYNC_LINES=0 + [ -f "$ASYNC_FILE" ] && ASYNC_LINES=$(wc -l < "$ASYNC_FILE") + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.healed" + [ -f "$ASYNC_FILE" ] && cat "$ASYNC_FILE" >> "$ASYNC_FILE.healed" + mv -f "$ASYNC_FILE.healed" "$ASYNC_FILE" && rm -f "$ASYNC_FILE.prev" + MERGED_LINES=$(wc -l < "$ASYNC_FILE") + echo " Recovered $PREV_LINES (prev) + $ASYNC_LINES (async) = $MERGED_LINES total rows" + fi + + # Case C: clean up orphaned temp files from prior crash + rm -f "$ASYNC_FILE.healed" "$ASYNC_FILE.restored" "$ASYNC_FILE.merged" +) || echo "[Self-heal] WARNING: recovery failed β€” continuing with available data" +REMAINING_INPUT="$OUTPUT_DIR/remaining_input_chunk$CHUNK_ID.jsonl" + +if ! PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.resume_filter "$ASYNC_FILE" "$INPUT_DATA" "$REMAINING_INPUT" 0; then + echo "ERROR: resume_filter failed" >&2 + exit 1 +fi + +if [ -f "$ASYNC_FILE" ] && [ ! -s "$REMAINING_INPUT" ]; then + echo "All rows already completed in -async -- finalizing." + cp -f "$ASYNC_FILE" "$OUTPUT_FILE" + touch "$DONE_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_FILE" "$ASYNC_FILE.prev" + echo "Done [$JOB_LABEL]." + exit 0 +fi + +HEAD_SERVER_PORT=$(find_free_port) + +cd "$GYM_PATH" + +echo "" +echo "[Step 2/3] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=$VLLM_URL" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=$MODEL_PATH" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+port_range_low=1024" \ + "+port_range_high=8999" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \ + "+judge_model.responses_api_models.vllm_model.base_url=http://judge:8001/v1" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 3/3] Collecting rollouts ..." +# Back up previous partial results before `gym eval run` clears the file. +ASYNC_BACKUP="" +PREV_MERGED=0 +if [ -s "$ASYNC_FILE" ]; then + ASYNC_BACKUP="$ASYNC_FILE.prev" + cp "$ASYNC_FILE" "$ASYNC_BACKUP" +fi +# Ensure clean slate for first attempt. Prior data is safe in .prev. +# Stale materialized_inputs from a prior Slurm job would cause +# resume_from_cache to load wrong task indexes. +rm -f "$ASYNC_FILE" +MATERIALIZED="$(dirname "$ASYNC_FILE")/$(basename "$ASYNC_FILE" .jsonl-async)_materialized_inputs.jsonl" +rm -f "$MATERIALIZED" +# Retry loop: the vLLM tokenizer race condition (RuntimeError: Already +# borrowed) can crash the client on the initial request burst. Retrying +# after a short delay shifts the timing and almost always succeeds. +# resume_from_cache=true ensures retries skip completed samples. +_NVFLOW_MAX_RETRIES=3 +_NVFLOW_RETRY_DELAY=15 +_NVFLOW_EXIT=0 +for _attempt in $(seq 1 $_NVFLOW_MAX_RETRIES); do + # Guard: truncate corrupted last line from SIGKILL mid-write + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + if [ "$(tail -c 1 "$ASYNC_FILE" | xxd -p)" != "0a" ]; then + head -n -1 "$ASYNC_FILE" > "$ASYNC_FILE.truncated" \ + && mv -f "$ASYNC_FILE.truncated" "$ASYNC_FILE" \ + || rm -f "$ASYNC_FILE.truncated" + echo "[nvflow] Truncated corrupted last line from $ASYNC_FILE" + fi + fi + set +e + gym eval run --no-serve \ + ${AGENT_NAME:++agent_name=$AGENT_NAME} \ + +input_jsonl_fpath=$REMAINING_INPUT \ + +output_jsonl_fpath=$ASYNC_FILE \ + +num_repeats=1 \ + +resume_from_cache=true \ + +num_samples_in_parallel=$NUM_PARALLEL \ + +head_server.host=127.0.0.1 \ + +head_server.port=$HEAD_SERVER_PORT \ + +responses_create_params.max_output_tokens=32768 \ + +responses_create_params.temperature=1.0 + _NVFLOW_EXIT=$? + set -e + [ $_NVFLOW_EXIT -eq 0 ] && break + # F6: `gym eval run` can exit non-zero AFTER all rollouts + # have already been written to ASYNC_FILE -- the gym's + # post-collection aggregate_metrics call returned 500 in the + # Nemotron-Nano smoke run, killing the process at 99% even + # though all 1000 rollouts were on disk. Retrying in that + # state is wasteful (re-runs everything) and dangerous if the + # container FS is being torn down (we hit + # ``/usr/bin/sleep: No such file or directory`` then). If + # ASYNC_FILE has at least as many rows as REMAINING_INPUT, we + # already have what we need; declare success and let finalize + # do its job. + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + _async_rows=$(wc -l < "$ASYNC_FILE" 2>&- || echo 0) + _input_rows=$(wc -l < "$REMAINING_INPUT" 2>&- || echo 0) + if [ "$_input_rows" -gt 0 ] && [ "$_async_rows" -ge "$_input_rows" ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT but $ASYNC_FILE has $_async_rows/$_input_rows rows -- treating as complete (post-collection error in gym, rollouts intact)." + _NVFLOW_EXIT=0 + break + fi + fi + if [ $_attempt -lt $_NVFLOW_MAX_RETRIES ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT (attempt $_attempt/$_NVFLOW_MAX_RETRIES). Retrying in ${_NVFLOW_RETRY_DELAY}s ..." + sleep $_NVFLOW_RETRY_DELAY + _NVFLOW_RETRY_DELAY=$((_NVFLOW_RETRY_DELAY * 2)) + fi +done +if [ $_NVFLOW_EXIT -ne 0 ]; then + echo "[nvflow] gym eval run failed after $_NVFLOW_MAX_RETRIES attempts." + exit $_NVFLOW_EXIT +fi + +# Merge previous partial results with new results. +if [ -n "$ASYNC_BACKUP" ] && [ -f "$ASYNC_BACKUP" ]; then + cat "$ASYNC_BACKUP" "$ASYNC_FILE" > "$ASYNC_FILE.merged" + mv -f "$ASYNC_FILE.merged" "$ASYNC_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_BACKUP" +fi +cp -f "$ASYNC_FILE" "$OUTPUT_FILE" +touch "$DONE_FILE" +# Safe to clean up β€” .done exists, chunk won't be rescheduled. +rm -f "$ASYNC_FILE" "$REMAINING_INPUT" +[ -n "$CHUNK_INPUT" ] && rm -f "$CHUNK_INPUT" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/fixtures/rollout/client_cmd_no_rcp.txt b/tests/fixtures/rollout/client_cmd_no_rcp.txt new file mode 100644 index 0000000..e6e7937 --- /dev/null +++ b/tests/fixtures/rollout/client_cmd_no_rcp.txt @@ -0,0 +1,267 @@ +set -e + +OUTPUT_DIR=/out/rollout +GYM_PATH=/opt/Gym +UV_VENV_DIR=/opt/Gym +MODEL_PATH=/hf_models/Qwen/Qwen3-30B-A3B +AGENT_NAME=finance_agent +INPUT_DATA=/data/train.jsonl +OUTPUT_FILE=/out/rollout/rs0/chunk_0.jsonl +DONE_FILE=/out/rollout/rs0/chunk_0.jsonl.done +CONFIG_PATHS=vllm.yaml,env.yaml,overlay.yaml +NUM_PARALLEL=512 +JOB_LABEL=rs0_chunk0 +VLLM_URL="http://policy:8000/v1" +JUDGE_URL="http://judge:8001/v1" +CHUNK_ID=0 +NUM_CHUNKS=8 + +mkdir -p "$OUTPUT_DIR/logs" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +NG_RUN_PID="" + +cleanup() { + local _nvflow_exit=$? + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + # Suppress stderr via ``2>&-`` (close fd) instead of + # ``2>/dev/null`` so the cleanup trap stays quiet even if + # the container's /dev/null disappeared mid-script -- which + # is exactly what happened in the Nemotron-Nano smoke run + # where pyxis tore down the container while bash was still + # in the cleanup path, producing ``/dev/null: No such file + # or directory`` noise on top of the original failure. + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>&- && wait $NG_RUN_PID 2>&- || true + # Best-effort merge of .prev into -async. On success, finalize already + # merged and removed .prev so this block is a no-op. On failure/kill, + # this is a first attempt; the self-heal at next startup is the guarantee. + # Chain with && so .prev is NEVER deleted unless the merge succeeds. + if [ -n "$ASYNC_FILE" ] && [ -f "$ASYNC_FILE.prev" ] && [ "${PREV_MERGED:-0}" -eq 0 ]; then + echo "[Cleanup] Restoring previous results into -async for resume ..." + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.restored" \ + && { [ ! -f "$ASYNC_FILE" ] || cat "$ASYNC_FILE" >> "$ASYNC_FILE.restored"; } \ + && mv -f "$ASYNC_FILE.restored" "$ASYNC_FILE" \ + && rm -f "$ASYNC_FILE.prev" \ + || echo "[Cleanup] WARNING: merge failed β€” self-heal will recover on next start" + fi + if [ $_nvflow_exit -ne 0 ]; then + echo "[nvflow] Client exited with code $_nvflow_exit β€” cancelling job ${SLURM_JOB_ID}" + scancel "${SLURM_JOB_ID}" 2>&- || kill 0 2>&- || true + fi +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Rollout Collection [$JOB_LABEL]" +echo "============================================================" +echo "Model: $MODEL_PATH" +echo "Agent: $AGENT_NAME" +echo "Input data: $INPUT_DATA" +echo "Output file: $OUTPUT_FILE" +echo "Policy URL: $VLLM_URL" +[ -n "$JUDGE_URL" ] && echo "Judge URL: $JUDGE_URL" +echo "============================================================" +if [ -f "$DONE_FILE" ]; then + echo "Chunk already complete (.done exists) β€” skipping." + exit 0 +fi + +echo "" +echo "[Step 1/3] Waiting for vLLM servers ..." +wait_for_server "http://policy:8000/v1/models" "Policy vLLM" $$ 400 /dev/null +wait_for_server "http://judge:8001/v1/models" "Judge vLLM" $$ 400 /dev/null + +CHUNK_INPUT="" +if [ $NUM_CHUNKS -gt 1 ]; then + echo "" + echo "[Step 1a] Extracting chunk slice ..." + TOTAL_LINES=$(wc -l < "$INPUT_DATA") + EFFECTIVE=$TOTAL_LINES + MAX_SAMPLES=0 + if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then + EFFECTIVE=$MAX_SAMPLES + fi + CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) + [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" + head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" + echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" + INPUT_DATA="$CHUNK_INPUT" +fi +ASYNC_FILE="$OUTPUT_FILE-async" +( + # Case A: output exists without .done β†’ restore to -async + if [ -f "$OUTPUT_FILE" ] && [ ! -f "$DONE_FILE" ]; then + echo "[Self-heal] Output file exists without .done β€” restoring to -async ..." + mv -f "$OUTPUT_FILE" "$ASYNC_FILE" + fi + + # Case B: orphaned .prev β†’ merge into -async + if [ -f "$ASYNC_FILE.prev" ]; then + echo "[Self-heal] Found orphaned .prev β€” merging into -async ..." + PREV_LINES=$(wc -l < "$ASYNC_FILE.prev") + ASYNC_LINES=0 + [ -f "$ASYNC_FILE" ] && ASYNC_LINES=$(wc -l < "$ASYNC_FILE") + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.healed" + [ -f "$ASYNC_FILE" ] && cat "$ASYNC_FILE" >> "$ASYNC_FILE.healed" + mv -f "$ASYNC_FILE.healed" "$ASYNC_FILE" && rm -f "$ASYNC_FILE.prev" + MERGED_LINES=$(wc -l < "$ASYNC_FILE") + echo " Recovered $PREV_LINES (prev) + $ASYNC_LINES (async) = $MERGED_LINES total rows" + fi + + # Case C: clean up orphaned temp files from prior crash + rm -f "$ASYNC_FILE.healed" "$ASYNC_FILE.restored" "$ASYNC_FILE.merged" +) || echo "[Self-heal] WARNING: recovery failed β€” continuing with available data" +REMAINING_INPUT="$OUTPUT_DIR/remaining_input_chunk$CHUNK_ID.jsonl" + +if ! PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.resume_filter "$ASYNC_FILE" "$INPUT_DATA" "$REMAINING_INPUT" 0; then + echo "ERROR: resume_filter failed" >&2 + exit 1 +fi + +if [ -f "$ASYNC_FILE" ] && [ ! -s "$REMAINING_INPUT" ]; then + echo "All rows already completed in -async -- finalizing." + cp -f "$ASYNC_FILE" "$OUTPUT_FILE" + touch "$DONE_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_FILE" "$ASYNC_FILE.prev" + echo "Done [$JOB_LABEL]." + exit 0 +fi + +HEAD_SERVER_PORT=$(find_free_port) + +cd "$GYM_PATH" + +echo "" +echo "[Step 2/3] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=$VLLM_URL" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=$MODEL_PATH" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+port_range_low=1024" \ + "+port_range_high=8999" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \ + "+judge_model.responses_api_models.vllm_model.base_url=http://judge:8001/v1" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 3/3] Collecting rollouts ..." +# Back up previous partial results before `gym eval run` clears the file. +ASYNC_BACKUP="" +PREV_MERGED=0 +if [ -s "$ASYNC_FILE" ]; then + ASYNC_BACKUP="$ASYNC_FILE.prev" + cp "$ASYNC_FILE" "$ASYNC_BACKUP" +fi +# Ensure clean slate for first attempt. Prior data is safe in .prev. +# Stale materialized_inputs from a prior Slurm job would cause +# resume_from_cache to load wrong task indexes. +rm -f "$ASYNC_FILE" +MATERIALIZED="$(dirname "$ASYNC_FILE")/$(basename "$ASYNC_FILE" .jsonl-async)_materialized_inputs.jsonl" +rm -f "$MATERIALIZED" +# Retry loop: the vLLM tokenizer race condition (RuntimeError: Already +# borrowed) can crash the client on the initial request burst. Retrying +# after a short delay shifts the timing and almost always succeeds. +# resume_from_cache=true ensures retries skip completed samples. +_NVFLOW_MAX_RETRIES=3 +_NVFLOW_RETRY_DELAY=15 +_NVFLOW_EXIT=0 +for _attempt in $(seq 1 $_NVFLOW_MAX_RETRIES); do + # Guard: truncate corrupted last line from SIGKILL mid-write + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + if [ "$(tail -c 1 "$ASYNC_FILE" | xxd -p)" != "0a" ]; then + head -n -1 "$ASYNC_FILE" > "$ASYNC_FILE.truncated" \ + && mv -f "$ASYNC_FILE.truncated" "$ASYNC_FILE" \ + || rm -f "$ASYNC_FILE.truncated" + echo "[nvflow] Truncated corrupted last line from $ASYNC_FILE" + fi + fi + set +e + gym eval run --no-serve \ + ${AGENT_NAME:++agent_name=$AGENT_NAME} \ + +input_jsonl_fpath=$REMAINING_INPUT \ + +output_jsonl_fpath=$ASYNC_FILE \ + +num_repeats=1 \ + +resume_from_cache=true \ + +num_samples_in_parallel=$NUM_PARALLEL \ + +head_server.host=127.0.0.1 \ + +head_server.port=$HEAD_SERVER_PORT + _NVFLOW_EXIT=$? + set -e + [ $_NVFLOW_EXIT -eq 0 ] && break + # F6: `gym eval run` can exit non-zero AFTER all rollouts + # have already been written to ASYNC_FILE -- the gym's + # post-collection aggregate_metrics call returned 500 in the + # Nemotron-Nano smoke run, killing the process at 99% even + # though all 1000 rollouts were on disk. Retrying in that + # state is wasteful (re-runs everything) and dangerous if the + # container FS is being torn down (we hit + # ``/usr/bin/sleep: No such file or directory`` then). If + # ASYNC_FILE has at least as many rows as REMAINING_INPUT, we + # already have what we need; declare success and let finalize + # do its job. + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + _async_rows=$(wc -l < "$ASYNC_FILE" 2>&- || echo 0) + _input_rows=$(wc -l < "$REMAINING_INPUT" 2>&- || echo 0) + if [ "$_input_rows" -gt 0 ] && [ "$_async_rows" -ge "$_input_rows" ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT but $ASYNC_FILE has $_async_rows/$_input_rows rows -- treating as complete (post-collection error in gym, rollouts intact)." + _NVFLOW_EXIT=0 + break + fi + fi + if [ $_attempt -lt $_NVFLOW_MAX_RETRIES ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT (attempt $_attempt/$_NVFLOW_MAX_RETRIES). Retrying in ${_NVFLOW_RETRY_DELAY}s ..." + sleep $_NVFLOW_RETRY_DELAY + _NVFLOW_RETRY_DELAY=$((_NVFLOW_RETRY_DELAY * 2)) + fi +done +if [ $_NVFLOW_EXIT -ne 0 ]; then + echo "[nvflow] gym eval run failed after $_NVFLOW_MAX_RETRIES attempts." + exit $_NVFLOW_EXIT +fi + +# Merge previous partial results with new results. +if [ -n "$ASYNC_BACKUP" ] && [ -f "$ASYNC_BACKUP" ]; then + cat "$ASYNC_BACKUP" "$ASYNC_FILE" > "$ASYNC_FILE.merged" + mv -f "$ASYNC_FILE.merged" "$ASYNC_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_BACKUP" +fi +cp -f "$ASYNC_FILE" "$OUTPUT_FILE" +touch "$DONE_FILE" +# Safe to clean up β€” .done exists, chunk won't be rescheduled. +rm -f "$ASYNC_FILE" "$REMAINING_INPUT" +[ -n "$CHUNK_INPUT" ] && rm -f "$CHUNK_INPUT" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/fixtures/rollout/client_cmd_policy_only.txt b/tests/fixtures/rollout/client_cmd_policy_only.txt new file mode 100644 index 0000000..69062a3 --- /dev/null +++ b/tests/fixtures/rollout/client_cmd_policy_only.txt @@ -0,0 +1,266 @@ +set -e + +OUTPUT_DIR=/out/rollout +GYM_PATH=/opt/Gym +UV_VENV_DIR=/opt/Gym +MODEL_PATH=/hf_models/Qwen/Qwen3-30B-A3B +AGENT_NAME=finance_agent +INPUT_DATA=/data/train.jsonl +OUTPUT_FILE=/out/rollout/rs0/chunk_0.jsonl +DONE_FILE=/out/rollout/rs0/chunk_0.jsonl.done +CONFIG_PATHS=vllm.yaml,env.yaml,overlay.yaml +NUM_PARALLEL=512 +JOB_LABEL=rs0_chunk0 +VLLM_URL="http://policy:8000/v1" +JUDGE_URL="" +CHUNK_ID=0 +NUM_CHUNKS=8 + +mkdir -p "$OUTPUT_DIR/logs" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +NG_RUN_PID="" + +cleanup() { + local _nvflow_exit=$? + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + # Suppress stderr via ``2>&-`` (close fd) instead of + # ``2>/dev/null`` so the cleanup trap stays quiet even if + # the container's /dev/null disappeared mid-script -- which + # is exactly what happened in the Nemotron-Nano smoke run + # where pyxis tore down the container while bash was still + # in the cleanup path, producing ``/dev/null: No such file + # or directory`` noise on top of the original failure. + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>&- && wait $NG_RUN_PID 2>&- || true + # Best-effort merge of .prev into -async. On success, finalize already + # merged and removed .prev so this block is a no-op. On failure/kill, + # this is a first attempt; the self-heal at next startup is the guarantee. + # Chain with && so .prev is NEVER deleted unless the merge succeeds. + if [ -n "$ASYNC_FILE" ] && [ -f "$ASYNC_FILE.prev" ] && [ "${PREV_MERGED:-0}" -eq 0 ]; then + echo "[Cleanup] Restoring previous results into -async for resume ..." + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.restored" \ + && { [ ! -f "$ASYNC_FILE" ] || cat "$ASYNC_FILE" >> "$ASYNC_FILE.restored"; } \ + && mv -f "$ASYNC_FILE.restored" "$ASYNC_FILE" \ + && rm -f "$ASYNC_FILE.prev" \ + || echo "[Cleanup] WARNING: merge failed β€” self-heal will recover on next start" + fi + if [ $_nvflow_exit -ne 0 ]; then + echo "[nvflow] Client exited with code $_nvflow_exit β€” cancelling job ${SLURM_JOB_ID}" + scancel "${SLURM_JOB_ID}" 2>&- || kill 0 2>&- || true + fi +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Rollout Collection [$JOB_LABEL]" +echo "============================================================" +echo "Model: $MODEL_PATH" +echo "Agent: $AGENT_NAME" +echo "Input data: $INPUT_DATA" +echo "Output file: $OUTPUT_FILE" +echo "Policy URL: $VLLM_URL" +[ -n "$JUDGE_URL" ] && echo "Judge URL: $JUDGE_URL" +echo "============================================================" +if [ -f "$DONE_FILE" ]; then + echo "Chunk already complete (.done exists) β€” skipping." + exit 0 +fi + +echo "" +echo "[Step 1/3] Waiting for vLLM servers ..." +wait_for_server "http://policy:8000/v1/models" "Policy vLLM" $$ 400 /dev/null + +CHUNK_INPUT="" +if [ $NUM_CHUNKS -gt 1 ]; then + echo "" + echo "[Step 1a] Extracting chunk slice ..." + TOTAL_LINES=$(wc -l < "$INPUT_DATA") + EFFECTIVE=$TOTAL_LINES + MAX_SAMPLES=0 + if [ $MAX_SAMPLES -gt 0 ] && [ $MAX_SAMPLES -lt $TOTAL_LINES ]; then + EFFECTIVE=$MAX_SAMPLES + fi + CHUNK_SIZE=$(( (EFFECTIVE + NUM_CHUNKS - 1) / NUM_CHUNKS )) + START_LINE=$(( CHUNK_ID * CHUNK_SIZE + 1 )) + END_LINE=$(( (CHUNK_ID + 1) * CHUNK_SIZE )) + [ $END_LINE -gt $EFFECTIVE ] && END_LINE=$EFFECTIVE + CHUNK_INPUT="$OUTPUT_DIR/chunk_input_chunk$CHUNK_ID.jsonl" + head -n $END_LINE "$INPUT_DATA" | tail -n +$START_LINE > "$CHUNK_INPUT" + echo " Chunk $CHUNK_ID/$NUM_CHUNKS: lines $START_LINE-$END_LINE ($((END_LINE - START_LINE + 1)) samples)" + INPUT_DATA="$CHUNK_INPUT" +fi +ASYNC_FILE="$OUTPUT_FILE-async" +( + # Case A: output exists without .done β†’ restore to -async + if [ -f "$OUTPUT_FILE" ] && [ ! -f "$DONE_FILE" ]; then + echo "[Self-heal] Output file exists without .done β€” restoring to -async ..." + mv -f "$OUTPUT_FILE" "$ASYNC_FILE" + fi + + # Case B: orphaned .prev β†’ merge into -async + if [ -f "$ASYNC_FILE.prev" ]; then + echo "[Self-heal] Found orphaned .prev β€” merging into -async ..." + PREV_LINES=$(wc -l < "$ASYNC_FILE.prev") + ASYNC_LINES=0 + [ -f "$ASYNC_FILE" ] && ASYNC_LINES=$(wc -l < "$ASYNC_FILE") + cat "$ASYNC_FILE.prev" > "$ASYNC_FILE.healed" + [ -f "$ASYNC_FILE" ] && cat "$ASYNC_FILE" >> "$ASYNC_FILE.healed" + mv -f "$ASYNC_FILE.healed" "$ASYNC_FILE" && rm -f "$ASYNC_FILE.prev" + MERGED_LINES=$(wc -l < "$ASYNC_FILE") + echo " Recovered $PREV_LINES (prev) + $ASYNC_LINES (async) = $MERGED_LINES total rows" + fi + + # Case C: clean up orphaned temp files from prior crash + rm -f "$ASYNC_FILE.healed" "$ASYNC_FILE.restored" "$ASYNC_FILE.merged" +) || echo "[Self-heal] WARNING: recovery failed β€” continuing with available data" +REMAINING_INPUT="$OUTPUT_DIR/remaining_input_chunk$CHUNK_ID.jsonl" + +if ! PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.resume_filter "$ASYNC_FILE" "$INPUT_DATA" "$REMAINING_INPUT" 0; then + echo "ERROR: resume_filter failed" >&2 + exit 1 +fi + +if [ -f "$ASYNC_FILE" ] && [ ! -s "$REMAINING_INPUT" ]; then + echo "All rows already completed in -async -- finalizing." + cp -f "$ASYNC_FILE" "$OUTPUT_FILE" + touch "$DONE_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_FILE" "$ASYNC_FILE.prev" + echo "Done [$JOB_LABEL]." + exit 0 +fi + +HEAD_SERVER_PORT=$(find_free_port) + +cd "$GYM_PATH" + +echo "" +echo "[Step 2/3] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=$VLLM_URL" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=$MODEL_PATH" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+port_range_low=1024" \ + "+port_range_high=8999" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 3/3] Collecting rollouts ..." +# Back up previous partial results before `gym eval run` clears the file. +ASYNC_BACKUP="" +PREV_MERGED=0 +if [ -s "$ASYNC_FILE" ]; then + ASYNC_BACKUP="$ASYNC_FILE.prev" + cp "$ASYNC_FILE" "$ASYNC_BACKUP" +fi +# Ensure clean slate for first attempt. Prior data is safe in .prev. +# Stale materialized_inputs from a prior Slurm job would cause +# resume_from_cache to load wrong task indexes. +rm -f "$ASYNC_FILE" +MATERIALIZED="$(dirname "$ASYNC_FILE")/$(basename "$ASYNC_FILE" .jsonl-async)_materialized_inputs.jsonl" +rm -f "$MATERIALIZED" +# Retry loop: the vLLM tokenizer race condition (RuntimeError: Already +# borrowed) can crash the client on the initial request burst. Retrying +# after a short delay shifts the timing and almost always succeeds. +# resume_from_cache=true ensures retries skip completed samples. +_NVFLOW_MAX_RETRIES=3 +_NVFLOW_RETRY_DELAY=15 +_NVFLOW_EXIT=0 +for _attempt in $(seq 1 $_NVFLOW_MAX_RETRIES); do + # Guard: truncate corrupted last line from SIGKILL mid-write + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + if [ "$(tail -c 1 "$ASYNC_FILE" | xxd -p)" != "0a" ]; then + head -n -1 "$ASYNC_FILE" > "$ASYNC_FILE.truncated" \ + && mv -f "$ASYNC_FILE.truncated" "$ASYNC_FILE" \ + || rm -f "$ASYNC_FILE.truncated" + echo "[nvflow] Truncated corrupted last line from $ASYNC_FILE" + fi + fi + set +e + gym eval run --no-serve \ + ${AGENT_NAME:++agent_name=$AGENT_NAME} \ + +input_jsonl_fpath=$REMAINING_INPUT \ + +output_jsonl_fpath=$ASYNC_FILE \ + +num_repeats=1 \ + +resume_from_cache=true \ + +num_samples_in_parallel=$NUM_PARALLEL \ + +head_server.host=127.0.0.1 \ + +head_server.port=$HEAD_SERVER_PORT \ + +responses_create_params.max_output_tokens=32768 \ + +responses_create_params.temperature=1.0 + _NVFLOW_EXIT=$? + set -e + [ $_NVFLOW_EXIT -eq 0 ] && break + # F6: `gym eval run` can exit non-zero AFTER all rollouts + # have already been written to ASYNC_FILE -- the gym's + # post-collection aggregate_metrics call returned 500 in the + # Nemotron-Nano smoke run, killing the process at 99% even + # though all 1000 rollouts were on disk. Retrying in that + # state is wasteful (re-runs everything) and dangerous if the + # container FS is being torn down (we hit + # ``/usr/bin/sleep: No such file or directory`` then). If + # ASYNC_FILE has at least as many rows as REMAINING_INPUT, we + # already have what we need; declare success and let finalize + # do its job. + if [ -f "$ASYNC_FILE" ] && [ -s "$ASYNC_FILE" ]; then + _async_rows=$(wc -l < "$ASYNC_FILE" 2>&- || echo 0) + _input_rows=$(wc -l < "$REMAINING_INPUT" 2>&- || echo 0) + if [ "$_input_rows" -gt 0 ] && [ "$_async_rows" -ge "$_input_rows" ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT but $ASYNC_FILE has $_async_rows/$_input_rows rows -- treating as complete (post-collection error in gym, rollouts intact)." + _NVFLOW_EXIT=0 + break + fi + fi + if [ $_attempt -lt $_NVFLOW_MAX_RETRIES ]; then + echo "[nvflow] gym eval run exited $_NVFLOW_EXIT (attempt $_attempt/$_NVFLOW_MAX_RETRIES). Retrying in ${_NVFLOW_RETRY_DELAY}s ..." + sleep $_NVFLOW_RETRY_DELAY + _NVFLOW_RETRY_DELAY=$((_NVFLOW_RETRY_DELAY * 2)) + fi +done +if [ $_NVFLOW_EXIT -ne 0 ]; then + echo "[nvflow] gym eval run failed after $_NVFLOW_MAX_RETRIES attempts." + exit $_NVFLOW_EXIT +fi + +# Merge previous partial results with new results. +if [ -n "$ASYNC_BACKUP" ] && [ -f "$ASYNC_BACKUP" ]; then + cat "$ASYNC_BACKUP" "$ASYNC_FILE" > "$ASYNC_FILE.merged" + mv -f "$ASYNC_FILE.merged" "$ASYNC_FILE" + PREV_MERGED=1 + rm -f "$ASYNC_BACKUP" +fi +cp -f "$ASYNC_FILE" "$OUTPUT_FILE" +touch "$DONE_FILE" +# Safe to clean up β€” .done exists, chunk won't be rescheduled. +rm -f "$ASYNC_FILE" "$REMAINING_INPUT" +[ -n "$CHUNK_INPUT" ] && rm -f "$CHUNK_INPUT" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/fixtures/rollout/filter_cmd_full.txt b/tests/fixtures/rollout/filter_cmd_full.txt new file mode 100644 index 0000000..f805aaa --- /dev/null +++ b/tests/fixtures/rollout/filter_cmd_full.txt @@ -0,0 +1,14 @@ +set -e +echo "Filter Training Data (reward-variance difficulty)" +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.filter_training_data \ + "/in/train.jsonl" \ + "/out/rollout/aggregate/difficulty.jsonl" \ + "/out" \ + --min-reward-std 1e-06 \ + --train-filename "train.jsonl" \ + --val-filename "validation.jsonl" \ + --report-filename "filter_report.json" \ + --validation-data "/in/val.jsonl" \ + --policy-model "/hf_models/Qwen/Qwen3-30B-A3B" \ + --judge-model "/hf_models/openai/gpt-oss-120b" +echo "Done. Filtered data in /out/" diff --git a/tests/fixtures/rollout/filter_cmd_minimal.txt b/tests/fixtures/rollout/filter_cmd_minimal.txt new file mode 100644 index 0000000..c0fc3de --- /dev/null +++ b/tests/fixtures/rollout/filter_cmd_minimal.txt @@ -0,0 +1,11 @@ +set -e +echo "Filter Training Data (reward-variance difficulty)" +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.filter_training_data \ + "/in/train.jsonl" \ + "/out/rollout/aggregate/difficulty.jsonl" \ + "/out" \ + --min-reward-std 1e-06 \ + --train-filename "train.jsonl" \ + --val-filename "validation.jsonl" \ + --report-filename "filter_report.json" +echo "Done. Filtered data in /out/" diff --git a/tests/fixtures/rollout/merge_cmd_1chunk.txt b/tests/fixtures/rollout/merge_cmd_1chunk.txt new file mode 100644 index 0000000..ae23811 --- /dev/null +++ b/tests/fixtures/rollout/merge_cmd_1chunk.txt @@ -0,0 +1,78 @@ +set -e + +MERGED_FILE=/out/rollout/output-rs0.jsonl +ANALYSIS_DIR=/out/rollout/analysis_rs0 +SEED_LABEL=rs0 +NUM_CHUNKS=1 +INPUT_DATA=/data/train.jsonl + +echo "============================================================" +echo "Merge Rollout Chunks [$SEED_LABEL]" +echo "============================================================" + +# Clean up stale temp file from a prior crashed merge +rm -f "$MERGED_FILE.tmp" + +# Precondition: ALL chunk .done markers must exist. A missing +# .done means upstream rollout never finalised (job failed before +# the cp/touch finalize step, or output was deleted). Fail loudly +# so Slurm marks the merge job FAILED and an operator notices -- +# silently exit-0'ing here is what produced the silent-success +# cascade where aggregate runs on N-1 seeds and the pipeline +# reports COMPLETED 0:0 despite missing training data. +for i in $(seq 0 $((NUM_CHUNKS - 1))); do + CHUNK_DONE="/out/rollout/rs0/chunk_$i.jsonl.done" + if [ ! -f "$CHUNK_DONE" ]; then + echo "ERROR: chunk $i .done missing for [$SEED_LABEL] -- upstream rollout failed before finalize. Aborting merge." + exit 1 + fi +done + +echo "[Step 1/3] Merging chunk files ..." +> "$MERGED_FILE.tmp" +for i in $(seq 0 $((NUM_CHUNKS - 1))); do + CHUNK_FILE="/out/rollout/rs0/chunk_$i.jsonl" + if [ ! -f "$CHUNK_FILE" ] || [ ! -s "$CHUNK_FILE" ]; then + echo "ERROR: chunk $i .done exists but file missing/empty β€” aborting." + rm -f "$MERGED_FILE.tmp" + exit 1 + fi + LINES=$(wc -l < "$CHUNK_FILE") + echo " Chunk $i: $LINES lines" + cat "$CHUNK_FILE" >> "$MERGED_FILE.tmp" +done + +TOTAL=$(wc -l < "$MERGED_FILE.tmp") +echo " Merged total: $TOTAL lines" +if [ "$TOTAL" -eq 0 ]; then + echo "ERROR: All chunks present but 0 lines merged." + rm -f "$MERGED_FILE.tmp" + exit 1 +fi + +# Atomic replace β€” old merged file untouched until this point +mv -f "$MERGED_FILE.tmp" "$MERGED_FILE" + +echo "" +echo "[Step 2/3] Enriching rollouts with input metadata ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.enrich_rollouts \ + "$INPUT_DATA" \ + "$MERGED_FILE" + +echo "" +echo "[Step 3/3] Analyzing rollouts ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.analyze_rollouts \ + "$MERGED_FILE" \ + "$ANALYSIS_DIR" + +touch /out/rollout/output-rs0.jsonl.done + +# Safe cleanup: delete chunk data files only (keep .done markers). +for i in $(seq 0 $((NUM_CHUNKS - 1))); do + CHUNK_FILE="/out/rollout/rs0/chunk_$i.jsonl" + rm -f "$CHUNK_FILE" +done +echo "Done [$SEED_LABEL]." +echo "" +echo "To browse rollouts interactively (in the nemo-gym container):" +echo " export PATH=/opt/gym-cli-venv/bin:$PATH && ng_viewer +jsonl_fpath=$MERGED_FILE" diff --git a/tests/fixtures/rollout/merge_cmd_8chunks.txt b/tests/fixtures/rollout/merge_cmd_8chunks.txt new file mode 100644 index 0000000..82e18f2 --- /dev/null +++ b/tests/fixtures/rollout/merge_cmd_8chunks.txt @@ -0,0 +1,78 @@ +set -e + +MERGED_FILE=/out/rollout/output-rs0.jsonl +ANALYSIS_DIR=/out/rollout/analysis_rs0 +SEED_LABEL=rs0 +NUM_CHUNKS=8 +INPUT_DATA=/data/train.jsonl + +echo "============================================================" +echo "Merge Rollout Chunks [$SEED_LABEL]" +echo "============================================================" + +# Clean up stale temp file from a prior crashed merge +rm -f "$MERGED_FILE.tmp" + +# Precondition: ALL chunk .done markers must exist. A missing +# .done means upstream rollout never finalised (job failed before +# the cp/touch finalize step, or output was deleted). Fail loudly +# so Slurm marks the merge job FAILED and an operator notices -- +# silently exit-0'ing here is what produced the silent-success +# cascade where aggregate runs on N-1 seeds and the pipeline +# reports COMPLETED 0:0 despite missing training data. +for i in $(seq 0 $((NUM_CHUNKS - 1))); do + CHUNK_DONE="/out/rollout/rs0/chunk_$i.jsonl.done" + if [ ! -f "$CHUNK_DONE" ]; then + echo "ERROR: chunk $i .done missing for [$SEED_LABEL] -- upstream rollout failed before finalize. Aborting merge." + exit 1 + fi +done + +echo "[Step 1/3] Merging chunk files ..." +> "$MERGED_FILE.tmp" +for i in $(seq 0 $((NUM_CHUNKS - 1))); do + CHUNK_FILE="/out/rollout/rs0/chunk_$i.jsonl" + if [ ! -f "$CHUNK_FILE" ] || [ ! -s "$CHUNK_FILE" ]; then + echo "ERROR: chunk $i .done exists but file missing/empty β€” aborting." + rm -f "$MERGED_FILE.tmp" + exit 1 + fi + LINES=$(wc -l < "$CHUNK_FILE") + echo " Chunk $i: $LINES lines" + cat "$CHUNK_FILE" >> "$MERGED_FILE.tmp" +done + +TOTAL=$(wc -l < "$MERGED_FILE.tmp") +echo " Merged total: $TOTAL lines" +if [ "$TOTAL" -eq 0 ]; then + echo "ERROR: All chunks present but 0 lines merged." + rm -f "$MERGED_FILE.tmp" + exit 1 +fi + +# Atomic replace β€” old merged file untouched until this point +mv -f "$MERGED_FILE.tmp" "$MERGED_FILE" + +echo "" +echo "[Step 2/3] Enriching rollouts with input metadata ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.enrich_rollouts \ + "$INPUT_DATA" \ + "$MERGED_FILE" + +echo "" +echo "[Step 3/3] Analyzing rollouts ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.analyze_rollouts \ + "$MERGED_FILE" \ + "$ANALYSIS_DIR" + +touch /out/rollout/output-rs0.jsonl.done + +# Safe cleanup: delete chunk data files only (keep .done markers). +for i in $(seq 0 $((NUM_CHUNKS - 1))); do + CHUNK_FILE="/out/rollout/rs0/chunk_$i.jsonl" + rm -f "$CHUNK_FILE" +done +echo "Done [$SEED_LABEL]." +echo "" +echo "To browse rollouts interactively (in the nemo-gym container):" +echo " export PATH=/opt/gym-cli-venv/bin:$PATH && ng_viewer +jsonl_fpath=$MERGED_FILE" diff --git a/tests/fixtures/verify/analysis_cmd_empty_entries.txt b/tests/fixtures/verify/analysis_cmd_empty_entries.txt new file mode 100644 index 0000000..015d80f --- /dev/null +++ b/tests/fixtures/verify/analysis_cmd_empty_entries.txt @@ -0,0 +1,6 @@ +set -e +echo "Reward Analysis" +echo "Done. Analysis complete." +echo "" +echo "To browse re-judged rollouts interactively (in the nemo-gym container):" +echo " export PATH=/opt/gym-cli-venv/bin:$PATH && ng_viewer +jsonl_fpath=/out/verify/rejudge/output-rs0.jsonl" diff --git a/tests/fixtures/verify/analysis_cmd_multi_seed.txt b/tests/fixtures/verify/analysis_cmd_multi_seed.txt new file mode 100644 index 0000000..f3f5aaa --- /dev/null +++ b/tests/fixtures/verify/analysis_cmd_multi_seed.txt @@ -0,0 +1,16 @@ +set -e +echo "Reward Analysis" +echo "Analyzing rs0 ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.analyze_rollouts \ + "/out/verify/rejudge/output-rs0.jsonl" \ + "/out/verify/rejudge/analysis_rs0" \ + "REWARD RE-COMPUTATION ANALYSIS" +echo "Analyzing rs1 ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.analyze_rollouts \ + "/out/verify/rejudge/output-rs1.jsonl" \ + "/out/verify/rejudge/analysis_rs1" \ + "REWARD RE-COMPUTATION ANALYSIS" +echo "Done. Analysis complete." +echo "" +echo "To browse re-judged rollouts interactively (in the nemo-gym container):" +echo " export PATH=/opt/gym-cli-venv/bin:$PATH && ng_viewer +jsonl_fpath=/out/verify/rejudge/output-rs0.jsonl" diff --git a/tests/fixtures/verify/analysis_cmd_single_seed.txt b/tests/fixtures/verify/analysis_cmd_single_seed.txt new file mode 100644 index 0000000..1068466 --- /dev/null +++ b/tests/fixtures/verify/analysis_cmd_single_seed.txt @@ -0,0 +1,11 @@ +set -e +echo "Reward Analysis" +echo "Analyzing rs0 ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.recipes.finance.utils.rl.analyze_rollouts \ + "/out/verify/rejudge/output-rs0.jsonl" \ + "/out/verify/rejudge/analysis_rs0" \ + "REWARD RE-COMPUTATION ANALYSIS" +echo "Done. Analysis complete." +echo "" +echo "To browse re-judged rollouts interactively (in the nemo-gym container):" +echo " export PATH=/opt/gym-cli-venv/bin:$PATH && ng_viewer +jsonl_fpath=/out/verify/rejudge/output-rs0.jsonl" diff --git a/tests/fixtures/verify/verify_cmd_local_judge.txt b/tests/fixtures/verify/verify_cmd_local_judge.txt new file mode 100644 index 0000000..4c1073d --- /dev/null +++ b/tests/fixtures/verify/verify_cmd_local_judge.txt @@ -0,0 +1,94 @@ +set -e + +OUTPUT_DIR="/out/verify" +GYM_PATH="/opt/Gym" +UV_VENV_DIR="/opt/Gym" +INPUT_FILE="/in/rollouts/output-rs0.jsonl" +OUTPUT_FILE="/out/verify/rejudge/output-rs0.jsonl" +DONE_FILE="/out/verify/rejudge/output-rs0.jsonl.done" +CONFIG_PATHS="vllm.yaml,env.yaml,overlay.yaml" +NUM_PARALLEL="8" +JOB_LABEL="rejudge_rs0" +ENVIRONMENT_NAME="finance_env" + +mkdir -p "$OUTPUT_DIR/logs" "$OUTPUT_DIR/rejudge" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +HEAD_SERVER_PORT=$(find_free_port) + +NG_RUN_PID="" + +cleanup() { + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>/dev/null && wait $NG_RUN_PID 2>/dev/null || true +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Compute Rewards (re-judge) [$JOB_LABEL]" +echo "============================================================" +echo "Input file: $INPUT_FILE" +echo "Output file: $OUTPUT_FILE" +echo "Judge mode: local_vllm" +echo "Environment: $ENVIRONMENT_NAME" +echo "============================================================" + +cd "$GYM_PATH" + +echo "" +echo "[Step 1/2] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=http://localhost:0/v1" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=unused" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + "+judge_model.responses_api_models.vllm_model.entrypoint=app.py" \ + "+judge_model.responses_api_models.vllm_model.base_url=http://127.0.0.1:$JUDGE_PORT/v1" \ + "+judge_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+judge_model.responses_api_models.vllm_model.model=/hf_models/openai/gpt-oss-120b" \ + "+judge_model.responses_api_models.vllm_model.return_token_id_information=false" \ + "+judge_model.responses_api_models.vllm_model.uses_reasoning_parser=true" \ + "+finance_env.resources_servers.finance_env.judge_model_server.name=judge_model" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 2/2] Re-judging rollouts ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.verify_worker \ + "$INPUT_FILE" \ + "$OUTPUT_FILE-async" \ + "127.0.0.1" \ + "$HEAD_SERVER_PORT" \ + "$ENVIRONMENT_NAME" \ + "$NUM_PARALLEL" + +mv "$OUTPUT_FILE-async" "$OUTPUT_FILE" +touch "$DONE_FILE" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/fixtures/verify/verify_cmd_openai_judge.txt b/tests/fixtures/verify/verify_cmd_openai_judge.txt new file mode 100644 index 0000000..834a268 --- /dev/null +++ b/tests/fixtures/verify/verify_cmd_openai_judge.txt @@ -0,0 +1,91 @@ +set -e + +OUTPUT_DIR="/out/verify" +GYM_PATH="/opt/Gym" +UV_VENV_DIR="/opt/Gym" +INPUT_FILE="/in/rollouts/output-rs0.jsonl" +OUTPUT_FILE="/out/verify/rejudge/output-rs0.jsonl" +DONE_FILE="/out/verify/rejudge/output-rs0.jsonl.done" +CONFIG_PATHS="vllm.yaml,env.yaml,overlay.yaml" +NUM_PARALLEL="8" +JOB_LABEL="rejudge_rs0" +ENVIRONMENT_NAME="finance_env" + +mkdir -p "$OUTPUT_DIR/logs" "$OUTPUT_DIR/rejudge" + +find_free_port() { + python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" +} + +HEAD_SERVER_PORT=$(find_free_port) + +NG_RUN_PID="" + +cleanup() { + echo "" + echo "[Cleanup] Shutting down NeMo-Gym servers ..." + [ -n "$NG_RUN_PID" ] && kill $NG_RUN_PID 2>/dev/null && wait $NG_RUN_PID 2>/dev/null || true +} +trap cleanup EXIT + +wait_for_server() { + local url="$1" name="$2" pid="$3" max_attempts="$4" log="$5" + echo " Waiting for $name at $url ..." + for i in $(seq 1 $max_attempts); do + if curl -s -m 5 "$url" > /dev/null 2>&1; then + echo " $name ready after $((i * 5))s" + return 0 + fi + if ! kill -0 $pid 2>/dev/null; then + echo "ERROR: $name died. Check $log" + exit 1 + fi + sleep 5 + done + echo "ERROR: $name did not start within $((max_attempts * 5))s" + exit 1 +} + +echo "============================================================" +echo "Compute Rewards (re-judge) [$JOB_LABEL]" +echo "============================================================" +echo "Input file: $INPUT_FILE" +echo "Output file: $OUTPUT_FILE" +echo "Judge mode: openai" +echo "Environment: $ENVIRONMENT_NAME" +echo "============================================================" + +cd "$GYM_PATH" + +echo "" +echo "[Step 1/2] Starting NeMo-Gym servers ..." +gym env start "+config_paths=[$CONFIG_PATHS]" \ + "+policy_model.responses_api_models.vllm_model.base_url=http://localhost:0/v1" \ + "+policy_model.responses_api_models.vllm_model.api_key=EMPTY" \ + "+policy_model.responses_api_models.vllm_model.model=unused" \ + "+head_server.host=127.0.0.1" \ + "+head_server.port=$HEAD_SERVER_PORT" \ + "+skip_venv_if_present=true" \ + "+uv_venv_dir=$UV_VENV_DIR" \ + "+judge_model.responses_api_models.openai_model.base_url=https://api.openai.com/v1" \ + "+judge_model.responses_api_models.openai_model.api_key_env_var=OPENAI_API_KEY" \ + "+judge_model.responses_api_models.openai_model.model=gpt-4o-mini" \ + "+finance_env.resources_servers.finance_env.judge_model_server.name=judge_model" \ + > "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" 2>&1 & +NG_RUN_PID=$! + +wait_for_server "http://127.0.0.1:$HEAD_SERVER_PORT/" "NeMo-Gym" $NG_RUN_PID 60 "$OUTPUT_DIR/logs/ng_run_$JOB_LABEL.log" + +echo "" +echo "[Step 2/2] Re-judging rollouts ..." +PYTHONPATH=/nemo_run/code python3 -m nvflow.lib.rl.verify_worker \ + "$INPUT_FILE" \ + "$OUTPUT_FILE-async" \ + "127.0.0.1" \ + "$HEAD_SERVER_PORT" \ + "$ENVIRONMENT_NAME" \ + "$NUM_PARALLEL" + +mv "$OUTPUT_FILE-async" "$OUTPUT_FILE" +touch "$DONE_FILE" +echo "Done [$JOB_LABEL]. Cleanup via trap." diff --git a/tests/requirements-ci.txt b/tests/requirements-ci.txt new file mode 100644 index 0000000..2463d0c --- /dev/null +++ b/tests/requirements-ci.txt @@ -0,0 +1,18 @@ +# Lightweight test environment, shared by .gitlab-ci.yml and +# .github/workflows/unit-tests.yml so the two pipelines cannot drift. +# +# Installed alongside the project itself (`uv pip install -e . --no-deps`), +# which keeps out the heavy core stack (nemo-skills, torch, ~200 packages). +# Tests that genuinely need nemo-skills call pytest.importorskip and are +# skipped in both pipelines; they run in the full-deps environment instead. +# +# Add new test-only dependencies here, not to an individual CI file. +pytest>=9.0.3 +pytest-cov>=4.1.0 +pytest-timeout>=2.2.0 +PyYAML +omegaconf +rich +orjson +pandas +pyarrow diff --git a/tests/test_aggregate_seeds.py b/tests/test_aggregate_seeds.py new file mode 100644 index 0000000..6594cd4 --- /dev/null +++ b/tests/test_aggregate_seeds.py @@ -0,0 +1,135 @@ +# 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.recipes.finance.utils.rl.aggregate_seeds. + +Pins the F2 contract: when ``--expected-seeds N`` is provided, missing +per-seed rollout files surface as a loud RuntimeError instead of +silently shrinking ``num_seeds`` in metrics.json. + +This is the second line of defense against the silent-success cascade +documented in the F1 commit message. The cluster's default Slurm dep +type is ``afterany`` (see cluster_configs/template-slurm.yaml note on +dependency_type), so a FAILED upstream merge does NOT prevent +aggregate from running. Without this validation, aggregate would +glob whatever ``output-rs*.jsonl`` files happened to be present and +proceed with the diminished set, producing partial difficulty data +that filter then passes through silently. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from nvflow.recipes.finance.utils.rl.aggregate_seeds import aggregate + + +def _write_seed_file(path: Path, num_rows: int = 3) -> None: + """Emit a minimal rollout file with the fields aggregate inspects.""" + rows = [ + {"uuid": f"q-{i}", "reward": float(i % 2), "question_type": "test"} for i in range(num_rows) + ] + path.write_text("".join(json.dumps(r) + "\n" for r in rows)) + + +def test_aggregate_passes_when_expected_seeds_matches_found(tmp_path: Path) -> None: + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + for s in range(3): + _write_seed_file(rollout_dir / f"output-rs{s}.jsonl") + out_dir = tmp_path / "agg" + + aggregate(str(rollout_dir), str(out_dir), expected_seeds=3) + + metrics = json.loads((out_dir / "metrics.json").read_text()) + assert metrics["num_seeds"] == 3 + + +def test_aggregate_raises_when_seed_missing(tmp_path: Path) -> None: + """Production scenario: rs0's merge job FAILED so output-rs0.jsonl is + absent, but rs1 and rs2 succeeded. Without --expected-seeds the + pre-F2 behaviour was to silently set num_seeds=2 and exit 0. With + F2 plumbed through (build_aggregate_cmd always passes + p.num_random_seeds), this raises so Slurm marks the aggregate job + FAILED -- visible signal that an upstream merge dropped a seed. + """ + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + # Only seeds 1 and 2 -- seed 0 missing. + _write_seed_file(rollout_dir / "output-rs1.jsonl") + _write_seed_file(rollout_dir / "output-rs2.jsonl") + out_dir = tmp_path / "agg" + + with pytest.raises(RuntimeError, match="Expected 3 per-seed rollout files"): + aggregate(str(rollout_dir), str(out_dir), expected_seeds=3) + + # No partial output should be written when the precondition fails. + assert not (out_dir / "metrics.json").exists() + assert not (out_dir / "summary.txt").exists() + + +def test_aggregate_raises_when_extra_seeds_present(tmp_path: Path) -> None: + """Symmetric guard: extra files (e.g. a stale output-rs7.jsonl from a + larger previous run) also fail the precondition. Otherwise an + operator could silently aggregate a mix of fresh + stale data. + """ + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + for s in range(5): # 5 seeds present but config expects 3 + _write_seed_file(rollout_dir / f"output-rs{s}.jsonl") + out_dir = tmp_path / "agg" + + with pytest.raises(RuntimeError, match="Expected 3 per-seed rollout files"): + aggregate(str(rollout_dir), str(out_dir), expected_seeds=3) + + +def test_aggregate_back_compat_no_expected_seeds(tmp_path: Path) -> None: + """Default expected_seeds=None preserves pre-F2 behaviour: aggregate + accepts whatever per-seed files exist. External callers (verify.py + rejudge path, ad-hoc scripts) that don't know the expected count + must continue to work. + """ + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + _write_seed_file(rollout_dir / "output-rs1.jsonl") + _write_seed_file(rollout_dir / "output-rs2.jsonl") + out_dir = tmp_path / "agg" + + aggregate(str(rollout_dir), str(out_dir)) # no expected_seeds kwarg + + metrics = json.loads((out_dir / "metrics.json").read_text()) + assert metrics["num_seeds"] == 2 # accepts the diminished set + + +def test_aggregate_excludes_chunk_and_async_files(tmp_path: Path) -> None: + """The seed-count comparison must be on canonical merged outputs + only. Per-chunk intermediates (output-rs0_chunk_0.jsonl) and + in-flight async files (output-rs0.jsonl-async) must be excluded + BEFORE the precondition check, otherwise stale intermediates + could mask a truly-missing seed. + """ + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + _write_seed_file(rollout_dir / "output-rs1.jsonl") + _write_seed_file(rollout_dir / "output-rs2.jsonl") + # Spurious intermediates that happen to glob-match. + _write_seed_file(rollout_dir / "output-rs0_chunk_0.jsonl") + _write_seed_file(rollout_dir / "output-rs0.jsonl-async") + out_dir = tmp_path / "agg" + + with pytest.raises(RuntimeError, match="found 2"): + aggregate(str(rollout_dir), str(out_dir), expected_seeds=3) diff --git a/tests/test_apply_validate_filter.py b/tests/test_apply_validate_filter.py new file mode 100644 index 0000000..7b6a05b --- /dev/null +++ b/tests/test_apply_validate_filter.py @@ -0,0 +1,196 @@ +# 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 ``apply_validate_filter`` post-S1. + +Pins the new contracts: +- ``raw_sdg_path`` is required (TypeError if omitted). +- ``--raw_sdg_source`` CLI flag is required (argparse exits non-zero). +- A VALID row whose ``problem`` is missing from the SDG file raises + :class:`MissingSdgRecordError` instead of silently emitting LLM bytes. +- An INVALID/missing-tag row whose ``problem`` is missing from the SDG + file is fine -- those rows go to the dropped stream and don't need + SDG-original bytes. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import orjson +import pytest + +from nvflow.recipes.finance.utils.rl.apply_validate_filter import ( + MissingSdgRecordError, + apply_validate_filter, +) + + +def _write_jsonl(path: Path, rows: list[dict]) -> None: + path.write_bytes(b"\n".join(orjson.dumps(r) for r in rows) + b"\n") + + +def _make_sdg_and_parsed( + tmp_path: Path, + *, + sdg_problems: list[str], + parsed: list[dict], +) -> tuple[Path, Path]: + sdg = tmp_path / "sdg.jsonl" + parsed_file = tmp_path / "parsed.jsonl" + _write_jsonl( + sdg, + [ + { + "problem": p, + "company_name": f"Co{p}", + "answer": f"A{p}", + "reasoning_content": f"R{p}", + } + for p in sdg_problems + ], + ) + _write_jsonl(parsed_file, parsed) + return sdg, parsed_file + + +# --------------------------------------------------------------------------- +# Required raw_sdg_path +# --------------------------------------------------------------------------- + + +def test_function_call_omitting_raw_sdg_path_raises_typeerror(tmp_path: Path) -> None: + """Calling the function without raw_sdg_path is a programming bug.""" + parsed_file = tmp_path / "parsed.jsonl" + parsed_file.write_bytes(b"") + with pytest.raises(TypeError, match="raw_sdg_path"): + apply_validate_filter( # type: ignore[call-arg] + input_file=str(parsed_file), + output_kept=str(tmp_path / "kept.jsonl"), + output_dropped=str(tmp_path / "dropped.jsonl"), + stats_file=str(tmp_path / "stats.json"), + ) + + +def test_cli_omitting_raw_sdg_source_exits_nonzero(tmp_path: Path) -> None: + """The CLI must require --raw_sdg_source so operators can't accidentally + fall back to a removed Mode B path. + """ + parsed_file = tmp_path / "parsed.jsonl" + parsed_file.write_bytes(b"") + proc = subprocess.run( + [ + sys.executable, + "-m", + "nvflow.recipes.finance.utils.rl.apply_validate_filter", + "--input_file", + str(parsed_file), + "--output_kept", + str(tmp_path / "kept.jsonl"), + "--output_dropped", + str(tmp_path / "dropped.jsonl"), + "--stats_file", + str(tmp_path / "stats.json"), + ], + capture_output=True, + text=True, + cwd=Path(__file__).resolve().parent.parent, + ) + assert proc.returncode != 0 + assert "--raw_sdg_source" in proc.stderr + + +# --------------------------------------------------------------------------- +# Missing-SDG-record contract +# --------------------------------------------------------------------------- + + +def test_valid_row_missing_from_sdg_raises(tmp_path: Path) -> None: + """A VALID row with no SDG counterpart is a Phase 1/2 mismatch -- raise.""" + sdg, parsed = _make_sdg_and_parsed( + tmp_path, + sdg_problems=["p_in_sdg"], + parsed=[ + {"problem": "p_NOT_in_sdg", "validate_tag": "VALID", "generation": "x"}, + ], + ) + with pytest.raises(MissingSdgRecordError, match="p_NOT_in_sdg"): + apply_validate_filter( + input_file=str(parsed), + output_kept=str(tmp_path / "kept.jsonl"), + output_dropped=str(tmp_path / "dropped.jsonl"), + stats_file=str(tmp_path / "stats.json"), + raw_sdg_path=str(sdg), + ) + + +def test_invalid_row_missing_from_sdg_is_fine(tmp_path: Path) -> None: + """Non-VALID rows go to the dropped stream regardless of SDG presence.""" + sdg, parsed = _make_sdg_and_parsed( + tmp_path, + sdg_problems=["p_in_sdg"], + parsed=[ + {"problem": "p_NOT_in_sdg", "validate_tag": "INVALID", "generation": "x"}, + {"problem": "p_NOT_in_sdg", "validate_tag": None, "generation": "y"}, + ], + ) + apply_validate_filter( + input_file=str(parsed), + output_kept=str(tmp_path / "kept.jsonl"), + output_dropped=str(tmp_path / "dropped.jsonl"), + stats_file=str(tmp_path / "stats.json"), + raw_sdg_path=str(sdg), + ) + stats = orjson.loads((tmp_path / "stats.json").read_bytes()) + assert stats["num_total"] == 2 + assert stats["num_kept"] == 0 + assert stats["num_dropped"] == 2 + # No VALID rows -> no missing-record exposure. + + +# --------------------------------------------------------------------------- +# Happy path: kept rows are byte-identical to SDG +# --------------------------------------------------------------------------- + + +def test_kept_rows_are_sdg_bytes_verbatim(tmp_path: Path) -> None: + sdg, parsed = _make_sdg_and_parsed( + tmp_path, + sdg_problems=["p1", "p2", "p3"], + parsed=[ + {"problem": "p1", "validate_tag": "VALID", "generation": "x"}, + {"problem": "p2", "validate_tag": "INVALID", "generation": "y"}, + {"problem": "p3", "validate_tag": "VALID", "generation": "z"}, + ], + ) + kept_file = tmp_path / "kept.jsonl" + apply_validate_filter( + input_file=str(parsed), + output_kept=str(kept_file), + output_dropped=str(tmp_path / "dropped.jsonl"), + stats_file=str(tmp_path / "stats.json"), + raw_sdg_path=str(sdg), + ) + + sdg_bytes_by_problem: dict[str, bytes] = {} + for line in sdg.read_bytes().splitlines(): + if line.strip(): + sdg_bytes_by_problem[orjson.loads(line)["problem"]] = line.strip() + + kept_lines = [line for line in kept_file.read_bytes().splitlines() if line.strip()] + assert len(kept_lines) == 2 + assert kept_lines[0] == sdg_bytes_by_problem["p1"] + assert kept_lines[1] == sdg_bytes_by_problem["p3"] diff --git a/tests/test_cli_cmd.py b/tests/test_cli_cmd.py new file mode 100644 index 0000000..f70c59f --- /dev/null +++ b/tests/test_cli_cmd.py @@ -0,0 +1,235 @@ +# 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.cli_cmd.build_python_cmd. + +Pins the shlex-quoting + interpreter contract so future refactors can't +accidentally reintroduce a path-injection vector via the rendered shell +command. This helper is shared across multiple stages +(validate_questions, data_transformation, apply_prompt_template, +convert_to_responses_api, prepare_data, prefetch_cache) so a regression +here would fan out across the entire data pipeline. +""" + +from __future__ import annotations + +import shlex +from pathlib import Path + +from nvflow.lib.cli_cmd import build_python_cmd, build_python_script_cmd + + +def test_emits_module_invocation() -> None: + out = build_python_cmd("foo.bar.baz", input_file="/a/b.jsonl") + assert out.startswith("python3 -m foo.bar.baz ") + # Should be a single space-joined string. + assert "\n" not in out + assert " " not in out # no doubled spaces + + +def test_renders_simple_paths_unchanged() -> None: + """Plain paths (no special chars) shouldn't gain extra quotes -- shlex.quote + only quotes when needed. This makes log-grepping the rendered command + straightforward. + """ + out = build_python_cmd( + "m", input_file="/lustre/foo/bar.jsonl", stats_file="/lustre/foo/stats.json" + ) + assert "--input_file /lustre/foo/bar.jsonl" in out + assert "--stats_file /lustre/foo/stats.json" in out + # No surrounding single quotes around the safe paths. + assert "'/lustre/foo/bar.jsonl'" not in out + + +def test_quotes_path_with_spaces() -> None: + """Spaces in a path must be properly quoted so the shell parses one arg.""" + out = build_python_cmd("m", input_file="/a path/with spaces.jsonl") + # shlex.quote surrounds the value with single quotes when needed. + assert "--input_file '/a path/with spaces.jsonl'" in out + # And shlex.split must round-trip to the original token. + tokens = shlex.split(out) + idx = tokens.index("--input_file") + assert tokens[idx + 1] == "/a path/with spaces.jsonl" + + +def test_quotes_path_with_single_quote() -> None: + """Single-quote in a value is the canonical injection vector for naive + f-string command builders -- shlex.quote handles it correctly. + """ + out = build_python_cmd("m", input_file="/a/file's name.jsonl") + tokens = shlex.split(out) + idx = tokens.index("--input_file") + assert tokens[idx + 1] == "/a/file's name.jsonl" + + +def test_quotes_shell_metacharacters() -> None: + """``$``, ``;``, ``|``, backticks, ``&`` etc. must not be interpreted by + the shell as control characters when they appear in a value. + """ + dangerous = "/path; rm -rf /;$(echo pwned)`evil`|cat&" + out = build_python_cmd("m", input_file=dangerous) + tokens = shlex.split(out) + idx = tokens.index("--input_file") + assert tokens[idx + 1] == dangerous + + +def test_accepts_pathlib_values() -> None: + """Stage code passes pathlib.Path -- helper must stringify them.""" + out = build_python_cmd("m", input_file=Path("/a/b.jsonl")) + assert "--input_file /a/b.jsonl" in out + + +def test_accepts_numeric_values() -> None: + """Stages pass ints (e.g. ``num_chunks=10``) and floats (e.g. + ``context_min_percentile=1.0``) directly -- helper must stringify them. + """ + out = build_python_cmd("m", num_chunks=10, context_min_percentile=1.0) + assert "--num_chunks 10" in out + assert "--context_min_percentile 1.0" in out + + +def test_preserves_flag_order() -> None: + """Ordered output keeps log-grep diffs minimal across reruns. Python + 3.7+ preserves kwarg order so this comes for free, but the test pins it. + """ + out = build_python_cmd("m", alpha="1", beta="2", gamma="3") + assert out.index("--alpha") < out.index("--beta") < out.index("--gamma") + + +def test_handles_empty_value() -> None: + """Empty string still needs quoting so the flag's value isn't lost.""" + out = build_python_cmd("m", input_file="") + tokens = shlex.split(out) + idx = tokens.index("--input_file") + assert tokens[idx + 1] == "" + + +def test_positional_args_emitted_before_flags() -> None: + """Positional inputs (e.g. dataset_transformer's input_files) must appear + between ``-m `` and the first flag, in argument order. + """ + out = build_python_cmd( + "m", + "/a/in1.jsonl", + "/a/in2.jsonl", + output_file="/a/out.jsonl", + ) + tokens = shlex.split(out) + assert tokens[:3] == ["python3", "-m", "m"] + assert tokens[3] == "/a/in1.jsonl" + assert tokens[4] == "/a/in2.jsonl" + assert tokens[5] == "--output_file" + assert tokens[6] == "/a/out.jsonl" + + +def test_positional_args_quoted() -> None: + """Positional args must use the same shlex.quote treatment as flag values + so ``input_files`` containing spaces or metacharacters do not break. + """ + out = build_python_cmd("m", "/a path/file.jsonl", "/b/'evil'.jsonl") + tokens = shlex.split(out) + assert tokens[3] == "/a path/file.jsonl" + assert tokens[4] == "/b/'evil'.jsonl" + + +def test_positional_args_accept_pathlib() -> None: + """Stage code may pass Path positionals -- they must be stringified.""" + out = build_python_cmd("m", Path("/a/b.jsonl"), output_file=Path("/a/out.jsonl")) + tokens = shlex.split(out) + # Layout: python3 -m m /a/b.jsonl --output_file /a/out.jsonl + # [0] [1] [2] [3] [4] [5] + assert tokens[3] == "/a/b.jsonl" + assert tokens[4] == "--output_file" + assert tokens[5] == "/a/out.jsonl" + + +def test_no_positional_no_flags_renders_bare_module() -> None: + """Bare module invocation (no args at all) is a valid edge case -- + the helper should not append trailing spaces or empty tokens. + """ + out = build_python_cmd("m") + assert out == "python3 -m m" + + +def test_uses_python3_interpreter_not_python() -> None: + """All cluster containers in this repo use ``python3`` -- standardising + avoids ambiguity around the unversioned ``python`` symlink (absent in + some minimal images). Pin this so future refactors don't silently + flip back to ``python``. + """ + out = build_python_cmd("any.module") + assert out.startswith("python3 ") + assert not out.startswith("python ") + + +# --- build_python_script_cmd (path-based variant for vendored external tools) + + +def test_script_cmd_emits_path_invocation() -> None: + """``python3