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..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 @@ -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 example from each L0 Gym environment through the real NeMo RL rollout boundary. +./tests/run_unit.sh \ + unit/environments/test_nemo_gym_rollout.py::test_l0_gym_environments_roll_out_through_nemo_rl \ + --nemo-gym-only diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 7415aaa985c..511c3751459 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -60,6 +60,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. @@ -1387,6 +1388,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/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..a852a6f3ff8 --- /dev/null +++ b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml @@ -0,0 +1,158 @@ +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 + rejected_generations: + - '\boxed{999999}' + 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 + rejected_generations: + - |- + ```python + print(0) + ``` + 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 + 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 + 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 + rejected_generations: + - '\boxed{A}' + 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 + rejected_generations: + - '\boxed{Ada Lovelace}' + rejected_judge_verdicts: + - "[[A!=B]]" + 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 + rejected_generations: + - I cannot extract that. + 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 + rejected_generations: + - Done + expected_reward: 1.0 + expected_result: {} + expected_tool_outputs: + - '{"output":"Email replied successfully."}' diff --git a/tests/unit/environments/test_nemo_gym_rollout.py b/tests/unit/environments/test_nemo_gym_rollout.py new file mode 100644 index 00000000000..6ed806f5dd3 --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_rollout.py @@ -0,0 +1,678 @@ +# 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 + +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 +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 ( + 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] +_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" +) +_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", + "max_new_tokens": 1024, + # 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, + "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"}}' +) +_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.", +} +_BYTE_TOKEN_OFFSET = 3 +_CONTINUATION_MARKER = "\x02nemo-rl-continuation\x00" +_REJECTED_ROLLOUT_MARKER = "[NEMO_RL_L0_REJECTED_ROLLOUT]" + + +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 _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=(",", ":")) + prompt_token_ids = _text_token_ids(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 a marked 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=(",", ":")) + _, first_generation = _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + return ( + _text_token_ids(initial_request) + + _text_token_ids(first_generation) + + _text_token_ids(_CONTINUATION_MARKER + serialized_request) + ) + + +class _ScriptedPolicyGeneration: + """Minimum policy-generation surface consumed by the Gym rollout path.""" + + cfg = _GENERATION_CONFIG + + +class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): + """Return deterministic policy outputs to real Gym model and agent services.""" + + protocol_version = "HTTP/1.1" + + @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["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) + 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 + 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: + message, _ = _tool_call_message( + "check_seat_availability", + json.loads(_TOOL_CALL_ARGUMENTS), + "call-l0-tool-use", + ) + 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": + self.send_error(404) + return + + 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 = _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, + "generation_token_ids": generation_token_ids, + "generation_log_probs": [-0.1] * len(generation_token_ids), + } + ) + response = { + "id": "chatcmpl-l0-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, "L0 Gym rollout acceptance matrix must not be empty" + required_fields = { + "name", + "config_path", + "data_path", + "example_index", + "example_sha256", + "agent_ref", + "expected_generations", + "rejected_generations", + "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_generations"] + assert case["rejected_generations"] + assert all( + isinstance(generation, str) + for key in ("expected_generations", "rejected_generations") + for generation in case[key] + ) + 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_L0_CASES, ( + f"L0 matrix must contain exactly {_REQUIRED_L0_CASES}, got {set(names)}" + ) + return 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) + _text_token_ids(first_generation) + + continuation_prompt = _prompt_token_ids(continuation_body) + + assert continuation_prompt[: len(seen_token_ids)] == seen_token_ids + 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 + ) + + +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 + + +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) + 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(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 = ByT5Tokenizer() + 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: + # The shared actor owns many Gym subprocesses, which are reaped + # sequentially during graceful shutdown. + ray.get( + env.shutdown.remote(), + timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + ) + finally: + ray.kill(env) + + +@pytest.mark.nemo_gym +@pytest.mark.timeout(900) +@pytest.mark.parametrize("case", _L0_CASES, ids=[case["name"] for case in _L0_CASES]) +@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 + 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"] + ) + result = run_nemo_gym_rollout_sync( + policy_generation=_ScriptedPolicyGeneration(), + 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"], + 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(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 + assert len(assistant_messages) == len(expected_generations) + for message, generation in zip( + assistant_messages, expected_generations, strict=True + ): + assert isinstance(message["token_ids"], torch.Tensor) + assert isinstance(message["generation_logprobs"], torch.Tensor) + 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"], + 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) + + 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]) + _assert_contract_preserved( + full_result["responses_create_params"], + expected_responses_create_params, + ) + if accepted: + for field, expected_value in case["expected_result"].items(): + assert full_result[field] == expected_value + assert full_result["response"]["output"] + 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"] + 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 == expected_judge_verdicts + generation_strings = [ + item["generation_str"] + for item in full_result["response"]["output"] + if "generation_str" in item + ] + 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(expected_generations) + assert all(case["expected_prompt_fragment"] in prompt for prompt in prompt_strings) + 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"] + 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 + ) diff --git a/tests/unit/environments/test_nemo_gym_utils.py b/tests/unit/environments/test_nemo_gym_utils.py index 09dca80f852..6103bd3bb3e 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, @@ -284,3 +285,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(num_gpu_nodes=0), + 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(num_gpu_nodes=0), + 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)