feat(gym): add Ultra environment prefetch config

docs/guides/nemotron-3-ultra.md

@@ -61,44 +61,42 @@ RLVR policy itself serves as the general teacher:
 
 ## Container
 
-Ultra runs on the stock NeMo RL container. The vLLM version is whatever `pyproject.toml` pins (currently the upstream
-aarch64 wheel), so the only requirement is an **aarch64 (arm64)** image for
-GB200 NVL72 nodes.
-
-The quickest option is to pull a prebuilt nightly image from
-[NGC](https://registry.ngc.nvidia.com/orgs/nvidia/containers/nemo-rl/tags) and
-skip the build entirely:
-
-```bash
-docker pull nvcr.io/nvidia/nemo-rl:<nightly-tag>
-```
-
-To build it yourself instead, from the root of the repo:
+Ultra uses vLLM and requires an **aarch64 (arm64)** image for GB200 NVL72
+nodes. Prebake the NeMo Gym virtual environments used by the recipes to avoid
+building them when each training job starts. From the root of the repository,
+build the image with the Gym virtual environments for the Ultra recipes:
 
 ```bash
 docker buildx build \
+  --platform linux/arm64 \
   --progress=plain \
   -f docker/Dockerfile \
   --target release \
-  -t nemo-rl-ultra:arm64 \
+  -t <your-registry>/nemo-rl:main-ultra-prefetched-venvs \
+  --push \
   --build-context nemo-rl=. \
   --build-arg MAX_JOBS=8 \
   --build-arg SKIP_SGLANG_BUILD=1 \
+  --build-arg SKIP_TRTLLM_BUILD=1 \
+  --build-arg NEMO_GYM_PREFETCH_CONFIGS="examples/nemo_gym/prefetch_ultra_all_envs.yaml" \
   .
 ```
 
 Build args:
+- `NEMO_GYM_PREFETCH_CONFIGS` — space-separated union configs whose Gym virtual
+  environments are baked into the image.
 - `SKIP_SGLANG_BUILD=1` — Ultra runs on vLLM; skip the SGLang build.
+- `SKIP_TRTLLM_BUILD=1` — Ultra does not use TensorRT-LLM; skip its build.
 - `MAX_JOBS` — parallel build jobs; tune to your machine.
 - `--build-context nemo-rl=.` — build from your local checkout (otherwise the
   Dockerfile pulls `NVIDIA-NeMo/RL.git#main`).
 
-
 To run on the cluster with Slurm, convert the image to a squashfs (`.sqsh`)
 with [enroot](https://github.com/NVIDIA/enroot):
 
 ```bash
-enroot import -o nemo-rl-container.sqsh dockerd://nemo-rl-ultra:arm64
+enroot import -o nemo-rl-container.sqsh \
+  docker://<your-registry>/nemo-rl:main-ultra-prefetched-venvs
 ```
 
 Pass the resulting image as `CONTAINER` in every launch command below (shown as

examples/nemo_gym/nemotron-3-ultra/mopd.yaml

@@ -453,7 +453,7 @@ env:
       - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
       - resources_servers/format_verification/configs/freeform_formatting.yaml
       - resources_servers/format_verification/configs/citation_format.yaml
-      - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml
+      - resources_servers/litmus_agent/configs/litmus_agent.yaml
       - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml
       - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml
       - resources_servers/swe_pivot/configs/swe_pivot.yaml
@@ -464,6 +464,11 @@ env:
       - resources_servers/format_verification/configs/citation_format.yaml
       - resources_servers/format_verification/configs/freeform_formatting.yaml
 
+    # The MOPD blend still uses the superseded rdkit_chemistry agent
+    # name. Keep that data reference while running the replacement Litmus agent.
+    rdkit_chemistry_agent:
+      _inherit_from: litmus_agent_agent
+
     # Increase the num workers to help with the serialization/deserialization overhead of token IDs
     policy_model:
       responses_api_models:
@@ -786,6 +791,7 @@ on_policy_distillation:
     lc_judge_simple_agent:                   ${_teachers.general}
     ns_tools_simple_agent:                   ${_teachers.reasoning}
     code_gen_simple_agent:                   ${_teachers.reasoning}
+    rdkit_chemistry_agent:                   ${_teachers.general}
 
     # IFBench / abstention / multichallenge (all share IFBench teacher)
     instruction_following_simple_agent:      ${_teachers.ifbench}

examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml

@@ -427,7 +427,6 @@ env:
       - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
       - resources_servers/format_verification/configs/freeform_formatting.yaml
       - resources_servers/format_verification/configs/citation_format.yaml
-      - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml
       - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml
       - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml
 

examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml

@@ -431,10 +431,15 @@ env:
       - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
       - resources_servers/format_verification/configs/freeform_formatting.yaml
       - resources_servers/format_verification/configs/citation_format.yaml
-      - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml
+      - resources_servers/litmus_agent/configs/litmus_agent.yaml
       - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml
       - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml
 
+    # The RLVR2 blend still uses the superseded rdkit_chemistry agent
+    # name. Keep that data reference while running the replacement Litmus agent.
+    rdkit_chemistry_agent:
+      _inherit_from: litmus_agent_agent
+
     policy_model:
       responses_api_models:
         vllm_model:

examples/nemo_gym/prefetch_ultra_all_envs.yaml

@@ -0,0 +1,132 @@
+# Prefetch config: union of NeMo Gym environments across the nemotron-3-ultra stages
+# (student RLVR + teacher stages + SWE + MOPD). Pass to the docker build to bake every
+# available Gym server venv into the image:
+#   docker buildx build -f docker/Dockerfile \
+#     --build-arg NEMO_GYM_PREFETCH_CONFIGS="examples/nemo_gym/prefetch_ultra_all_envs.yaml" .
+# Only env.nemo_gym is read (dry-run NemoGym._spinup); the dummy values below satisfy the bundled
+# configs' interpolations during the dry run. swerl_gen is omitted because it is disabled in the
+# Ultra recipes.
+env:
+  should_use_nemo_gym: true
+  nemo_gym:
+    skip_venv_if_present: true
+    port_range_low: 5000
+    port_range_high: 5999
+
+    # Dummy interpolation values for the bundled configs (unused in the dry run).
+    policy_model_name: dummy-model
+    policy_base_url: http://localhost:8000/v1
+    policy_api_key: dummy
+    judge_model_name: dummy-judge
+    judge_base_url: http://localhost:8000/v1
+    judge_api_key: dummy
+    genrm_model_name: dummy-genrm
+    anthropic_model_name: dummy-anthropic
+    anthropic_base_url: http://localhost:8000/v1
+    anthropic_api_key: dummy
+
+    invalid_tool_call_patterns:
+      - "<tool_call>"
+      - "</tool_call>"
+      - "<function_call>"
+      - "</function_call>"
+    thinking_tags:
+      - "<think>"
+      - "</think>"
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - resources_servers/math_with_judge/configs/math_with_judge.yaml
+    - resources_servers/code_gen/configs/code_gen.yaml
+    - resources_servers/workplace_assistant/configs/workplace_assistant.yaml
+    - resources_servers/mcqa/configs/mcqa.yaml
+    - resources_servers/instruction_following/configs/instruction_following.yaml
+    - resources_servers/equivalence_llm_judge/configs/lc_judge.yaml
+    - resources_servers/calendar/configs/calendar.yaml
+    - resources_servers/genrm_compare/configs/genrm_compare.yaml
+    - resources_servers/equivalence_llm_judge/configs/nl2bash-equivalency.yaml
+    - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml
+    - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml
+    - resources_servers/reasoning_gym/configs/reasoning_gym.yaml
+    - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml
+    - resources_servers/ns_tools/configs/ns_tools.yaml
+    - resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml
+    - resources_servers/multichallenge/configs/multichallenge.yaml
+    - resources_servers/inverse_if/configs/inverse_if.yaml
+    - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml
+    - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml
+    - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml
+    - resources_servers/abstention/configs/abstention.yaml
+    - resources_servers/nvarc/configs/inductive.yaml
+    - resources_servers/nvarc/configs/transductive.yaml
+    - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml
+    - resources_servers/equivalence_rule/configs/lc.yaml
+    - resources_servers/ether0/configs/ether0.yaml
+    - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml
+    - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml
+    - resources_servers/format_verification/configs/freeform_formatting.yaml
+    - resources_servers/format_verification/configs/citation_format.yaml
+    - resources_servers/litmus_agent/configs/litmus_agent.yaml
+    - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml
+    - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml
+    - resources_servers/swe_pivot/configs/swe_pivot.yaml
+    - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_stirrup.yaml
+    - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_agent006.yaml
+    - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_opencode.yaml
+    - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_codex.yaml
+    - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml
+
+    # Ultra data still names the removed rdkit_chemistry agent. This alias keeps
+    # that data compatible while prefetching the replacement Litmus environment.
+    rdkit_chemistry_agent:
+      _inherit_from: litmus_agent_agent
+
+    # Thin vllm_model clients for the judge refs the bundled configs expect (so refs resolve and
+    # the vllm_model venv builds). genrm_model is already defined by the bundled genrm_compare.yaml.
+    safety_judge_model:
+      responses_api_models:
+        vllm_model:
+          entrypoint: app.py
+          base_url: http://localhost:8000/v1
+          api_key: dummy
+          model: dummy-judge
+          spinup_server: false
+    judge_model:
+      responses_api_models:
+        vllm_model:
+          entrypoint: app.py
+          base_url: http://localhost:8000/v1
+          api_key: dummy
+          model: dummy-judge
+          spinup_server: false
+    # Thin client for the *_reasoning_off agent refs.
+    policy_model_reasoning_off:
+      responses_api_models:
+        vllm_model:
+          entrypoint: app.py
+          base_url: http://localhost:8000/v1
+          api_key: dummy
+          model: dummy-model
+          spinup_server: false
+    # Standalone server defined only to bake the local_vllm_model venv (the runtime GPU judges
+    # use this type). dry_run builds the venv without loading a model.
+    prefetch_local_vllm_model:
+      responses_api_models:
+        local_vllm_model:
+          entrypoint: app.py
+          model: dummy-model
+          vllm_serve_kwargs:
+            tensor_parallel_size: 1
+            data_parallel_size: 1
+            pipeline_parallel_size: 1
+
+    # swe_agents marks container_formatter/dataset_path mandatory; dummies satisfy the dry run.
+    swe_agents_train:
+      responses_api_agents:
+        swe_agents:
+          container_formatter: dummy
+          dataset_path: /dev/null
+    swe_agents_val:
+      responses_api_agents:
+        swe_agents:
+          container_formatter: dummy
+          dataset_path: /dev/null