diff --git a/docker-bake.hcl b/docker-bake.hcl index 6f8b4f7db5..58ac77f624 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -128,7 +128,7 @@ variable "NEMO_RL_REPO" { # RL pins Gym as a git submodule (-> soluwalana/Gym over https), so Gym rides in with the RL git ADD # - no separate Gym pin needed. variable "NEMO_RL_REF" { - default = "874f94737d5358680607735e339a8c06c01495f7" # soluwalana/RL nmp/customizer + default = "eab835b111b78d90c3629b26af9b3f12dc0d0d29" # soluwalana/RL nmp/customizer } variable "RL_BASE_CONTEXT" { default = "" diff --git a/docs/customizer/grpo-training.mdx b/docs/customizer/grpo-training.mdx index b4e5084565..9d9ff85fb6 100644 --- a/docs/customizer/grpo-training.mdx +++ b/docs/customizer/grpo-training.mdx @@ -84,7 +84,14 @@ There are two installs at job start: | `wheels-v1` | The package's `wheels/`, for anything it vendors | No, **if the closure also covers step 1** | | `adapter-wheels-v1` | The package's `wheels/`, plus the agent harness's own requirements | Yes — see below | -A `wheels-v1` job runs with no egress only when `wheels/` carries the server's full requirement closure, `nemo-gym` at exactly the training-image version, and Gym's pinned `ray[default]` and `openai` — the venv is seeded empty, so nothing carries over from the image. Completeness of `wheels/` is what makes an offline run work, not the format name. Full detail: [GRPO Environment Packages](/documentation/customizer-reference/tutorials/grpo-environment-packages#dependencies). +A `wheels-v1` job runs with no egress only when `wheels/` carries all of: + +- the server's own requirement closure; +- `nemo-gym` at exactly the training-image version; +- `ray[default]` and `openai` at the versions the image pins, which Gym appends to every install; +- **`pip`** — `uv venv --seed` installs it into each new virtualenv *before* any dependency install, resolving it like any other package. + +The virtualenv is created empty, so nothing carries over from the image. Completeness of `wheels/` is what makes an offline run work, not the format name. Full detail: [GRPO Environment Packages](/documentation/customizer-reference/tutorials/grpo-environment-packages#dependencies). `adapter-wheels-v1` requires network access when the job starts. The package `wheels/` directory covers the hub environment; `verifiers_agent` also installs `verifiers` from GitHub. @@ -105,7 +112,7 @@ ARCH=x86_64 # amd64 nodes The interpreter is **Python 3.13** on every architecture. A closure resolved for the wrong architecture passes `--validate-only`, which checks layout rather than wheel tags, and then fails on the cluster with `has no wheels with a matching platform tag`. ```bash -pip download --dest my-env/wheels \ +pip download pip --dest my-env/wheels \ --only-binary=:all: \ --python-version 3.13 \ --platform "manylinux_2_39_$ARCH" \ diff --git a/docs/customizer/tutorials/grpo-environment-packages.mdx b/docs/customizer/tutorials/grpo-environment-packages.mdx index d8c21fd03e..80f1393f1b 100644 --- a/docs/customizer/tutorials/grpo-environment-packages.mdx +++ b/docs/customizer/tutorials/grpo-environment-packages.mdx @@ -122,6 +122,7 @@ For a `wheels-v1` job to start with no egress, `wheels/` must contain: - everything the server’s `requirements.txt` names, and their transitive dependencies; - `nemo-gym` at exactly the version the training image reports; - `ray[default]` and `openai` at the versions the image pins — the virtualenv is created with `uv venv --seed`, so nothing is inherited from the image. +- **`pip`** — `uv venv --seed` installs it into each new virtualenv *before* any `uv pip install` runs, resolving it like any other package. Without it in `wheels/` the seed step reaches for an index and the job fails with `Failed to install seed packages into virtual environment` / `No solution found when resolving: pip`. `uv venv` honours `UV_FIND_LINKS`, so vendoring the wheel is all that is needed. Python 3.13 seeds only `pip`. Completeness of `wheels/` is what makes a sandboxed run work; the format name alone does not. @@ -370,7 +371,7 @@ pip download --dest weather-env/wheels \ --platform "manylinux_2_28_$ARCH" \ --platform "manylinux_2_17_$ARCH" \ --platform "manylinux2014_$ARCH" \ - "nemo-gym==$GYM_VERSION" "ray[default]==$RAY_VERSION" "openai==$OPENAI_VERSION" \ + "nemo-gym==$GYM_VERSION" "ray[default]==$RAY_VERSION" "openai==$OPENAI_VERSION" pip \ -r resources_servers/weather_verifier/requirements.txt uv run --package nmp-rl pi-to-gym-conversion --validate-only ./weather-env ``` diff --git a/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/gym-environments.md b/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/gym-environments.md index aa0a09712e..5aca727406 100644 --- a/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/gym-environments.md +++ b/plugins/nemo-customizer/src/nemo_customizer/skills/nemo-customizer/references/gym-environments.md @@ -127,6 +127,7 @@ So `--no-index` is real, but it is step 2 only. **A job is offline-clean only wh - everything the server's `requirements.txt` names, and their transitive deps; - **`nemo-gym` at exactly the version the training image reports** (below); - `ray[default]` and `openai` at the versions Gym pins as head-server deps — the sub-venv starts from `uv venv --seed`, so nothing is inherited from the image. +- **`pip`** — `uv venv --seed` installs it into each new virtualenv *before* any `uv pip install` runs, resolving it like any other package. Without it in `wheels/` the seed step reaches for an index and the job fails with `Failed to install seed packages into virtual environment` / `No solution found when resolving: pip`. `uv venv` honours `UV_FIND_LINKS`, so vendoring the wheel is all that is needed. Python 3.13 seeds only `pip`. Completeness of `wheels/` is what makes a sandboxed run work. The format name alone does not. @@ -485,7 +486,7 @@ pip download --dest my-env/wheels \ --platform "manylinux_2_28_$ARCH" \ --platform "manylinux_2_17_$ARCH" \ --platform "manylinux2014_$ARCH" \ - "nemo-gym==$GYM_VERSION" "ray[default]==$RAY_VERSION" "openai==$OPENAI_VERSION" \ + "nemo-gym==$GYM_VERSION" "ray[default]==$RAY_VERSION" "openai==$OPENAI_VERSION" pip \ -r resources_servers/my_env/requirements.txt uv run --package nmp-rl pi-to-gym-conversion --validate-only ./my-env ``` diff --git a/scripts/grpo-examples/gym_to_env_package.py b/scripts/grpo-examples/gym_to_env_package.py index fb914b2648..639757d1d5 100644 --- a/scripts/grpo-examples/gym_to_env_package.py +++ b/scripts/grpo-examples/gym_to_env_package.py @@ -28,6 +28,7 @@ import argparse import json +import re import shutil import subprocess import sys @@ -201,6 +202,47 @@ def build_fork_wheel(gym_root: Path, wheels: Path, expect_version: str | None) - return target +_NO_WHEEL_RE = re.compile(r"No matching distribution found for (\S+)") + + +def _build_pure_wheel(requirement: str, wheels: Path) -> None: + """Build a wheel for a pin that publishes none, e.g. antlr4-python3-runtime. + + Only pure-Python projects are safe to build here: the build host is not the target, so + anything compiling an extension would produce a wheel for the wrong platform. + """ + before = set(wheels.glob("*.whl")) + subprocess.run( + [sys.executable, "-m", "pip", "wheel", "--no-deps", "--no-cache-dir", "--wheel-dir", str(wheels), requirement], + check=True, + ) + for built in set(wheels.glob("*.whl")) - before: + if not built.stem.endswith("-py3-none-any") and not built.stem.endswith("-py2.py3-none-any"): + built.unlink() + raise SystemExit( + f"{requirement} has no wheel on the index and does not build a pure-Python one " + f"({built.name}). Building it here would target the build host, not the cluster." + ) + print(f"Built from sdist: {built.name}", flush=True) + + +def _download_with_sdist_fallback(download_cmd: list[str], wheels: Path, max_builds: int = 8) -> None: + """Run pip download, building any pin that publishes no wheel, then retrying.""" + builds = 0 + while True: + result = subprocess.run(download_cmd, stderr=subprocess.PIPE, text=True) + if result.returncode == 0: + return + sys.stderr.write(result.stderr) + match = _NO_WHEEL_RE.search(result.stderr) + if not match: + raise SystemExit("pip download failed; see the error above.") + if builds >= max_builds: + raise SystemExit(f"still missing wheels after building {max_builds} package(s).") + _build_pure_wheel(match.group(1), wheels) + builds += 1 + + def vendor_wheels( out_dir: Path, gym_root: Path, @@ -210,52 +252,104 @@ def vendor_wheels( ray_version: str, openai_version: str, ) -> Path: - """Resolve the closure from nemo-gym, Gym's head-server deps, and the server's requirements. - - Gym appends ``ray[default]==`` and ``openai==`` (read from the Gym host process) to every - per-server venv install, and the venv is created with ``uv venv --seed``, so nothing carries - over from the image. Omitting them leaves the closure incomplete and the venv build reaches - for an index -- which fails on a cluster with no egress, the case wheels-v1 exists for. + """Resolve and download the offline closure for the training nodes. + + Resolution and download are separate steps. ``uv pip compile --python-platform`` picks + the versions, because environment markers are evaluated during resolution: resolving on + the build host omits a linux-only dependency (sqlalchemy's greenlet) and can add a + darwin-only one. ``pip download`` then fetches exactly those pins for the target tags. + + Beyond the server's own requirements the closure needs: + + * ``nemo-gym[dev]`` at the image's version, built from ``gym_root`` so a fork is not + replaced by the same-versioned upstream release. The ``dev`` extra is required because + the image's own servers depend on it; + * ``ray[default]`` and ``openai`` at the image's versions, which Gym appends to every + per-server install; + * ``pip``, installed by ``uv venv --seed`` into each venv before anything else; + * ``setuptools`` and ``setuptools-scm``, Gym's ``build-system.requires``. Servers that + resolve to the Gym tree take uv's editable branch and build ``nemo-gym`` from source, + which needs a PEP 517 build environment. """ wheels = out_dir / "wheels" # Rebuild from empty: pip copies by filename, so a wheel from an earlier run survives - # whenever the new resolution picks a different version, leaving two versions of one - # distribution and no guarantee about which gets installed. + # whenever the new resolution picks a different version. if wheels.exists(): shutil.rmtree(wheels) wheels.mkdir(parents=True) fork_wheel = build_fork_wheel(gym_root, wheels, expect_version) - cmd = [ - sys.executable, - "-m", + requirements = [ + # The [dev] extra, not the bare wheel: servers that resolve into the Gym tree install + # `nemo-gym[dev]` (its own requirements.txt, and vllm_model's pyproject), so the + # closure has to cover pre-commit, mypy, ruff and the pytest set too. + f"nemo-gym[dev] @ file://{fork_wheel}", + f"ray[default]=={ray_version}", + f"openai=={openai_version}", "pip", - "download", - "--dest", - str(wheels), - "--no-cache-dir", - "--only-binary", - ":all:", - "--python-version", - TARGET_PYTHON_VERSION, + "setuptools>=61", + "setuptools-scm", ] - for tag in (f"manylinux_2_39_{arch}", f"manylinux_2_28_{arch}", f"manylinux_2_17_{arch}", f"manylinux2014_{arch}"): - cmd += ["--platform", tag] - cmd += [str(fork_wheel), f"ray[default]=={ray_version}", f"openai=={openai_version}"] - # The sub-venv is created with `uv venv --seed`, so everything the server imports must be - # present. Its own requirements.txt names those; the editable nemo-gym line is not a real - # requirement outside a checkout, and Gym rewrites it at spin-up. reqs = pkg_server_dir / "requirements.txt" if reqs.is_file(): - kept = [ln for ln in reqs.read_text(encoding="utf-8").splitlines() if ln.strip() and "../.." not in ln] - if kept: - with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as fh: - fh.write("\n".join(kept) + "\n") - cmd += ["-r", fh.name] - print("Running:", " ".join(cmd), flush=True) - subprocess.run(cmd, check=True) - - stray = [p.name for p in wheels.iterdir() if p.is_file() and p.suffix != ".whl"] + # The editable nemo-gym line is meaningless outside a checkout; Gym rewrites it. + requirements += [ln for ln in reqs.read_text(encoding="utf-8").splitlines() if ln.strip() and "../.." not in ln] + + with tempfile.TemporaryDirectory(prefix="closure-") as tmp: + req_in = Path(tmp) / "requirements.in" + req_in.write_text("\n".join(requirements) + "\n", encoding="utf-8") + pinned = Path(tmp) / "requirements.txt" + compile_cmd = [ + "uv", + "pip", + "compile", + str(req_in), + "--output-file", + str(pinned), + "--no-header", + # Ignore the ambient project's [tool.uv] policy, which would repin the closure. + "--no-config", + "--python-platform", + f"{arch}-unknown-linux-gnu", + "--python-version", + TARGET_PYTHON_VERSION, + ] + print("Running:", " ".join(compile_cmd), flush=True) + subprocess.run(compile_cmd, check=True) + + download_cmd = [ + sys.executable, + "-m", + "pip", + "download", + "--dest", + str(wheels), + "--no-cache-dir", + "--only-binary", + ":all:", + # Pick up anything built from an sdist below. + "--find-links", + str(wheels), + "--python-version", + TARGET_PYTHON_VERSION, + # pip defaults this to the build host's interpreter; the target is CPython + # regardless of what runs this script. The ABI set is left to pip, which + # derives it from the implementation and version. + "--implementation", + "cp", + ] + for tag in ( + f"manylinux_2_39_{arch}", + f"manylinux_2_28_{arch}", + f"manylinux_2_17_{arch}", + f"manylinux2014_{arch}", + ): + download_cmd += ["--platform", tag] + download_cmd += ["-r", str(pinned)] + print("Running:", " ".join(download_cmd), flush=True) + _download_with_sdist_fallback(download_cmd, wheels) + + stray = [f.name for f in wheels.iterdir() if f.is_file() and f.suffix != ".whl"] if stray: raise SystemExit(f"wheels/ must contain only .whl files, got: {stray}") return wheels diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py index ec6622fd5e..dda52c1862 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_config.py @@ -185,18 +185,33 @@ def _count_jsonl_rows(path: Path) -> int: return count -def _read_manifest_config_paths(environment_path: str | None) -> list[str]: - """Read ``config_paths`` from the environment package manifest on job storage.""" +# Formats whose wheels/ is a complete closure, so the job can resolve without an index. +# adapter-wheels-v1 ships wheels too, but its agent harness still installs from GitHub. +OFFLINE_ENVIRONMENT_FORMATS = frozenset({"wheels-v1"}) + + +def _read_manifest(environment_path: str | None) -> dict: + """Load the environment package manifest from job storage, or {} when absent.""" if not environment_path: - return [] + return {} manifest_path = Path(environment_path) / "nemo-environment.yaml" if not manifest_path.is_file(): logger.warning("No nemo-environment.yaml at %s; Gym will start with no config_paths", environment_path) - return [] + return {} manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) - if not isinstance(manifest, dict) or not manifest.get("config_paths"): - return [] - return list(manifest["config_paths"]) + return manifest if isinstance(manifest, dict) else {} + + +def _environment_is_offline(environment_path: str | None) -> bool: + """Whether the package promises a self-sufficient wheelhouse.""" + + return _read_manifest(environment_path).get("format") in OFFLINE_ENVIRONMENT_FORMATS + + +def _read_manifest_config_paths(environment_path: str | None) -> list[str]: + """Read ``config_paths`` from the environment package manifest on job storage.""" + manifest = _read_manifest(environment_path) + return list(manifest["config_paths"]) if manifest.get("config_paths") else [] def _resolve_gym_paths( @@ -264,14 +279,25 @@ def _build_nemo_gym_env_config( # FileNotFoundError on e.g. /opt/nemo-rl/configs/.yaml. So anchor them to wherever # the package is actually visible to the process that loads them: the sandbox mount in # mode B, the job-storage copy in mode A. Absolute entries are passed through untouched. - package_root = SANDBOX_ENVIRONMENT_PATH if sandboxed else (gym.environment_path or DEFAULT_ENVIRONMENT_PATH) + # Where the manifest is readable from THIS process, which is the job-storage copy in + # both modes. Distinct from package_root below, which is where the servers will see it. + manifest_root = gym.environment_path or DEFAULT_ENVIRONMENT_PATH + # What the Gym host will treat as the package root in mode B. NeMo-RL reads + # `sandboxed.environment_path` and only falls back to the mount, so anchoring + # config_paths to the bare constant would point them somewhere Gym never looks. + sandbox_root = gym.sandbox_environment_path or SANDBOX_ENVIRONMENT_PATH + package_root = sandbox_root if sandboxed else manifest_root config_paths = [ path if Path(path).is_absolute() else str(Path(package_root) / path) - for path in _read_manifest_config_paths(gym.environment_path) + for path in _read_manifest_config_paths(manifest_root) ] if config_paths: nemo_gym["config_paths"] = config_paths + offline_environment = _environment_is_offline(manifest_root) + if offline_environment: + nemo_gym["environment_offline"] = True + if sandboxed: # Fail loudly rather than substituting a default here: a wrong sandbox # image surfaces as an opaque client-side ReadTimeout once the sandbox pod @@ -339,7 +365,8 @@ def _build_nemo_gym_env_config( sandbox_cfg = NemoGymSandboxedConfig( sandboxed=True, host_provider="opensandbox", - environment_path=gym.sandbox_environment_path or SANDBOX_ENVIRONMENT_PATH, + environment_path=sandbox_root, + environment_offline=offline_environment, job_id=job_ctx.job_id, sandbox=sandbox, ) diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/sandbox_config.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/sandbox_config.py index d397335361..738db34cb8 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/sandbox_config.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/sandbox_config.py @@ -93,6 +93,9 @@ class NemoGymSandboxedConfig(BaseModel): sandboxed: bool = True host_provider: str = "opensandbox" environment_path: str | None = None + # Whether that package vendors a complete closure. NeMo-RL forwards it to the sandbox as + # NMP_ENVIRONMENT_OFFLINE so uv resolves from wheels/ instead of an index it cannot reach. + environment_offline: bool = False sandbox: SandboxConfig | None = None # Stamped onto every sandbox pod as a label. Upstream defaults this to a shared # constant when absent, which makes concurrent jobs indistinguishable.