From 0341f3be03f197656eae7f062953419c4bb03fbc Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 22:01:42 -0700 Subject: [PATCH 01/16] test: add Gym rollout acceptance matrix Signed-off-by: yaoyu-33 --- .../run_nemo_gym_single_node_sanity_tests.sh | 13 +- .../p0_rollout_acceptance.yaml | 70 ++++ .../test_nemo_gym_rollout_acceptance.py | 359 ++++++++++++++++++ 3 files changed, 440 insertions(+), 2 deletions(-) create mode 100644 tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml create mode 100644 tests/unit/environments/test_nemo_gym_rollout_acceptance.py diff --git a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh index b12905694dc..b6dc544abb0 100755 --- a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh +++ b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh @@ -41,7 +41,16 @@ uv run python -c "import ray; ray.shutdown()" ./tests/run_unit.sh unit/environments/test_math_environment.py::test_math_env_step_basic # NeMo Gym integrates directly into NeMo RL as an Environment since that is the cleanest way. This tests the NeMo Gym integration logic and correctness. -./tests/run_unit.sh unit/environments/test_nemo_gym.py::test_nemo_gym_sanity +./tests/run_unit.sh \ + unit/environments/test_nemo_gym.py::test_nemo_gym_sanity \ + --nemo-gym-only # NeMo Gym uses a separate rollout loop inside grpo_train in NeMo RL. This tests the e2e rollout functionality and correctness. -./tests/run_unit.sh unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout +./tests/run_unit.sh \ + unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout \ + --nemo-gym-only + +# Exercise one pinned Gym example from each P0 environment through the real NeMo RL rollout boundary. +./tests/run_unit.sh \ + unit/environments/test_nemo_gym_rollout_acceptance.py::test_p0_gym_environments_roll_out_through_nemo_rl \ + --nemo-gym-only diff --git a/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml new file mode 100644 index 00000000000..3e6e1443929 --- /dev/null +++ b/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml @@ -0,0 +1,70 @@ +cases: + - name: math_with_judge + config_path: resources_servers/math_with_judge/configs/math_with_judge.yaml + data_path: resources_servers/math_with_judge/data/example.jsonl + example_index: 0 + example_sha256: 5e8d77674814a6b1fdcd0f2dbef0c8597cc109814763056070ff2eeb0c9d70ef + agent_ref: + type: responses_api_agents + name: math_with_judge_simple_agent + expected_generation: '\boxed{32}' + expected_prompt_fragment: 1000 digit numbers + expected_reward: 1.0 + expected_result: + extracted_answer: "32" + library_reward: 1.0 + + - name: code_gen + config_path: resources_servers/code_gen/configs/code_gen.yaml + data_path: resources_servers/code_gen/data/example.jsonl + example_index: 4 + example_sha256: ceea8ed47c6615783bc9df8140110b2aa86729438448ee04e74c5d4850ebff24 + agent_ref: + type: responses_api_agents + name: code_gen_simple_agent + expected_generation: |- + ```python + import sys + + values = iter(map(int, sys.stdin.buffer.read().split())) + test_count = next(values) + answers = [] + for _ in range(test_count): + n = next(values) + positions = [0] * (n + 1) + for index in range(n): + positions[next(values)] = index + left = right = positions[1] + bits = [] + for value in range(1, n + 1): + left = min(left, positions[value]) + right = max(right, positions[value]) + bits.append("1" if right - left + 1 == value else "0") + answers.append("".join(bits)) + sys.stdout.write("\n".join(answers)) + ``` + expected_prompt_fragment: beautiful number + expected_reward: 1.0 + expected_result: + result: + - true + - true + - true + + - name: single_step_tool_use_with_argument_comparison + config_path: resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + data_path: resources_servers/single_step_tool_use_with_argument_comparison/data/example.jsonl + example_index: 4 + example_sha256: 2e24731dfeda02e16351dd5abd5c1e3fd8307554f4824cf4fb3630672c63f5ca + agent_ref: + type: responses_api_agents + name: single_step_tool_use_with_argument_comparison_agent + expected_generation: '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' + expected_prompt_fragment: SHOW24 + expected_reward: 1.0 + expected_result: + category: A tool call that matches the expected tool call was found + expected_action: + type: function_call + name: check_seat_availability + arguments: '{"event_id": "SHOW24", "section": "Medical Zone"}' diff --git a/tests/unit/environments/test_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_nemo_gym_rollout_acceptance.py new file mode 100644 index 00000000000..916ee653870 --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_rollout_acceptance.py @@ -0,0 +1,359 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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. + +import json +import math +from copy import deepcopy +from hashlib import sha256 +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from threading import Thread +from typing import Any +from zlib import compress, decompress + +import pytest +import ray +import torch +from yaml import safe_load + +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.datasets.response_datasets import NemoGymDataset +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data.processors import nemo_gym_data_processor +from nemo_rl.distributed.virtual_cluster import _get_node_ip_local +from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor +from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync + +_REPO_ROOT = Path(__file__).parents[3] +_GYM_ROOT = _REPO_ROOT / "3rdparty/Gym-workspace/Gym" +_CASES_PATH = Path(__file__).parent / "nemo_gym_test_data/p0_rollout_acceptance.yaml" +_POLICY_MODEL_CONFIG = ( + "responses_api_models/vllm_model/configs/vllm_model_for_training.yaml" +) +_REQUIRED_P0_CASES = { + "code_gen", + "math_with_judge", + "single_step_tool_use_with_argument_comparison", +} +_GENERATION_CONFIG = { + "backend": "test", + "max_new_tokens": 1024, + "max_total_sequence_length": 16384, + "temperature": 0.0, + "top_p": 1.0, + "top_k": None, + "stop_token_ids": None, + "stop_strings": None, +} +_CODE_GEN_COMPLETION = """```python +import sys + +values = iter(map(int, sys.stdin.buffer.read().split())) +test_count = next(values) +answers = [] +for _ in range(test_count): + n = next(values) + positions = [0] * (n + 1) + for index in range(n): + positions[next(values)] = index + left = right = positions[1] + bits = [] + for value in range(1, n + 1): + left = min(left, positions[value]) + right = max(right, positions[value]) + bits.append("1" if right - left + 1 == value else "0") + answers.append("".join(bits)) +sys.stdout.write("\\n".join(answers)) +```""" +_TOOL_CALL_ARGUMENTS = '{"event_id": "SHOW24", "section": "Medical Zone"}' +_TOOL_CALL_GENERATION = ( + '{"name":"check_seat_availability","arguments":' + '{"event_id":"SHOW24","section":"Medical Zone"}}' +) + + +class _ScriptedPolicyGeneration: + """Minimum policy-generation surface consumed by the Gym rollout path.""" + + cfg = _GENERATION_CONFIG + + +class _ScriptedTokenizer: + """Reversibly decode request and generation metadata without model weights.""" + + pad_token_id = 256 + + def batch_decode(self, batch: list[list[int]]) -> list[str]: + decoded = [] + for token_ids in batch: + payload = bytes(int(token_id) for token_id in token_ids) + if payload.startswith(b"\x00"): + decoded.append(decompress(payload[1:]).decode()) + elif payload.startswith(b"\x01"): + decoded.append(payload[1:].decode()) + else: + raise ValueError("Unrecognized scripted token sequence") + return decoded + + +class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): + """Return deterministic policy outputs to real Gym model and agent services.""" + + protocol_version = "HTTP/1.1" + + @staticmethod + def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: + serialized_body = json.dumps(body) + if "1000 digit numbers" in serialized_body: + content = r"\boxed{32}" + return {"role": "assistant", "content": content}, content + if "python programming language only" in serialized_body: + return { + "role": "assistant", + "content": _CODE_GEN_COMPLETION, + }, _CODE_GEN_COMPLETION + if "SHOW24" in serialized_body and "Medical Zone" in serialized_body: + return ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-p0-acceptance", + "type": "function", + "function": { + "name": "check_seat_availability", + "arguments": _TOOL_CALL_ARGUMENTS, + }, + } + ], + }, + _TOOL_CALL_GENERATION, + ) + raise AssertionError("P0 scripted policy received an unrecognized prompt") + + def do_POST(self) -> None: # noqa: N802 + if self.path != "/v1/chat/completions": + self.send_error(404) + return + + content_length = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(content_length)) + message, generation_text = self._message_for(body) + serialized_request = json.dumps( + body, sort_keys=True, separators=(",", ":") + ).encode() + prompt_token_ids = list(b"\x00" + compress(serialized_request)) + generation_token_ids = list(b"\x01" + generation_text.encode()) + message.update( + { + "prompt_token_ids": prompt_token_ids, + "generation_token_ids": generation_token_ids, + "generation_log_probs": [-0.1] * len(generation_token_ids), + } + ) + response = { + "id": "chatcmpl-p0-acceptance", + "object": "chat.completion", + "created": 0, + "model": body.get("model", "scripted-model"), + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls" + if message.get("tool_calls") + else "stop", + "message": message, + } + ], + "usage": { + "prompt_tokens": len(prompt_token_ids), + "completion_tokens": len(generation_token_ids), + "total_tokens": len(prompt_token_ids) + len(generation_token_ids), + }, + } + payload = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: Any) -> None: + pass + + +def _load_acceptance_cases() -> list[dict[str, Any]]: + with _CASES_PATH.open() as case_file: + cases = safe_load(case_file)["cases"] + + assert cases, "P0 Gym rollout acceptance matrix must not be empty" + required_fields = { + "name", + "config_path", + "data_path", + "example_index", + "example_sha256", + "agent_ref", + "expected_generation", + "expected_prompt_fragment", + "expected_result", + "expected_reward", + } + for case in cases: + assert required_fields <= case.keys(), ( + f"acceptance case is missing fields: {required_fields - case.keys()}" + ) + assert (_GYM_ROOT / case["config_path"]).is_file(), ( + f"{case['name']}: missing Gym config {case['config_path']}" + ) + data_path = _GYM_ROOT / case["data_path"] + assert data_path.is_file(), ( + f"{case['name']}: missing Gym data {case['data_path']}" + ) + with data_path.open("rb") as data_file: + examples = [line.rstrip(b"\r\n") for line in data_file if line.strip()] + assert 0 <= case["example_index"] < len(examples), ( + f"{case['name']}: example_index {case['example_index']} is outside a {len(examples)}-row dataset" + ) + actual_sha256 = sha256(examples[case["example_index"]]).hexdigest() + assert actual_sha256 == case["example_sha256"], ( + f"{case['name']}: pinned example changed; review the row and update its golden values" + ) + assert case["agent_ref"].keys() >= {"type", "name"} + assert case["expected_result"] + assert math.isfinite(case["expected_reward"]) + + names = [case["name"] for case in cases] + assert len(names) == len(set(names)), "acceptance case names must be unique" + assert set(names) == _REQUIRED_P0_CASES, ( + f"P0 matrix must contain exactly {_REQUIRED_P0_CASES}, got {set(names)}" + ) + return cases + + +_P0_CASES = _load_acceptance_cases() + + +def _load_case_datum(case: dict[str, Any]) -> DatumSpec: + dataset = NemoGymDataset(str(_GYM_ROOT / case["data_path"])) + datum = nemo_gym_data_processor( + dataset.dataset[case["example_index"]], None, None, None, 0 + ) + datum["extra_env_info"]["agent_ref"] = deepcopy(case["agent_ref"]) + return datum + + +@pytest.fixture(scope="module") +def scripted_openai_base_url(): + server = ThreadingHTTPServer(("0.0.0.0", 0), _ScriptedOpenAIHandler) + server_thread = Thread(target=server.serve_forever, daemon=True) + server_thread.start() + try: + yield f"http://{_get_node_ip_local()}:{server.server_port}/v1" + finally: + server.shutdown() + server.server_close() + server_thread.join(timeout=5) + + +@pytest.fixture +def p0_nemo_gym(scripted_openai_base_url, case): + """Start the selected real Gym environment with case-local attribution.""" + config_paths = [_POLICY_MODEL_CONFIG, case["config_path"]] + tokenizer = _ScriptedTokenizer() + env = spinup_nemo_gym_actor( + { + "nemo_gym": { + "config_paths": config_paths, + "skip_venv_if_present": True, + } + }, + base_urls=[scripted_openai_base_url], + model_name="scripted-model", + tokenizer=tokenizer, + enable_router_replay=False, + use_fastokens=False, + ) + try: + yield env + finally: + try: + ray.get(env.shutdown.remote(), timeout=10) + finally: + ray.kill(env) + + +@pytest.mark.nemo_gym +@pytest.mark.timeout(900) +@pytest.mark.parametrize("case", _P0_CASES, ids=[case["name"] for case in _P0_CASES]) +def test_p0_gym_environments_roll_out_through_nemo_rl(p0_nemo_gym, case): + """A pinned Gym example must preserve its contract across the NeMo RL boundary.""" + tokenizer = _ScriptedTokenizer() + result = run_nemo_gym_rollout_sync( + policy_generation=_ScriptedPolicyGeneration(), + input_batch=rl_collate_fn([_load_case_datum(case)]), + tokenizer=tokenizer, + task_to_env={"nemo_gym": p0_nemo_gym}, + max_seq_len=_GENERATION_CONFIG["max_total_sequence_length"], + generation_config=deepcopy(_GENERATION_CONFIG), + log_full_result_tables=True, + ) + + final_batch = result.final_batch + assert final_batch.size == 1 + assert final_batch["agent_ref"] == [case["agent_ref"]] + + reward = final_batch["total_reward"].item() + assert math.isfinite(reward) + assert reward == pytest.approx(case["expected_reward"]) + assert final_batch["length"].item() > 0 + + assistant_messages = [ + message + for message in final_batch["message_log"][0] + if message["role"] == "assistant" + ] + assert assistant_messages + for message in assistant_messages: + assert isinstance(message["token_ids"], torch.Tensor) + assert isinstance(message["generation_logprobs"], torch.Tensor) + assert len(message["token_ids"]) > 0 + assert len(message["token_ids"]) == len(message["generation_logprobs"]) + assert torch.isfinite(message["generation_logprobs"]).all() + + metric_prefix = f"{case['agent_ref']['name']}/reward/" + assert any(key.startswith(metric_prefix) for key in result.rollout_metrics) + + full_result_key = f"{case['agent_ref']['name']}/full_result" + assert full_result_key in result.rollout_metrics + full_result_table = result.rollout_metrics[full_result_key] + assert len(full_result_table.data) == 1 + full_result = json.loads(full_result_table.data[0][0]) + for field, expected_value in case["expected_result"].items(): + assert full_result[field] == expected_value + assert full_result["response"]["output"] + generation_strings = [ + item["generation_str"] + for item in full_result["response"]["output"] + if "generation_str" in item + ] + assert generation_strings == [case["expected_generation"]] + prompt_strings = [ + item["prompt_str"] + for item in full_result["response"]["output"] + if "prompt_str" in item + ] + assert len(prompt_strings) == 1 + assert case["expected_prompt_fragment"] in prompt_strings[0] From e4a9d932d0ed307cc173eadf2de8e233b13122c3 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 22:04:06 -0700 Subject: [PATCH 02/16] test: allowlist acceptance row hashes Signed-off-by: yaoyu-33 --- .../nemo_gym_test_data/p0_rollout_acceptance.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml index 3e6e1443929..0e57e86e747 100644 --- a/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml +++ b/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml @@ -3,7 +3,7 @@ cases: config_path: resources_servers/math_with_judge/configs/math_with_judge.yaml data_path: resources_servers/math_with_judge/data/example.jsonl example_index: 0 - example_sha256: 5e8d77674814a6b1fdcd0f2dbef0c8597cc109814763056070ff2eeb0c9d70ef + example_sha256: 5e8d77674814a6b1fdcd0f2dbef0c8597cc109814763056070ff2eeb0c9d70ef # pragma: allowlist secret agent_ref: type: responses_api_agents name: math_with_judge_simple_agent @@ -18,7 +18,7 @@ cases: config_path: resources_servers/code_gen/configs/code_gen.yaml data_path: resources_servers/code_gen/data/example.jsonl example_index: 4 - example_sha256: ceea8ed47c6615783bc9df8140110b2aa86729438448ee04e74c5d4850ebff24 + example_sha256: ceea8ed47c6615783bc9df8140110b2aa86729438448ee04e74c5d4850ebff24 # pragma: allowlist secret agent_ref: type: responses_api_agents name: code_gen_simple_agent @@ -55,7 +55,7 @@ cases: config_path: resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml data_path: resources_servers/single_step_tool_use_with_argument_comparison/data/example.jsonl example_index: 4 - example_sha256: 2e24731dfeda02e16351dd5abd5c1e3fd8307554f4824cf4fb3630672c63f5ca + example_sha256: 2e24731dfeda02e16351dd5abd5c1e3fd8307554f4824cf4fb3630672c63f5ca # pragma: allowlist secret agent_ref: type: responses_api_agents name: single_step_tool_use_with_argument_comparison_agent From d5251d9e3a4b484a75a6e475094b339a4a5b58ef Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 22:35:36 -0700 Subject: [PATCH 03/16] test: expand L0 Gym rollout acceptance Signed-off-by: yaoyu-33 --- .../run_nemo_gym_single_node_sanity_tests.sh | 4 +- .../l0_rollout_acceptance.yaml | 139 +++++++++++ .../p0_rollout_acceptance.yaml | 70 ------ ...=> test_l0_nemo_gym_rollout_acceptance.py} | 236 +++++++++++++++--- 4 files changed, 340 insertions(+), 109 deletions(-) create mode 100644 tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml delete mode 100644 tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml rename tests/unit/environments/{test_nemo_gym_rollout_acceptance.py => test_l0_nemo_gym_rollout_acceptance.py} (61%) diff --git a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh index b6dc544abb0..c3b1571e468 100755 --- a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh +++ b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh @@ -50,7 +50,7 @@ uv run python -c "import ray; ray.shutdown()" unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout \ --nemo-gym-only -# Exercise one pinned Gym example from each P0 environment through the real NeMo RL rollout boundary. +# Exercise one pinned example from each L0 Gym environment through the real NeMo RL rollout boundary. ./tests/run_unit.sh \ - unit/environments/test_nemo_gym_rollout_acceptance.py::test_p0_gym_environments_roll_out_through_nemo_rl \ + unit/environments/test_l0_nemo_gym_rollout_acceptance.py::test_l0_gym_environments_roll_out_through_nemo_rl \ --nemo-gym-only diff --git a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml new file mode 100644 index 00000000000..c9efd1f68a7 --- /dev/null +++ b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml @@ -0,0 +1,139 @@ +cases: + - name: math_with_judge + config_path: resources_servers/math_with_judge/configs/math_with_judge.yaml + data_path: resources_servers/math_with_judge/data/example.jsonl + example_index: 0 + example_sha256: 5e8d77674814a6b1fdcd0f2dbef0c8597cc109814763056070ff2eeb0c9d70ef # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: math_with_judge_simple_agent + expected_generations: + - '\boxed{32}' + expected_prompt_fragment: 1000 digit numbers + expected_reward: 1.0 + expected_result: + extracted_answer: "32" + library_reward: 1.0 + + - name: code_gen + config_path: resources_servers/code_gen/configs/code_gen.yaml + data_path: resources_servers/code_gen/data/example.jsonl + example_index: 4 + example_sha256: ceea8ed47c6615783bc9df8140110b2aa86729438448ee04e74c5d4850ebff24 # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: code_gen_simple_agent + expected_generations: + - |- + ```python + import sys + + values = iter(map(int, sys.stdin.buffer.read().split())) + test_count = next(values) + answers = [] + for _ in range(test_count): + n = next(values) + positions = [0] * (n + 1) + for index in range(n): + positions[next(values)] = index + left = right = positions[1] + bits = [] + for value in range(1, n + 1): + left = min(left, positions[value]) + right = max(right, positions[value]) + bits.append("1" if right - left + 1 == value else "0") + answers.append("".join(bits)) + sys.stdout.write("\n".join(answers)) + ``` + expected_prompt_fragment: beautiful number + expected_reward: 1.0 + expected_result: + result: + - true + - true + - true + + - name: single_step_tool_use_with_argument_comparison + config_path: resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + data_path: resources_servers/single_step_tool_use_with_argument_comparison/data/example.jsonl + example_index: 4 + example_sha256: 2e24731dfeda02e16351dd5abd5c1e3fd8307554f4824cf4fb3630672c63f5ca # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: single_step_tool_use_with_argument_comparison_agent + expected_generations: + - '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' + expected_prompt_fragment: SHOW24 + expected_reward: 1.0 + expected_result: + category: A tool call that matches the expected tool call was found + expected_action: + type: function_call + name: check_seat_availability + arguments: '{"event_id": "SHOW24", "section": "Medical Zone"}' + + - name: mcqa + config_path: resources_servers/mcqa/configs/mcqa.yaml + data_path: resources_servers/mcqa/data/example.jsonl + example_index: 0 + example_sha256: 3e3bd43ed3707cc35f84edf0a57357d0cc0f8dfd7692991d0daf05ad421d0846 # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: mcqa_simple_agent + expected_generations: + - '\boxed{B}' + expected_prompt_fragment: cystic fibrosis + expected_reward: 1.0 + expected_result: + expected_answer: B + extracted_answer: B + + - name: equivalence_llm_judge + config_path: resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + data_path: resources_servers/equivalence_llm_judge/data/example.jsonl + example_index: 0 + example_sha256: efabdc976b384f6f48ee5e42cf855a66e6fba84feedd1b15fd30f256f56ecd9b # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: equivalence_llm_judge_simple_agent + expected_generations: + - '\boxed{Charles Darwin}' + expected_prompt_fragment: theory of evolution by natural selection + expected_reward: 1.0 + expected_result: + expected_answer: darwin + expected_judge_verdicts: + - '[[A=B]]' + + - name: structured_outputs_v4 + config_path: resources_servers/structured_outputs/configs/structured_outputs_v4.yaml + data_path: resources_servers/structured_outputs/data/structured_outputs_v4_example.jsonl + example_index: 0 + example_sha256: 138b2f4383ec722989fdefd60ea0f1724712037e46e9d62ca2df0cc14308e1a4 # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: structured_outputs_v4_simple_agent + expected_generations: + - '{"name":"response_tool_8","arguments":{"extraction":{"name":"Dizer Kola","native_name":"ديزركلا","romanized_name":"Dizer Kola","settlement_type":"village","country":"Iran","province":"Mazandaran","county":"Nowshahr","bakhsh":"Central","rural_district":"Baladeh Kojur","coordinates":{"latitude":36.55694,"longitude":51.79389},"population_total":250,"population_year":2006,"number_of_families":64,"timezone_standard":"UTC+3:30 (IRST)","timezone_dst":"UTC+4:30 (IRDT)"}}}' + expected_prompt_fragment: Dizer Kola + expected_reward: 1.0 + expected_result: + schema_type: json + error_type: null + + - name: workplace_assistant + config_path: resources_servers/workplace_assistant/configs/workplace_assistant.yaml + data_path: resources_servers/workplace_assistant/data/example.jsonl + example_index: 0 + example_sha256: c7a80ae027d7bf89dea4700acc00a20d37400a91a19dbe9908d4b945367b657a # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: workplace_assistant_simple_agent + expected_generations: + - '{"name":"email_reply_email","arguments":{"email_id":"00000057","body":"Thanks for the update - I will get back to you tomorrow."}}' + - Done + expected_prompt_fragment: Task Update on Develop prototype + expected_reward: 1.0 + expected_result: {} + expected_tool_outputs: + - '{"output":"Email replied successfully."}' diff --git a/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml deleted file mode 100644 index 0e57e86e747..00000000000 --- a/tests/unit/environments/nemo_gym_test_data/p0_rollout_acceptance.yaml +++ /dev/null @@ -1,70 +0,0 @@ -cases: - - name: math_with_judge - config_path: resources_servers/math_with_judge/configs/math_with_judge.yaml - data_path: resources_servers/math_with_judge/data/example.jsonl - example_index: 0 - example_sha256: 5e8d77674814a6b1fdcd0f2dbef0c8597cc109814763056070ff2eeb0c9d70ef # pragma: allowlist secret - agent_ref: - type: responses_api_agents - name: math_with_judge_simple_agent - expected_generation: '\boxed{32}' - expected_prompt_fragment: 1000 digit numbers - expected_reward: 1.0 - expected_result: - extracted_answer: "32" - library_reward: 1.0 - - - name: code_gen - config_path: resources_servers/code_gen/configs/code_gen.yaml - data_path: resources_servers/code_gen/data/example.jsonl - example_index: 4 - example_sha256: ceea8ed47c6615783bc9df8140110b2aa86729438448ee04e74c5d4850ebff24 # pragma: allowlist secret - agent_ref: - type: responses_api_agents - name: code_gen_simple_agent - expected_generation: |- - ```python - import sys - - values = iter(map(int, sys.stdin.buffer.read().split())) - test_count = next(values) - answers = [] - for _ in range(test_count): - n = next(values) - positions = [0] * (n + 1) - for index in range(n): - positions[next(values)] = index - left = right = positions[1] - bits = [] - for value in range(1, n + 1): - left = min(left, positions[value]) - right = max(right, positions[value]) - bits.append("1" if right - left + 1 == value else "0") - answers.append("".join(bits)) - sys.stdout.write("\n".join(answers)) - ``` - expected_prompt_fragment: beautiful number - expected_reward: 1.0 - expected_result: - result: - - true - - true - - true - - - name: single_step_tool_use_with_argument_comparison - config_path: resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml - data_path: resources_servers/single_step_tool_use_with_argument_comparison/data/example.jsonl - example_index: 4 - example_sha256: 2e24731dfeda02e16351dd5abd5c1e3fd8307554f4824cf4fb3630672c63f5ca # pragma: allowlist secret - agent_ref: - type: responses_api_agents - name: single_step_tool_use_with_argument_comparison_agent - expected_generation: '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' - expected_prompt_fragment: SHOW24 - expected_reward: 1.0 - expected_result: - category: A tool call that matches the expected tool call was found - expected_action: - type: function_call - name: check_seat_availability - arguments: '{"event_id": "SHOW24", "section": "Medical Zone"}' diff --git a/tests/unit/environments/test_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py similarity index 61% rename from tests/unit/environments/test_nemo_gym_rollout_acceptance.py rename to tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 916ee653870..d091361ae09 100644 --- a/tests/unit/environments/test_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -37,14 +37,18 @@ _REPO_ROOT = Path(__file__).parents[3] _GYM_ROOT = _REPO_ROOT / "3rdparty/Gym-workspace/Gym" -_CASES_PATH = Path(__file__).parent / "nemo_gym_test_data/p0_rollout_acceptance.yaml" +_CASES_PATH = Path(__file__).parent / "nemo_gym_test_data/l0_rollout_acceptance.yaml" _POLICY_MODEL_CONFIG = ( "responses_api_models/vllm_model/configs/vllm_model_for_training.yaml" ) -_REQUIRED_P0_CASES = { +_REQUIRED_L0_CASES = { "code_gen", + "equivalence_llm_judge", "math_with_judge", + "mcqa", "single_step_tool_use_with_argument_comparison", + "structured_outputs_v4", + "workplace_assistant", } _GENERATION_CONFIG = { "backend": "test", @@ -81,6 +85,93 @@ '{"name":"check_seat_availability","arguments":' '{"event_id":"SHOW24","section":"Medical Zone"}}' ) +_MCQA_COMPLETION = r"\boxed{B}" +_EQUIVALENCE_COMPLETION = r"\boxed{Charles Darwin}" +_EQUIVALENCE_JUDGE_COMPLETION = ( + "The candidate identifies Darwin, so it matches the reference.\n\n" + "[[A=B]] they are equivalent" +) +_STRUCTURED_OUTPUTS_ARGUMENTS = { + "name": "Dizer Kola", + "native_name": "ديزركلا", + "romanized_name": "Dizer Kola", + "settlement_type": "village", + "country": "Iran", + "province": "Mazandaran", + "county": "Nowshahr", + "bakhsh": "Central", + "rural_district": "Baladeh Kojur", + "coordinates": {"latitude": 36.55694, "longitude": 51.79389}, + "population_total": 250, + "population_year": 2006, + "number_of_families": 64, + "timezone_standard": "UTC+3:30 (IRST)", + "timezone_dst": "UTC+4:30 (IRDT)", +} +_WORKPLACE_ARGUMENTS = { + "email_id": "00000057", + "body": "Thanks for the update - I will get back to you tomorrow.", +} + + +def _tool_call_message( + name: str, arguments: dict[str, Any], call_id: str +) -> tuple[dict[str, Any], str]: + message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, ensure_ascii=False), + }, + } + ], + } + generation = json.dumps( + {"name": name, "arguments": arguments}, + ensure_ascii=False, + separators=(",", ":"), + ) + return message, generation + + +def _prompt_token_ids(body: dict[str, Any]) -> list[int]: + serialized_request = json.dumps( + body, sort_keys=True, separators=(",", ":") + ).encode() + prompt_token_ids = list(b"\x00" + compress(serialized_request)) + + messages = body.get("messages", []) + if not any(message.get("role") == "tool" for message in messages): + return prompt_token_ids + + # The RL postprocessor requires each later prompt to start with every token + # already observed. Reconstruct the first workplace request, then append the + # first generation and an opaque encoding of the continuation request. + first_assistant_index = next( + index + for index, message in enumerate(messages) + if message.get("role") == "assistant" + ) + initial_body = deepcopy(body) + initial_body["messages"] = messages[:first_assistant_index] + initial_request = json.dumps( + initial_body, sort_keys=True, separators=(",", ":") + ).encode() + _, first_generation = _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + return ( + list(b"\x00" + compress(initial_request)) + + list(b"\x01" + first_generation.encode()) + + list(b"\x02" + compress(serialized_request)) + ) class _ScriptedPolicyGeneration: @@ -115,6 +206,11 @@ class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): @staticmethod def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: serialized_body = json.dumps(body) + if "GOLD:" in serialized_body and "CANDIDATE:" in serialized_body: + return { + "role": "assistant", + "content": _EQUIVALENCE_JUDGE_COMPLETION, + }, _EQUIVALENCE_JUDGE_COMPLETION if "1000 digit numbers" in serialized_body: content = r"\boxed{32}" return {"role": "assistant", "content": content}, content @@ -124,24 +220,40 @@ def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: "content": _CODE_GEN_COMPLETION, }, _CODE_GEN_COMPLETION if "SHOW24" in serialized_body and "Medical Zone" in serialized_body: - return ( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call-p0-acceptance", - "type": "function", - "function": { - "name": "check_seat_availability", - "arguments": _TOOL_CALL_ARGUMENTS, - }, - } - ], - }, - _TOOL_CALL_GENERATION, + message, _ = _tool_call_message( + "check_seat_availability", + json.loads(_TOOL_CALL_ARGUMENTS), + "call-l0-tool-use", ) - raise AssertionError("P0 scripted policy received an unrecognized prompt") + return message, _TOOL_CALL_GENERATION + if "cystic fibrosis" in serialized_body: + return { + "role": "assistant", + "content": _MCQA_COMPLETION, + }, _MCQA_COMPLETION + if "theory of evolution by natural selection" in serialized_body: + return { + "role": "assistant", + "content": _EQUIVALENCE_COMPLETION, + }, _EQUIVALENCE_COMPLETION + if "Dizer Kola" in serialized_body: + return _tool_call_message( + "response_tool_8", + {"extraction": _STRUCTURED_OUTPUTS_ARGUMENTS}, + "call-l0-structured-output", + ) + if "Task Update on Develop prototype" in serialized_body: + if any( + message.get("role") == "tool" for message in body.get("messages", []) + ): + content = "Done" + return {"role": "assistant", "content": content}, content + return _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + raise AssertionError("L0 scripted policy received an unrecognized prompt") def do_POST(self) -> None: # noqa: N802 if self.path != "/v1/chat/completions": @@ -151,10 +263,7 @@ def do_POST(self) -> None: # noqa: N802 content_length = int(self.headers.get("Content-Length", "0")) body = json.loads(self.rfile.read(content_length)) message, generation_text = self._message_for(body) - serialized_request = json.dumps( - body, sort_keys=True, separators=(",", ":") - ).encode() - prompt_token_ids = list(b"\x00" + compress(serialized_request)) + prompt_token_ids = _prompt_token_ids(body) generation_token_ids = list(b"\x01" + generation_text.encode()) message.update( { @@ -164,7 +273,7 @@ def do_POST(self) -> None: # noqa: N802 } ) response = { - "id": "chatcmpl-p0-acceptance", + "id": "chatcmpl-l0-acceptance", "object": "chat.completion", "created": 0, "model": body.get("model", "scripted-model"), @@ -198,7 +307,7 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: with _CASES_PATH.open() as case_file: cases = safe_load(case_file)["cases"] - assert cases, "P0 Gym rollout acceptance matrix must not be empty" + assert cases, "L0 Gym rollout acceptance matrix must not be empty" required_fields = { "name", "config_path", @@ -206,7 +315,7 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: "example_index", "example_sha256", "agent_ref", - "expected_generation", + "expected_generations", "expected_prompt_fragment", "expected_result", "expected_reward", @@ -232,18 +341,58 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: f"{case['name']}: pinned example changed; review the row and update its golden values" ) assert case["agent_ref"].keys() >= {"type", "name"} - assert case["expected_result"] + assert case["expected_generations"] + assert all( + isinstance(generation, str) for generation in case["expected_generations"] + ) assert math.isfinite(case["expected_reward"]) names = [case["name"] for case in cases] assert len(names) == len(set(names)), "acceptance case names must be unique" - assert set(names) == _REQUIRED_P0_CASES, ( - f"P0 matrix must contain exactly {_REQUIRED_P0_CASES}, got {set(names)}" + assert set(names) == _REQUIRED_L0_CASES, ( + f"L0 matrix must contain exactly {_REQUIRED_L0_CASES}, got {set(names)}" ) return cases -_P0_CASES = _load_acceptance_cases() +_L0_CASES = _load_acceptance_cases() + + +@pytest.mark.nemo_gym +def test_l0_scripted_multiturn_tokens_are_contiguous(): + initial_body = { + "model": "scripted-model", + "messages": [ + {"role": "system", "content": "system"}, + { + "role": "user", + "content": "Task Update on Develop prototype", + }, + ], + } + continuation_body = deepcopy(initial_body) + continuation_body["messages"].extend( + [ + {"role": "assistant", "content": None, "tool_calls": []}, + { + "role": "tool", + "tool_call_id": "call-l0-workplace", + "content": '{"output":"Email replied successfully."}', + }, + ] + ) + _, first_generation = _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + seen_token_ids = _prompt_token_ids(initial_body) + list( + b"\x01" + first_generation.encode() + ) + + continuation_prompt = _prompt_token_ids(continuation_body) + + assert continuation_prompt[: len(seen_token_ids)] == seen_token_ids def _load_case_datum(case: dict[str, Any]) -> DatumSpec: @@ -269,7 +418,7 @@ def scripted_openai_base_url(): @pytest.fixture -def p0_nemo_gym(scripted_openai_base_url, case): +def l0_nemo_gym(scripted_openai_base_url, case): """Start the selected real Gym environment with case-local attribution.""" config_paths = [_POLICY_MODEL_CONFIG, case["config_path"]] tokenizer = _ScriptedTokenizer() @@ -297,15 +446,15 @@ def p0_nemo_gym(scripted_openai_base_url, case): @pytest.mark.nemo_gym @pytest.mark.timeout(900) -@pytest.mark.parametrize("case", _P0_CASES, ids=[case["name"] for case in _P0_CASES]) -def test_p0_gym_environments_roll_out_through_nemo_rl(p0_nemo_gym, case): +@pytest.mark.parametrize("case", _L0_CASES, ids=[case["name"] for case in _L0_CASES]) +def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): """A pinned Gym example must preserve its contract across the NeMo RL boundary.""" tokenizer = _ScriptedTokenizer() result = run_nemo_gym_rollout_sync( policy_generation=_ScriptedPolicyGeneration(), input_batch=rl_collate_fn([_load_case_datum(case)]), tokenizer=tokenizer, - task_to_env={"nemo_gym": p0_nemo_gym}, + task_to_env={"nemo_gym": l0_nemo_gym}, max_seq_len=_GENERATION_CONFIG["max_total_sequence_length"], generation_config=deepcopy(_GENERATION_CONFIG), log_full_result_tables=True, @@ -344,16 +493,29 @@ def test_p0_gym_environments_roll_out_through_nemo_rl(p0_nemo_gym, case): for field, expected_value in case["expected_result"].items(): assert full_result[field] == expected_value assert full_result["response"]["output"] + if "expected_tool_outputs" in case: + tool_outputs = [ + item["output"] + for item in full_result["response"]["output"] + if item["type"] == "function_call_output" + ] + assert tool_outputs == case["expected_tool_outputs"] + if "expected_judge_verdicts" in case: + verdicts = [ + evaluation["verdict_label"] + for evaluation in full_result["judge_evaluations"] + ] + assert verdicts == case["expected_judge_verdicts"] generation_strings = [ item["generation_str"] for item in full_result["response"]["output"] if "generation_str" in item ] - assert generation_strings == [case["expected_generation"]] + assert generation_strings == case["expected_generations"] prompt_strings = [ item["prompt_str"] for item in full_result["response"]["output"] if "prompt_str" in item ] - assert len(prompt_strings) == 1 - assert case["expected_prompt_fragment"] in prompt_strings[0] + assert len(prompt_strings) == len(case["expected_generations"]) + assert all(case["expected_prompt_fragment"] in prompt for prompt in prompt_strings) From 2beac64c96cb20a413f67406df8084edfadb400b Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 23:29:35 -0700 Subject: [PATCH 04/16] test: tighten Gym rollout transport assertions Signed-off-by: yaoyu-33 --- .../test_l0_nemo_gym_rollout_acceptance.py | 52 +++++++++++++++---- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index d091361ae09..3dee25bb2c8 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -112,6 +112,7 @@ "email_id": "00000057", "body": "Thanks for the update - I will get back to you tomorrow.", } +_CONTINUATION_FRAME = b"\x02nemo-rl-continuation\x00" def _tool_call_message( @@ -170,7 +171,7 @@ def _prompt_token_ids(body: dict[str, Any]) -> list[int]: return ( list(b"\x00" + compress(initial_request)) + list(b"\x01" + first_generation.encode()) - + list(b"\x02" + compress(serialized_request)) + + list(_CONTINUATION_FRAME + serialized_request) ) @@ -190,7 +191,10 @@ def batch_decode(self, batch: list[list[int]]) -> list[str]: for token_ids in batch: payload = bytes(int(token_id) for token_id in token_ids) if payload.startswith(b"\x00"): - decoded.append(decompress(payload[1:]).decode()) + _, frame, continuation = payload.rpartition(_CONTINUATION_FRAME) + decoded.append( + continuation.decode() if frame else decompress(payload[1:]).decode() + ) elif payload.startswith(b"\x01"): decoded.append(payload[1:].decode()) else: @@ -203,6 +207,16 @@ class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" + @staticmethod + def _assert_request_contract(body: dict[str, Any]) -> None: + assert body["model"] == "scripted-model" + assert body["temperature"] == _GENERATION_CONFIG["temperature"] + assert body["top_p"] == _GENERATION_CONFIG["top_p"] + assert body["max_tokens"] == _GENERATION_CONFIG["max_new_tokens"] + assert body["logprobs"] is True + assert body["top_logprobs"] == 0 + assert body["return_tokens_as_token_ids"] is True + @staticmethod def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: serialized_body = json.dumps(body) @@ -262,6 +276,7 @@ def do_POST(self) -> None: # noqa: N802 content_length = int(self.headers.get("Content-Length", "0")) body = json.loads(self.rfile.read(content_length)) + self._assert_request_contract(body) message, generation_text = self._message_for(body) prompt_token_ids = _prompt_token_ids(body) generation_token_ids = list(b"\x01" + generation_text.encode()) @@ -393,6 +408,9 @@ def test_l0_scripted_multiturn_tokens_are_contiguous(): continuation_prompt = _prompt_token_ids(continuation_body) assert continuation_prompt[: len(seen_token_ids)] == seen_token_ids + assert json.loads(_ScriptedTokenizer().batch_decode([continuation_prompt])[0]) == ( + continuation_body + ) def _load_case_datum(case: dict[str, Any]) -> DatumSpec: @@ -417,10 +435,10 @@ def scripted_openai_base_url(): server_thread.join(timeout=5) -@pytest.fixture -def l0_nemo_gym(scripted_openai_base_url, case): - """Start the selected real Gym environment with case-local attribution.""" - config_paths = [_POLICY_MODEL_CONFIG, case["config_path"]] +@pytest.fixture(scope="module") +def l0_nemo_gym(scripted_openai_base_url): + """Start all acceptance environments in one real Gym actor.""" + config_paths = [_POLICY_MODEL_CONFIG] + [case["config_path"] for case in _L0_CASES] tokenizer = _ScriptedTokenizer() env = spinup_nemo_gym_actor( { @@ -475,12 +493,18 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): if message["role"] == "assistant" ] assert assistant_messages - for message in assistant_messages: + assert len(assistant_messages) == len(case["expected_generations"]) + for message, generation in zip( + assistant_messages, case["expected_generations"], strict=True + ): assert isinstance(message["token_ids"], torch.Tensor) assert isinstance(message["generation_logprobs"], torch.Tensor) - assert len(message["token_ids"]) > 0 - assert len(message["token_ids"]) == len(message["generation_logprobs"]) - assert torch.isfinite(message["generation_logprobs"]).all() + expected_token_ids = torch.tensor(list(b"\x01" + generation.encode())) + torch.testing.assert_close(message["token_ids"], expected_token_ids) + torch.testing.assert_close( + message["generation_logprobs"], + torch.full((len(expected_token_ids),), -0.1), + ) metric_prefix = f"{case['agent_ref']['name']}/reward/" assert any(key.startswith(metric_prefix) for key in result.rollout_metrics) @@ -519,3 +543,11 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): ] assert len(prompt_strings) == len(case["expected_generations"]) assert all(case["expected_prompt_fragment"] in prompt for prompt in prompt_strings) + if case["name"] == "workplace_assistant": + continuation_messages = json.loads(prompt_strings[1])["messages"] + assert any(message.get("tool_calls") for message in continuation_messages) + assert any( + message.get("role") == "tool" + and "Email replied successfully" in message.get("content", "") + for message in continuation_messages + ) From 85e137efddf7d59ac2c48aa8cbbe853519487eac Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Fri, 4 Sep 2026 23:44:54 -0700 Subject: [PATCH 05/16] test: distinguish rollout and judge contracts Signed-off-by: yaoyu-33 --- .../test_l0_nemo_gym_rollout_acceptance.py | 33 +++++++++++++++++-- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 3dee25bb2c8..c3d3d4168b0 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -209,14 +209,41 @@ class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): @staticmethod def _assert_request_contract(body: dict[str, Any]) -> None: + serialized_body = json.dumps(body) + is_judge_request = ( + "GOLD:" in serialized_body and "CANDIDATE:" in serialized_body + ) + assert body["model"] == "scripted-model" - assert body["temperature"] == _GENERATION_CONFIG["temperature"] - assert body["top_p"] == _GENERATION_CONFIG["top_p"] - assert body["max_tokens"] == _GENERATION_CONFIG["max_new_tokens"] assert body["logprobs"] is True assert body["top_logprobs"] == 0 assert body["return_tokens_as_token_ids"] is True + if is_judge_request: + assert "temperature" not in body + assert "top_p" not in body + assert "max_tokens" not in body + else: + assert body["temperature"] == _GENERATION_CONFIG["temperature"] + assert body["top_p"] == _GENERATION_CONFIG["top_p"] + assert body["max_tokens"] == _GENERATION_CONFIG["max_new_tokens"] + + expected_tool_name = None + if "SHOW24" in serialized_body and "Medical Zone" in serialized_body: + expected_tool_name = "check_seat_availability" + elif "Dizer Kola" in serialized_body: + expected_tool_name = "response_tool_8" + elif "Task Update on Develop prototype" in serialized_body: + expected_tool_name = "email_reply_email" + + if expected_tool_name is not None: + tool_names = { + tool["function"]["name"] + for tool in body.get("tools", []) + if tool.get("type") == "function" + } + assert expected_tool_name in tool_names + @staticmethod def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: serialized_body = json.dumps(body) From bce2879cfd0418b42594ec2ca20d5a528a92ef46 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 11:29:05 -0700 Subject: [PATCH 06/16] test: use Ray-serializable tokenizer for Gym acceptance Signed-off-by: yaoyu-33 --- .../test_l0_nemo_gym_rollout_acceptance.py | 77 ++++++++----------- 1 file changed, 33 insertions(+), 44 deletions(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index c3d3d4168b0..2550c064a70 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -20,11 +20,11 @@ from pathlib import Path from threading import Thread from typing import Any -from zlib import compress, decompress import pytest import ray import torch +from transformers import ByT5Tokenizer from yaml import safe_load from nemo_rl.data.collate_fn import rl_collate_fn @@ -53,7 +53,7 @@ _GENERATION_CONFIG = { "backend": "test", "max_new_tokens": 1024, - "max_total_sequence_length": 16384, + "max_total_sequence_length": 65536, "temperature": 0.0, "top_p": 1.0, "top_k": None, @@ -112,7 +112,8 @@ "email_id": "00000057", "body": "Thanks for the update - I will get back to you tomorrow.", } -_CONTINUATION_FRAME = b"\x02nemo-rl-continuation\x00" +_BYTE_TOKEN_OFFSET = 3 +_CONTINUATION_MARKER = "\x02nemo-rl-continuation\x00" def _tool_call_message( @@ -140,11 +141,14 @@ def _tool_call_message( return message, generation +def _text_token_ids(text: str) -> list[int]: + """Match ByT5's byte vocabulary without deriving goldens from the tokenizer.""" + return [byte + _BYTE_TOKEN_OFFSET for byte in text.encode()] + + def _prompt_token_ids(body: dict[str, Any]) -> list[int]: - serialized_request = json.dumps( - body, sort_keys=True, separators=(",", ":") - ).encode() - prompt_token_ids = list(b"\x00" + compress(serialized_request)) + serialized_request = json.dumps(body, sort_keys=True, separators=(",", ":")) + prompt_token_ids = _text_token_ids(serialized_request) messages = body.get("messages", []) if not any(message.get("role") == "tool" for message in messages): @@ -152,7 +156,7 @@ def _prompt_token_ids(body: dict[str, Any]) -> list[int]: # The RL postprocessor requires each later prompt to start with every token # already observed. Reconstruct the first workplace request, then append the - # first generation and an opaque encoding of the continuation request. + # first generation and a marked encoding of the continuation request. first_assistant_index = next( index for index, message in enumerate(messages) @@ -160,18 +164,16 @@ def _prompt_token_ids(body: dict[str, Any]) -> list[int]: ) initial_body = deepcopy(body) initial_body["messages"] = messages[:first_assistant_index] - initial_request = json.dumps( - initial_body, sort_keys=True, separators=(",", ":") - ).encode() + initial_request = json.dumps(initial_body, sort_keys=True, separators=(",", ":")) _, first_generation = _tool_call_message( "email_reply_email", _WORKPLACE_ARGUMENTS, "call-l0-workplace", ) return ( - list(b"\x00" + compress(initial_request)) - + list(b"\x01" + first_generation.encode()) - + list(_CONTINUATION_FRAME + serialized_request) + _text_token_ids(initial_request) + + _text_token_ids(first_generation) + + _text_token_ids(_CONTINUATION_MARKER + serialized_request) ) @@ -181,27 +183,6 @@ class _ScriptedPolicyGeneration: cfg = _GENERATION_CONFIG -class _ScriptedTokenizer: - """Reversibly decode request and generation metadata without model weights.""" - - pad_token_id = 256 - - def batch_decode(self, batch: list[list[int]]) -> list[str]: - decoded = [] - for token_ids in batch: - payload = bytes(int(token_id) for token_id in token_ids) - if payload.startswith(b"\x00"): - _, frame, continuation = payload.rpartition(_CONTINUATION_FRAME) - decoded.append( - continuation.decode() if frame else decompress(payload[1:]).decode() - ) - elif payload.startswith(b"\x01"): - decoded.append(payload[1:].decode()) - else: - raise ValueError("Unrecognized scripted token sequence") - return decoded - - class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): """Return deterministic policy outputs to real Gym model and agent services.""" @@ -306,7 +287,11 @@ def do_POST(self) -> None: # noqa: N802 self._assert_request_contract(body) message, generation_text = self._message_for(body) prompt_token_ids = _prompt_token_ids(body) - generation_token_ids = list(b"\x01" + generation_text.encode()) + generation_token_ids = _text_token_ids(generation_text) + assert ( + len(prompt_token_ids) + len(generation_token_ids) + <= _GENERATION_CONFIG["max_total_sequence_length"] + ) message.update( { "prompt_token_ids": prompt_token_ids, @@ -428,14 +413,14 @@ def test_l0_scripted_multiturn_tokens_are_contiguous(): _WORKPLACE_ARGUMENTS, "call-l0-workplace", ) - seen_token_ids = _prompt_token_ids(initial_body) + list( - b"\x01" + first_generation.encode() - ) + seen_token_ids = _prompt_token_ids(initial_body) + _text_token_ids(first_generation) continuation_prompt = _prompt_token_ids(continuation_body) assert continuation_prompt[: len(seen_token_ids)] == seen_token_ids - assert json.loads(_ScriptedTokenizer().batch_decode([continuation_prompt])[0]) == ( + decoded_continuation = ByT5Tokenizer().batch_decode([continuation_prompt])[0] + assert _CONTINUATION_MARKER in decoded_continuation + assert json.loads(decoded_continuation.rpartition(_CONTINUATION_MARKER)[2]) == ( continuation_body ) @@ -466,7 +451,7 @@ def scripted_openai_base_url(): def l0_nemo_gym(scripted_openai_base_url): """Start all acceptance environments in one real Gym actor.""" config_paths = [_POLICY_MODEL_CONFIG] + [case["config_path"] for case in _L0_CASES] - tokenizer = _ScriptedTokenizer() + tokenizer = ByT5Tokenizer() env = spinup_nemo_gym_actor( { "nemo_gym": { @@ -494,7 +479,7 @@ def l0_nemo_gym(scripted_openai_base_url): @pytest.mark.parametrize("case", _L0_CASES, ids=[case["name"] for case in _L0_CASES]) def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): """A pinned Gym example must preserve its contract across the NeMo RL boundary.""" - tokenizer = _ScriptedTokenizer() + tokenizer = ByT5Tokenizer() result = run_nemo_gym_rollout_sync( policy_generation=_ScriptedPolicyGeneration(), input_batch=rl_collate_fn([_load_case_datum(case)]), @@ -526,7 +511,9 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): ): assert isinstance(message["token_ids"], torch.Tensor) assert isinstance(message["generation_logprobs"], torch.Tensor) - expected_token_ids = torch.tensor(list(b"\x01" + generation.encode())) + expected_token_ids = torch.tensor( + [byte + _BYTE_TOKEN_OFFSET for byte in generation.encode()] + ) torch.testing.assert_close(message["token_ids"], expected_token_ids) torch.testing.assert_close( message["generation_logprobs"], @@ -571,7 +558,9 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): assert len(prompt_strings) == len(case["expected_generations"]) assert all(case["expected_prompt_fragment"] in prompt for prompt in prompt_strings) if case["name"] == "workplace_assistant": - continuation_messages = json.loads(prompt_strings[1])["messages"] + assert _CONTINUATION_MARKER in prompt_strings[1] + continuation_json = prompt_strings[1].rpartition(_CONTINUATION_MARKER)[2] + continuation_messages = json.loads(continuation_json)["messages"] assert any(message.get("tool_calls") for message in continuation_messages) assert any( message.get("role") == "tool" From 5533d83efc42e19a538634dbd6525ba7082cd901 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 13:21:55 -0700 Subject: [PATCH 07/16] test: allow long Gym tool continuation prompts Signed-off-by: yaoyu-33 --- .../unit/environments/test_l0_nemo_gym_rollout_acceptance.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 2550c064a70..dfefcf0f840 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -53,7 +53,9 @@ _GENERATION_CONFIG = { "backend": "test", "max_new_tokens": 1024, - "max_total_sequence_length": 65536, + # Tool-heavy environments such as workplace_assistant serialize more than + # 64K byte-level tokens before a continuation generation. + "max_total_sequence_length": 131072, "temperature": 0.0, "top_p": 1.0, "top_k": None, From 436845451bd6691166e1e05fedd560ecb46d2357 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 13:56:36 -0700 Subject: [PATCH 08/16] test: allow Gym subprocesses to shut down Signed-off-by: yaoyu-33 --- .../unit/environments/test_l0_nemo_gym_rollout_acceptance.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index dfefcf0f840..ae75b0c395d 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -471,7 +471,9 @@ def l0_nemo_gym(scripted_openai_base_url): yield env finally: try: - ray.get(env.shutdown.remote(), timeout=10) + # The shared actor owns many Gym subprocesses, which are reaped + # sequentially during graceful shutdown. + ray.get(env.shutdown.remote(), timeout=60) finally: ray.kill(env) From 8954ff2307b4d531d63d0d4f924cae243d906217 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 14:47:21 -0700 Subject: [PATCH 09/16] test: cover rejected Gym rollout paths Signed-off-by: yaoyu-33 --- .../l0_rollout_acceptance.yaml | 19 +++++ .../test_l0_nemo_gym_rollout_acceptance.py | 71 +++++++++++++++---- 2 files changed, 75 insertions(+), 15 deletions(-) diff --git a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml index c9efd1f68a7..a852a6f3ff8 100644 --- a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml +++ b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml @@ -10,6 +10,8 @@ cases: expected_generations: - '\boxed{32}' expected_prompt_fragment: 1000 digit numbers + rejected_generations: + - '\boxed{999999}' expected_reward: 1.0 expected_result: extracted_answer: "32" @@ -46,6 +48,11 @@ cases: sys.stdout.write("\n".join(answers)) ``` expected_prompt_fragment: beautiful number + rejected_generations: + - |- + ```python + print(0) + ``` expected_reward: 1.0 expected_result: result: @@ -64,6 +71,8 @@ cases: expected_generations: - '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' expected_prompt_fragment: SHOW24 + rejected_generations: + - I cannot call the requested tool. expected_reward: 1.0 expected_result: category: A tool call that matches the expected tool call was found @@ -83,6 +92,8 @@ cases: expected_generations: - '\boxed{B}' expected_prompt_fragment: cystic fibrosis + rejected_generations: + - '\boxed{A}' expected_reward: 1.0 expected_result: expected_answer: B @@ -99,6 +110,10 @@ cases: expected_generations: - '\boxed{Charles Darwin}' expected_prompt_fragment: theory of evolution by natural selection + rejected_generations: + - '\boxed{Ada Lovelace}' + rejected_judge_verdicts: + - "[[A!=B]]" expected_reward: 1.0 expected_result: expected_answer: darwin @@ -116,6 +131,8 @@ cases: expected_generations: - '{"name":"response_tool_8","arguments":{"extraction":{"name":"Dizer Kola","native_name":"ديزركلا","romanized_name":"Dizer Kola","settlement_type":"village","country":"Iran","province":"Mazandaran","county":"Nowshahr","bakhsh":"Central","rural_district":"Baladeh Kojur","coordinates":{"latitude":36.55694,"longitude":51.79389},"population_total":250,"population_year":2006,"number_of_families":64,"timezone_standard":"UTC+3:30 (IRST)","timezone_dst":"UTC+4:30 (IRDT)"}}}' expected_prompt_fragment: Dizer Kola + rejected_generations: + - I cannot extract that. expected_reward: 1.0 expected_result: schema_type: json @@ -133,6 +150,8 @@ cases: - '{"name":"email_reply_email","arguments":{"email_id":"00000057","body":"Thanks for the update - I will get back to you tomorrow."}}' - Done expected_prompt_fragment: Task Update on Develop prototype + rejected_generations: + - Done expected_reward: 1.0 expected_result: {} expected_tool_outputs: diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index ae75b0c395d..4c44519add4 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -116,6 +116,7 @@ } _BYTE_TOKEN_OFFSET = 3 _CONTINUATION_MARKER = "\x02nemo-rl-continuation\x00" +_REJECTED_ROLLOUT_MARKER = "[NEMO_RL_L0_REJECTED_ROLLOUT]" def _tool_call_message( @@ -231,10 +232,31 @@ def _assert_request_contract(body: dict[str, Any]) -> None: def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: serialized_body = json.dumps(body) if "GOLD:" in serialized_body and "CANDIDATE:" in serialized_body: + if r"\\boxed{Ada Lovelace}" in serialized_body: + content = "The candidate does not match the reference.\n\n[[A!=B]]" + return {"role": "assistant", "content": content}, content return { "role": "assistant", "content": _EQUIVALENCE_JUDGE_COMPLETION, }, _EQUIVALENCE_JUDGE_COMPLETION + if _REJECTED_ROLLOUT_MARKER in serialized_body: + if "1000 digit numbers" in serialized_body: + content = r"\boxed{999999}" + elif "python programming language only" in serialized_body: + content = "```python\nprint(0)\n```" + elif "SHOW24" in serialized_body and "Medical Zone" in serialized_body: + content = "I cannot call the requested tool." + elif "cystic fibrosis" in serialized_body: + content = r"\boxed{A}" + elif "theory of evolution by natural selection" in serialized_body: + content = r"\boxed{Ada Lovelace}" + elif "Dizer Kola" in serialized_body: + content = "I cannot extract that." + elif "Task Update on Develop prototype" in serialized_body: + content = "Done" + else: + raise AssertionError("L0 rejected policy received an unrecognized prompt") + return {"role": "assistant", "content": content}, content if "1000 digit numbers" in serialized_body: content = r"\boxed{32}" return {"role": "assistant", "content": content}, content @@ -345,6 +367,7 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: "example_sha256", "agent_ref", "expected_generations", + "rejected_generations", "expected_prompt_fragment", "expected_result", "expected_reward", @@ -371,8 +394,11 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: ) assert case["agent_ref"].keys() >= {"type", "name"} assert case["expected_generations"] + assert case["rejected_generations"] assert all( - isinstance(generation, str) for generation in case["expected_generations"] + isinstance(generation, str) + for key in ("expected_generations", "rejected_generations") + for generation in case[key] ) assert math.isfinite(case["expected_reward"]) @@ -427,12 +453,18 @@ def test_l0_scripted_multiturn_tokens_are_contiguous(): ) -def _load_case_datum(case: dict[str, Any]) -> DatumSpec: +def _load_case_datum(case: dict[str, Any], *, accepted: bool) -> DatumSpec: dataset = NemoGymDataset(str(_GYM_ROOT / case["data_path"])) datum = nemo_gym_data_processor( dataset.dataset[case["example_index"]], None, None, None, 0 ) datum["extra_env_info"]["agent_ref"] = deepcopy(case["agent_ref"]) + if not accepted: + input_items = datum["extra_env_info"]["responses_create_params"]["input"] + final_user_item = next( + item for item in reversed(input_items) if item.get("role") == "user" + ) + final_user_item["content"] += f"\n\n{_REJECTED_ROLLOUT_MARKER}" return datum @@ -481,12 +513,17 @@ def l0_nemo_gym(scripted_openai_base_url): @pytest.mark.nemo_gym @pytest.mark.timeout(900) @pytest.mark.parametrize("case", _L0_CASES, ids=[case["name"] for case in _L0_CASES]) -def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): +@pytest.mark.parametrize("accepted", [True, False], ids=["accepted", "rejected"]) +def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepted): """A pinned Gym example must preserve its contract across the NeMo RL boundary.""" tokenizer = ByT5Tokenizer() + expected_generations = case[ + "expected_generations" if accepted else "rejected_generations" + ] + expected_reward = case["expected_reward"] if accepted else 0.0 result = run_nemo_gym_rollout_sync( policy_generation=_ScriptedPolicyGeneration(), - input_batch=rl_collate_fn([_load_case_datum(case)]), + input_batch=rl_collate_fn([_load_case_datum(case, accepted=accepted)]), tokenizer=tokenizer, task_to_env={"nemo_gym": l0_nemo_gym}, max_seq_len=_GENERATION_CONFIG["max_total_sequence_length"], @@ -500,7 +537,7 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): reward = final_batch["total_reward"].item() assert math.isfinite(reward) - assert reward == pytest.approx(case["expected_reward"]) + assert reward == pytest.approx(expected_reward) assert final_batch["length"].item() > 0 assistant_messages = [ @@ -509,9 +546,9 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): if message["role"] == "assistant" ] assert assistant_messages - assert len(assistant_messages) == len(case["expected_generations"]) + assert len(assistant_messages) == len(expected_generations) for message, generation in zip( - assistant_messages, case["expected_generations"], strict=True + assistant_messages, expected_generations, strict=True ): assert isinstance(message["token_ids"], torch.Tensor) assert isinstance(message["generation_logprobs"], torch.Tensor) @@ -532,36 +569,40 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case): full_result_table = result.rollout_metrics[full_result_key] assert len(full_result_table.data) == 1 full_result = json.loads(full_result_table.data[0][0]) - for field, expected_value in case["expected_result"].items(): - assert full_result[field] == expected_value + if accepted: + for field, expected_value in case["expected_result"].items(): + assert full_result[field] == expected_value assert full_result["response"]["output"] - if "expected_tool_outputs" in case: + if accepted and "expected_tool_outputs" in case: tool_outputs = [ item["output"] for item in full_result["response"]["output"] if item["type"] == "function_call_output" ] assert tool_outputs == case["expected_tool_outputs"] - if "expected_judge_verdicts" in case: + expected_judge_verdicts = case.get( + "expected_judge_verdicts" if accepted else "rejected_judge_verdicts" + ) + if expected_judge_verdicts is not None: verdicts = [ evaluation["verdict_label"] for evaluation in full_result["judge_evaluations"] ] - assert verdicts == case["expected_judge_verdicts"] + assert verdicts == expected_judge_verdicts generation_strings = [ item["generation_str"] for item in full_result["response"]["output"] if "generation_str" in item ] - assert generation_strings == case["expected_generations"] + assert generation_strings == expected_generations prompt_strings = [ item["prompt_str"] for item in full_result["response"]["output"] if "prompt_str" in item ] - assert len(prompt_strings) == len(case["expected_generations"]) + assert len(prompt_strings) == len(expected_generations) assert all(case["expected_prompt_fragment"] in prompt for prompt in prompt_strings) - if case["name"] == "workplace_assistant": + if accepted and case["name"] == "workplace_assistant": assert _CONTINUATION_MARKER in prompt_strings[1] continuation_json = prompt_strings[1].rpartition(_CONTINUATION_MARKER)[2] continuation_messages = json.loads(continuation_json)["messages"] From b77e034404480cd24719db6f61e6e2dde9a297ec Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 15:12:06 -0700 Subject: [PATCH 10/16] style: format L2 rollout acceptance test Signed-off-by: yaoyu-33 --- .../unit/environments/test_l0_nemo_gym_rollout_acceptance.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 4c44519add4..50cd9083287 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -255,7 +255,9 @@ def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: elif "Task Update on Develop prototype" in serialized_body: content = "Done" else: - raise AssertionError("L0 rejected policy received an unrecognized prompt") + raise AssertionError( + "L0 rejected policy received an unrecognized prompt" + ) return {"role": "assistant", "content": content}, content if "1000 digit numbers" in serialized_body: content = r"\boxed{32}" From 220faac8af8a973e5e761de12e1121f36b473823 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 15:29:37 -0700 Subject: [PATCH 11/16] test: preserve L2 verifier metadata contracts Signed-off-by: yaoyu-33 --- .../l0_rollout_acceptance.yaml | 28 +++++++++++++++++++ .../test_l0_nemo_gym_rollout_acceptance.py | 19 ++++++++++++- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml index a852a6f3ff8..65f557c3cfc 100644 --- a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml +++ b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml @@ -10,6 +10,9 @@ cases: expected_generations: - '\boxed{32}' expected_prompt_fragment: 1000 digit numbers + metadata_fields: + - question + - expected_answer rejected_generations: - '\boxed{999999}' expected_reward: 1.0 @@ -48,6 +51,8 @@ cases: sys.stdout.write("\n".join(answers)) ``` expected_prompt_fragment: beautiful number + metadata_fields: + - verifier_metadata rejected_generations: - |- ```python @@ -71,6 +76,8 @@ cases: expected_generations: - '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' expected_prompt_fragment: SHOW24 + metadata_fields: + - expected_action rejected_generations: - I cannot call the requested tool. expected_reward: 1.0 @@ -92,6 +99,11 @@ cases: expected_generations: - '\boxed{B}' expected_prompt_fragment: cystic fibrosis + metadata_fields: + - uuid + - options + - expected_answer + - grading_mode rejected_generations: - '\boxed{A}' expected_reward: 1.0 @@ -110,6 +122,8 @@ cases: expected_generations: - '\boxed{Charles Darwin}' expected_prompt_fragment: theory of evolution by natural selection + metadata_fields: + - expected_answer rejected_generations: - '\boxed{Ada Lovelace}' rejected_judge_verdicts: @@ -131,6 +145,15 @@ cases: expected_generations: - '{"name":"response_tool_8","arguments":{"extraction":{"name":"Dizer Kola","native_name":"ديزركلا","romanized_name":"Dizer Kola","settlement_type":"village","country":"Iran","province":"Mazandaran","county":"Nowshahr","bakhsh":"Central","rural_district":"Baladeh Kojur","coordinates":{"latitude":36.55694,"longitude":51.79389},"population_total":250,"population_year":2006,"number_of_families":64,"timezone_standard":"UTC+3:30 (IRST)","timezone_dst":"UTC+4:30 (IRDT)"}}}' expected_prompt_fragment: Dizer Kola + metadata_fields: + - schema_str + - schema_type + - response_mode + - problem_type + - schema_repr + - tool_name + - tool_schema_mode + - tool_payload_key rejected_generations: - I cannot extract that. expected_reward: 1.0 @@ -150,6 +173,11 @@ cases: - '{"name":"email_reply_email","arguments":{"email_id":"00000057","body":"Thanks for the update - I will get back to you tomorrow."}}' - Done expected_prompt_fragment: Task Update on Develop prototype + metadata_fields: + - id + - ground_truth + - category + - environment_name rejected_generations: - Done expected_reward: 1.0 diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 50cd9083287..0176f85f394 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -38,6 +38,8 @@ _REPO_ROOT = Path(__file__).parents[3] _GYM_ROOT = _REPO_ROOT / "3rdparty/Gym-workspace/Gym" _CASES_PATH = Path(__file__).parent / "nemo_gym_test_data/l0_rollout_acceptance.yaml" +# "L0" names the RL CI tier chosen for this test. The exercised boundary is +# the conceptual L2 contract: a real Gym environment rolls out through NeMo RL. _POLICY_MODEL_CONFIG = ( "responses_api_models/vllm_model/configs/vllm_model_for_training.yaml" ) @@ -371,6 +373,7 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: "expected_generations", "rejected_generations", "expected_prompt_fragment", + "metadata_fields", "expected_result", "expected_reward", } @@ -395,6 +398,8 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: f"{case['name']}: pinned example changed; review the row and update its golden values" ) assert case["agent_ref"].keys() >= {"type", "name"} + assert case["metadata_fields"] + assert all(isinstance(field, str) for field in case["metadata_fields"]) assert case["expected_generations"] assert case["rejected_generations"] assert all( @@ -523,9 +528,17 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte "expected_generations" if accepted else "rejected_generations" ] expected_reward = case["expected_reward"] if accepted else 0.0 + datum = _load_case_datum(case, accepted=accepted) + extra_env_info = datum["extra_env_info"] + expected_metadata = { + field: deepcopy(extra_env_info[field]) for field in case["metadata_fields"] + } + assert case["expected_prompt_fragment"] in json.dumps( + extra_env_info["responses_create_params"] + ) result = run_nemo_gym_rollout_sync( policy_generation=_ScriptedPolicyGeneration(), - input_batch=rl_collate_fn([_load_case_datum(case, accepted=accepted)]), + input_batch=rl_collate_fn([datum]), tokenizer=tokenizer, task_to_env={"nemo_gym": l0_nemo_gym}, max_seq_len=_GENERATION_CONFIG["max_total_sequence_length"], @@ -571,6 +584,10 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte full_result_table = result.rollout_metrics[full_result_key] assert len(full_result_table.data) == 1 full_result = json.loads(full_result_table.data[0][0]) + for field, expected_value in extra_env_info["responses_create_params"].items(): + assert full_result["responses_create_params"][field] == expected_value + for field, expected_value in expected_metadata.items(): + assert full_result[field] == expected_value if accepted: for field, expected_value in case["expected_result"].items(): assert full_result[field] == expected_value From 3545d7a51519ed02ac90cb271d437c7c94ab2b94 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 19:10:15 -0700 Subject: [PATCH 12/16] test: allow Gym response schema defaults Signed-off-by: yaoyu-33 --- .../test_l0_nemo_gym_rollout_acceptance.py | 40 ++++++++++++++++++- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 0176f85f394..238dfe36213 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -475,6 +475,40 @@ def _load_case_datum(case: dict[str, Any], *, accepted: bool) -> DatumSpec: return datum +def _assert_contract_preserved( + actual: Any, expected: Any, *, path: str = "responses_create_params" +) -> None: + """Require every source value while allowing Gym to add schema defaults.""" + if isinstance(expected, dict): + assert isinstance(actual, dict), ( + f"{path}: expected a mapping, got {type(actual).__name__}" + ) + missing_keys = expected.keys() - actual.keys() + assert not missing_keys, f"{path}: missing keys {sorted(missing_keys)}" + for key, expected_value in expected.items(): + _assert_contract_preserved( + actual[key], expected_value, path=f"{path}.{key}" + ) + return + + if isinstance(expected, list): + assert isinstance(actual, list), ( + f"{path}: expected a list, got {type(actual).__name__}" + ) + assert len(actual) == len(expected), ( + f"{path}: expected {len(expected)} items, got {len(actual)}" + ) + for index, (actual_value, expected_value) in enumerate( + zip(actual, expected, strict=True) + ): + _assert_contract_preserved( + actual_value, expected_value, path=f"{path}[{index}]" + ) + return + + assert actual == expected, f"{path}: expected {expected!r}, got {actual!r}" + + @pytest.fixture(scope="module") def scripted_openai_base_url(): server = ThreadingHTTPServer(("0.0.0.0", 0), _ScriptedOpenAIHandler) @@ -584,8 +618,10 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte full_result_table = result.rollout_metrics[full_result_key] assert len(full_result_table.data) == 1 full_result = json.loads(full_result_table.data[0][0]) - for field, expected_value in extra_env_info["responses_create_params"].items(): - assert full_result["responses_create_params"][field] == expected_value + _assert_contract_preserved( + full_result["responses_create_params"], + extra_env_info["responses_create_params"], + ) for field, expected_value in expected_metadata.items(): assert full_result[field] == expected_value if accepted: From 6b9e45094683ea62809aeffc646c46cb67148e4f Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 21:46:18 -0700 Subject: [PATCH 13/16] test: assert Gym verifier outputs Signed-off-by: yaoyu-33 --- .../l0_rollout_acceptance.yaml | 28 ------------------- .../test_l0_nemo_gym_rollout_acceptance.py | 8 ------ 2 files changed, 36 deletions(-) diff --git a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml index 65f557c3cfc..a852a6f3ff8 100644 --- a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml +++ b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml @@ -10,9 +10,6 @@ cases: expected_generations: - '\boxed{32}' expected_prompt_fragment: 1000 digit numbers - metadata_fields: - - question - - expected_answer rejected_generations: - '\boxed{999999}' expected_reward: 1.0 @@ -51,8 +48,6 @@ cases: sys.stdout.write("\n".join(answers)) ``` expected_prompt_fragment: beautiful number - metadata_fields: - - verifier_metadata rejected_generations: - |- ```python @@ -76,8 +71,6 @@ cases: expected_generations: - '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' expected_prompt_fragment: SHOW24 - metadata_fields: - - expected_action rejected_generations: - I cannot call the requested tool. expected_reward: 1.0 @@ -99,11 +92,6 @@ cases: expected_generations: - '\boxed{B}' expected_prompt_fragment: cystic fibrosis - metadata_fields: - - uuid - - options - - expected_answer - - grading_mode rejected_generations: - '\boxed{A}' expected_reward: 1.0 @@ -122,8 +110,6 @@ cases: expected_generations: - '\boxed{Charles Darwin}' expected_prompt_fragment: theory of evolution by natural selection - metadata_fields: - - expected_answer rejected_generations: - '\boxed{Ada Lovelace}' rejected_judge_verdicts: @@ -145,15 +131,6 @@ cases: expected_generations: - '{"name":"response_tool_8","arguments":{"extraction":{"name":"Dizer Kola","native_name":"ديزركلا","romanized_name":"Dizer Kola","settlement_type":"village","country":"Iran","province":"Mazandaran","county":"Nowshahr","bakhsh":"Central","rural_district":"Baladeh Kojur","coordinates":{"latitude":36.55694,"longitude":51.79389},"population_total":250,"population_year":2006,"number_of_families":64,"timezone_standard":"UTC+3:30 (IRST)","timezone_dst":"UTC+4:30 (IRDT)"}}}' expected_prompt_fragment: Dizer Kola - metadata_fields: - - schema_str - - schema_type - - response_mode - - problem_type - - schema_repr - - tool_name - - tool_schema_mode - - tool_payload_key rejected_generations: - I cannot extract that. expected_reward: 1.0 @@ -173,11 +150,6 @@ cases: - '{"name":"email_reply_email","arguments":{"email_id":"00000057","body":"Thanks for the update - I will get back to you tomorrow."}}' - Done expected_prompt_fragment: Task Update on Develop prototype - metadata_fields: - - id - - ground_truth - - category - - environment_name rejected_generations: - Done expected_reward: 1.0 diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index 238dfe36213..deb1e35b17f 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -373,7 +373,6 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: "expected_generations", "rejected_generations", "expected_prompt_fragment", - "metadata_fields", "expected_result", "expected_reward", } @@ -398,8 +397,6 @@ def _load_acceptance_cases() -> list[dict[str, Any]]: f"{case['name']}: pinned example changed; review the row and update its golden values" ) assert case["agent_ref"].keys() >= {"type", "name"} - assert case["metadata_fields"] - assert all(isinstance(field, str) for field in case["metadata_fields"]) assert case["expected_generations"] assert case["rejected_generations"] assert all( @@ -564,9 +561,6 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte expected_reward = case["expected_reward"] if accepted else 0.0 datum = _load_case_datum(case, accepted=accepted) extra_env_info = datum["extra_env_info"] - expected_metadata = { - field: deepcopy(extra_env_info[field]) for field in case["metadata_fields"] - } assert case["expected_prompt_fragment"] in json.dumps( extra_env_info["responses_create_params"] ) @@ -622,8 +616,6 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte full_result["responses_create_params"], extra_env_info["responses_create_params"], ) - for field, expected_value in expected_metadata.items(): - assert full_result[field] == expected_value if accepted: for field, expected_value in case["expected_result"].items(): assert full_result[field] == expected_value From 52328392efb0b666ff704f3818b9e446193fb409 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 22:23:59 -0700 Subject: [PATCH 14/16] fix: harden Gym acceptance lifecycle Signed-off-by: yaoyu-33 --- nemo_rl/environments/nemo_gym.py | 26 +++++- .../test_l0_nemo_gym_rollout_acceptance.py | 23 ++++- .../unit/environments/test_nemo_gym_utils.py | 86 +++++++++++++++++++ 3 files changed, 130 insertions(+), 5 deletions(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index b96ae1f5470..0551fdb129c 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -59,6 +59,7 @@ from nemo_rl.utils.venvs import make_actor_runtime_env NEMO_GYM_ACTOR_FQN = "nemo_rl.environments.nemo_gym.NemoGym" +NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S = 120 # Kept local so the Gym actor does not depend on model-config dtype resolution. # Must cover every name resolve_routed_experts_dtype can produce. @@ -1082,6 +1083,27 @@ def spinup_nemo_gym_actor( ) actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg) - ray.get(actor._spinup.remote()) - ray.get(actor.set_tokenizer.remote(tokenizer)) + try: + ray.get(actor._spinup.remote()) + ray.get(actor.set_tokenizer.remote(tokenizer)) + except Exception: + # _spinup can fail after RunHelper has started some Gym subprocesses. + # Ask the actor to reap anything it owns, then force-stop the actor as a + # final safety net. Cleanup errors must not hide the startup failure. + try: + ray.get( + actor.shutdown.remote(), + timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + ) + except Exception as cleanup_error: + print( + f"Warning: NeMo-Gym actor cleanup after startup failure failed: {cleanup_error}" + ) + try: + ray.kill(actor) + except Exception as kill_error: + print( + f"Warning: NeMo-Gym actor kill after startup failure failed: {kill_error}" + ) + raise return actor diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py index deb1e35b17f..6ed806f5dd3 100644 --- a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py +++ b/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py @@ -32,7 +32,10 @@ from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data.processors import nemo_gym_data_processor from nemo_rl.distributed.virtual_cluster import _get_node_ip_local -from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor +from nemo_rl.environments.nemo_gym import ( + NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + spinup_nemo_gym_actor, +) from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync _REPO_ROOT = Path(__file__).parents[3] @@ -543,7 +546,10 @@ def l0_nemo_gym(scripted_openai_base_url): try: # The shared actor owns many Gym subprocesses, which are reaped # sequentially during graceful shutdown. - ray.get(env.shutdown.remote(), timeout=60) + ray.get( + env.shutdown.remote(), + timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + ) finally: ray.kill(env) @@ -561,6 +567,17 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte expected_reward = case["expected_reward"] if accepted else 0.0 datum = _load_case_datum(case, accepted=accepted) extra_env_info = datum["extra_env_info"] + expected_responses_create_params = deepcopy( + extra_env_info["responses_create_params"] + ) + expected_responses_create_params["temperature"] = _GENERATION_CONFIG["temperature"] + expected_responses_create_params["top_p"] = _GENERATION_CONFIG["top_p"] + row_max_output_tokens = expected_responses_create_params.get("max_output_tokens") + expected_responses_create_params["max_output_tokens"] = ( + min(row_max_output_tokens, _GENERATION_CONFIG["max_new_tokens"]) + if row_max_output_tokens is not None + else _GENERATION_CONFIG["max_new_tokens"] + ) assert case["expected_prompt_fragment"] in json.dumps( extra_env_info["responses_create_params"] ) @@ -614,7 +631,7 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte full_result = json.loads(full_result_table.data[0][0]) _assert_contract_preserved( full_result["responses_create_params"], - extra_env_info["responses_create_params"], + expected_responses_create_params, ) if accepted: for field, expected_value in case["expected_result"].items(): diff --git a/tests/unit/environments/test_nemo_gym_utils.py b/tests/unit/environments/test_nemo_gym_utils.py index c562553ef93..6a9a26b102f 100644 --- a/tests/unit/environments/test_nemo_gym_utils.py +++ b/tests/unit/environments/test_nemo_gym_utils.py @@ -25,6 +25,7 @@ from nemo_rl.environments import nemo_gym as nemo_gym_mod from nemo_rl.environments.nemo_gym import ( NEMO_GYM_ACTOR_FQN, + NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, _detect_invalid_tool_call_and_malformed_thinking, build_nemo_gym_config, get_nemo_gym_uv_cache_dir, @@ -267,3 +268,88 @@ def test_spinup_nemo_gym_actor(detected_uv_dirs, num_gpu_nodes): actor._spinup.remote.assert_called_once_with() actor.set_tokenizer.remote.assert_called_once_with(tokenizer) assert mock_ray.get.call_args_list == [call("spinup-ref"), call("tokenizer-ref")] + + +@pytest.mark.parametrize("failed_ref", ["spinup-ref", "tokenizer-ref"]) +def test_spinup_nemo_gym_actor_cleans_up_after_startup_failure( + detected_uv_dirs, failed_ref +): + actor = MagicMock() + actor._spinup.remote.return_value = "spinup-ref" + actor.set_tokenizer.remote.return_value = "tokenizer-ref" + actor.shutdown.remote.return_value = "shutdown-ref" + + def get_or_fail(ref, **_kwargs): + if ref == failed_ref: + raise RuntimeError("startup failed") + return None + + with ( + patch.object(nemo_gym_mod, "make_actor_runtime_env", return_value={}), + patch.object(nemo_gym_mod, "NemoGym") as mock_cls, + patch.object(nemo_gym_mod, "ray") as mock_ray, + ): + mock_cls.options.return_value.remote.return_value = actor + mock_ray.get.side_effect = get_or_fail + + with pytest.raises(RuntimeError, match="startup failed"): + spinup_nemo_gym_actor( + _env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + tokenizer=MagicMock(), + enable_router_replay=False, + use_fastokens=False, + ) + + actor.shutdown.remote.assert_called_once_with() + mock_ray.kill.assert_called_once_with(actor) + assert ( + call("shutdown-ref", timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S) + in mock_ray.get.call_args_list + ) + + +@pytest.mark.parametrize("cleanup_failure", ["shutdown-remote", "shutdown-get", "kill"]) +def test_spinup_nemo_gym_actor_preserves_startup_error_when_cleanup_fails( + detected_uv_dirs, cleanup_failure +): + actor = MagicMock() + actor._spinup.remote.return_value = "spinup-ref" + actor.shutdown.remote.return_value = "shutdown-ref" + startup_error = RuntimeError("startup failed") + cleanup_error = RuntimeError("cleanup failed") + + if cleanup_failure == "shutdown-remote": + actor.shutdown.remote.side_effect = cleanup_error + + def get_or_fail(ref, **_kwargs): + if ref == "spinup-ref": + raise startup_error + if ref == "shutdown-ref" and cleanup_failure == "shutdown-get": + raise cleanup_error + return None + + with ( + patch.object(nemo_gym_mod, "make_actor_runtime_env", return_value={}), + patch.object(nemo_gym_mod, "NemoGym") as mock_cls, + patch.object(nemo_gym_mod, "ray") as mock_ray, + ): + mock_cls.options.return_value.remote.return_value = actor + mock_ray.get.side_effect = get_or_fail + if cleanup_failure == "kill": + mock_ray.kill.side_effect = cleanup_error + + with pytest.raises(RuntimeError) as exc_info: + spinup_nemo_gym_actor( + _env_configs(), + base_urls=["http://vllm-0"], + model_name="test-model", + tokenizer=MagicMock(), + enable_router_replay=False, + use_fastokens=False, + ) + + assert exc_info.value is startup_error + actor.shutdown.remote.assert_called_once_with() + mock_ray.kill.assert_called_once_with(actor) From b897b7a3eae1826e8c269876007820e9f27b0a03 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 5 Sep 2026 23:31:48 -0700 Subject: [PATCH 15/16] test: isolate Gym actor cleanup failures Signed-off-by: yaoyu-33 --- tests/unit/environments/test_nemo_gym_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/environments/test_nemo_gym_utils.py b/tests/unit/environments/test_nemo_gym_utils.py index 806d166e64f..6103bd3bb3e 100644 --- a/tests/unit/environments/test_nemo_gym_utils.py +++ b/tests/unit/environments/test_nemo_gym_utils.py @@ -311,7 +311,7 @@ def get_or_fail(ref, **_kwargs): with pytest.raises(RuntimeError, match="startup failed"): spinup_nemo_gym_actor( - _env_configs(), + _env_configs(num_gpu_nodes=0), base_urls=["http://vllm-0"], model_name="test-model", tokenizer=MagicMock(), @@ -359,7 +359,7 @@ def get_or_fail(ref, **_kwargs): with pytest.raises(RuntimeError) as exc_info: spinup_nemo_gym_actor( - _env_configs(), + _env_configs(num_gpu_nodes=0), base_urls=["http://vllm-0"], model_name="test-model", tokenizer=MagicMock(), From b8588ae999f1b3f447aa1e8f0d14f55719ddd351 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Tue, 8 Sep 2026 11:42:36 -0700 Subject: [PATCH 16/16] test: use tier-neutral NeMo Gym rollout module name Signed-off-by: yaoyu-33 --- examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh | 2 +- ..._nemo_gym_rollout_acceptance.py => test_nemo_gym_rollout.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename tests/unit/environments/{test_l0_nemo_gym_rollout_acceptance.py => test_nemo_gym_rollout.py} (100%) diff --git a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh index c3b1571e468..6e600ec023f 100755 --- a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh +++ b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh @@ -52,5 +52,5 @@ uv run python -c "import ray; ray.shutdown()" # Exercise one pinned example from each L0 Gym environment through the real NeMo RL rollout boundary. ./tests/run_unit.sh \ - unit/environments/test_l0_nemo_gym_rollout_acceptance.py::test_l0_gym_environments_roll_out_through_nemo_rl \ + unit/environments/test_nemo_gym_rollout.py::test_l0_gym_environments_roll_out_through_nemo_rl \ --nemo-gym-only diff --git a/tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py b/tests/unit/environments/test_nemo_gym_rollout.py similarity index 100% rename from tests/unit/environments/test_l0_nemo_gym_rollout_acceptance.py rename to tests/unit/environments/test_nemo_gym_rollout.py