diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index b11414c71b1..5ed97996cc2 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit b11414c71b15e54d333eb49346ed199f20fa9021 +Subproject commit 5ed97996cc2b422904d18179375b6d7366915097 diff --git a/nemo_rl/distributed/held_port.py b/nemo_rl/distributed/held_port.py index 18d68be06b6..9c493758120 100644 --- a/nemo_rl/distributed/held_port.py +++ b/nemo_rl/distributed/held_port.py @@ -39,6 +39,7 @@ def receive_held_socket(port: int) -> socket.socket: fds: list[int] = [] received_socket: socket.socket try: + # Retrieve the reserved socket fd's from HeldPortReservation. client.connect(_held_port_uds_name(port)) _, fds, _, _ = socket.recv_fds(client, 1024, 1) if not fds: @@ -50,10 +51,12 @@ def receive_held_socket(port: int) -> socket.socket: # reservation socket and fail with EADDRINUSE. released = client.recv(1) if released != _HANDOFF_RELEASED: + # Unexpected contract, the reserved socket should be released. socket.close(fds.pop()) raise RuntimeError( f"Port holder for port {port} did not confirm releasing its socket." ) + # Manufacture a new socket from the FD. received_socket = socket.socket(fileno=fds.pop()) except OSError as e: for fd in fds: @@ -95,8 +98,10 @@ def address(self) -> tuple[str, int]: return self._node_ip, self._port def _serve_fd_once(self) -> None: + # Wait until connection. conn, _ = self._uds.accept() try: + # Send reserved TCP socket across Unix socket for hand-over. socket.send_fds(conn, [b"s"], [self._sock.fileno()]) # The receiver holds a duplicate fd, so the port remains reserved. # Close this copy before acknowledging the handoff; the receiver may diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 2df187b58c4..7415aaa985c 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -341,6 +341,9 @@ def __init__(self, cfg: NemoGymConfig): # here rather than in _spinup so a second spinup cannot wipe an installed # tokenizer and then report that set_tokenizer was never called. self._tokenizer: Optional[PreTrainedTokenizerBase] = None + # _spinup replaces this from cfg. Keep restarted/unspun actors internally + # complete so diagnostics and focused tests do not fail with AttributeError. + self._token_capture_enabled = False self._pad_dynamic_image_shapes = bool(cfg.get("pad_dynamic_image_shapes")) # Reconstruct the processor inside the actor (rather than serializing it # per rollout call) for full-trajectory multimodal postprocessing. @@ -587,7 +590,7 @@ async def run_rollouts( nemo_gym_examples: list[dict], timer_prefix: str, deduplicate_multimodal_data: bool = False, - ) -> AsyncGenerator[tuple[int, dict, dict | None], None]: + ) -> AsyncGenerator[tuple[int, dict, dict, dict | None], None]: """Stream postprocessed rollouts as NeMo-Gym tasks complete.""" self._require_spinup() if not nemo_gym_examples: @@ -602,18 +605,20 @@ async def run_rollouts( maybe_patch_fastokens(bool(self.cfg.get("use_fastokens"))) - timer = Timer() - counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) - # Normalize local media before shipping requests to vLLM. Helper is a no-op # for text-only rows and already-qualified URLs. # Megatron's HTTP backend consumes the same normalized Responses payload. normalize_media_in_examples(nemo_gym_examples) + timer = Timer() timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( examples=nemo_gym_examples, head_server_config=self.head_server_config ) + # Current Gym collates data with ``task_source`` rather than a baked-in + # ``agent_ref``. ``run_examples`` resolves that routing synchronously and + # stamps each input row before returning its result iterator. + counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) num_results = 0 for task in nemo_gym_result_iterator: @@ -681,7 +686,15 @@ async def run_rollouts( file=sys.stderr, ) - yield nemo_gym_row["_rowidx"], nemo_rl_result, timing_metrics + # task_source is resolved to agent_ref inside this Ray actor, after + # the caller's row was serialized. Return the resolved ref explicitly + # so the caller can hydrate its own row copy before postprocessing. + yield ( + nemo_gym_row["_rowidx"], + nemo_gym_row["agent_ref"], + nemo_rl_result, + timing_metrics, + ) async def _postprocess_receipt_mode( self, nemo_gym_row: dict, nemo_gym_result: dict diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index ed1cb68cfc9..b2e16b7cd06 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -858,10 +858,22 @@ async def run_rollout( timer.stop(f"{timer_prefix}/total") rollout_metrics.update(timer.get_timing_metrics("sum")) + resolved_agent_ref = rollout_inputs[0].get("agent_ref") + if not isinstance(resolved_agent_ref, dict): + raise ValueError("NeMo-Gym did not return a resolved agent_ref") + if any( + row.get("agent_ref") != resolved_agent_ref for row in rollout_inputs[1:] + ): + raise ValueError( + "NeMo-Gym resolved one prompt group to inconsistent agent_ref values" + ) + record_extra_env_info = copy.deepcopy(input_sample["extra_env_info"]) + record_extra_env_info["agent_ref"] = copy.deepcopy(resolved_agent_ref) + return PromptGroupRecord( prompt_idx=input_sample["idx"], prompt=prompt_message_log, - extra_env_info=input_sample["extra_env_info"], + extra_env_info=record_extra_env_info, metadata={"task_name": "nemo_gym"}, completions=completions, rollout_metrics=rollout_metrics, @@ -954,13 +966,14 @@ async def _stream_rows( The environment's timing metrics, or None if the stream ended without them. """ dispatched = {row["_rowidx"] for row in pending} + pending_by_rowidx = {row["_rowidx"]: row for row in pending} received: set[int] = set() env_timing_metrics: Optional[dict[str, Any]] = None async for result_ref in nemo_gym_env.run_rollouts.options( num_returns="streaming" ).remote(pending, timer_prefix): - rowidx, result, timing_metrics = await result_ref + rowidx, resolved_agent_ref, result, timing_metrics = await result_ref # Validated against the original group, not the pending subset: on a # re-dispatch the row keeps its original index so results stay ordered. if not isinstance(rowidx, int) or not 0 <= rowidx < total_rows: @@ -976,6 +989,7 @@ async def _stream_rows( if rowidx in received: raise ValueError(f"NeMo-Gym returned duplicate row index {rowidx}") received.add(rowidx) + pending_by_rowidx[rowidx]["agent_ref"] = resolved_agent_ref results[rowidx] = result if timing_metrics is not None: env_timing_metrics = timing_metrics diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index c17148fe549..a0688be3620 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -1857,7 +1857,12 @@ def __init__( def is_complete(self) -> bool: return len(self._received_row_indices) == len(self._rows) - def add(self, row_index: int, result: dict) -> _CompletedNemoGymGroup | None: + def add( + self, + row_index: int, + result: dict, + resolved_agent_ref: dict, + ) -> _CompletedNemoGymGroup | None: """Add one streamed row and return its group when that group is complete.""" if not isinstance(row_index, int): raise TypeError( @@ -1871,6 +1876,7 @@ def add(self, row_index: int, result: dict) -> _CompletedNemoGymGroup | None: if row_index in self._received_row_indices: raise ValueError(f"NeMo-Gym returned duplicate row index {row_index}") + self._rows[row_index]["agent_ref"] = resolved_agent_ref self._received_row_indices.add(row_index) group_index = row_index // self._num_generations group_results = self._pending_results[group_index] @@ -2616,14 +2622,14 @@ async def run_async_nemo_gym_rollout( except StopAsyncIteration: stream_finished = True else: - rowidx, result, timing_metrics = await future + rowidx, resolved_agent_ref, result, timing_metrics = await future # Measure the received streaming Ray value in the caller. In # async training this runs in the collector actor; validation # runs in the driver, so the two phases cannot share a metric # accumulator even when they share the NeMo-Gym actor. print_multimodal_payload_metrics( collect_multimodal_payload_metrics( - (rowidx, result, timing_metrics), + (rowidx, resolved_agent_ref, result, timing_metrics), "nemo_gym_return", enabled=debug_payload_metrics, ) @@ -2634,7 +2640,9 @@ async def run_async_nemo_gym_rollout( actor_timing_metrics = timing_metrics _tensorize_nemo_gym_result(result) - completed_group = accumulator.add(rowidx, result) + completed_group = accumulator.add( + rowidx, result, resolved_agent_ref=resolved_agent_ref + ) if original_message_logs is not None: _reattach_static_multimodal_payloads_to_result( result, original_message_logs[rowidx] diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 590a8cb44f1..882544e4347 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -57,6 +57,7 @@ from nemo_rl.data.multimodal_utils import CACHED_VIDEO_FRAME_MANIFEST_MAGIC from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.held_port import receive_held_socket from nemo_rl.models.generation.interfaces import ( GenerationDatumSpec, GenerationOutputSpec, @@ -90,7 +91,7 @@ class MegatronGenerationMixin: - megatron_tokenizer: tokenizer for inference. - processor: optional multimodal processor. - is_generation_colocated: Whether colocated or distributed. - - _reserved_http_server_socket: driver-reserved server socket, or None. + - _reserved_http_server_port: driver-reserved server port, or None. """ # Colocated-reshard hosts assign the dedicated inference-layout model here @@ -554,10 +555,16 @@ def _setup_openai_api_server(self) -> str: ) ip = _get_node_ip_local() - reserved_socket = self._reserved_http_server_socket - if reserved_socket is not None: + reserved_port = self._reserved_http_server_port + if reserved_port is not None: + # Defer socket handoff until immediately before server startup. + # Holding this listener across model initialization can fork + # into a persistent child process, which then receives Gym + # traffic despite never accepting HTTP requests. + reserved_socket = receive_held_socket(reserved_port) server_port = reserved_socket.getsockname()[1] else: + reserved_socket = None server_port = _get_free_port_local() start_text_gen_server( diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index a8fd91e9bd4..93dcfbae127 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -16,7 +16,6 @@ import logging import os import re -import socket import time import warnings from collections import OrderedDict, defaultdict @@ -58,7 +57,6 @@ ) from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.held_port import receive_held_socket from nemo_rl.distributed.named_sharding import NamedSharding from nemo_rl.models.generation.interfaces import GenerationDatumSpec from nemo_rl.models.generation.megatron.megatron_worker import ( @@ -481,14 +479,12 @@ def __init__( self.rank = get_rank_safe() self.timer = Timer(context={"worker": "megatron_policy", "rank": self.rank}) - # Adopt the driver-reserved OpenAI server socket before any heavy init. - # The port holder has kept it bound and listening since reservation, so - # there was no window in which the pre-published URL could be stolen. - self._reserved_http_server_socket: Optional[socket.socket] = None - if reserved_http_server_port is not None and self.rank == 0: - self._reserved_http_server_socket = receive_held_socket( - reserved_http_server_port - ) + # Store the reserved HTTP server port for inference server initialization. + # Megatron-LLM's inference server lives on Rank 0 only. + # TODO: Multiple inference servers for each MP coordinator. + self._reserved_http_server_port = ( + reserved_http_server_port if self.rank == 0 else None + ) # Step 1: Setup distributed setup_distributed(config) @@ -1462,16 +1458,12 @@ def train_microbatch( any of the packing/CP capability flags, so a multimodal model would silently be handed CP-sliced rows it believes are full. """ - if self.media_placeholder_token_id is not None or ( - self.model_slices_context_parallel_inputs - ): + if self.media_placeholder_token_id is not None: raise NotImplementedError( "train_microbatch does not support multimodal models: its " "microbatch iterator is built without " - "attach_media_token_validity_mask, delegate_pack_to_model, " - "delegate_mtp_loss_mask_to_model or " - "model_slices_context_parallel_inputs, all of which the train / " - "get_logprobs / get_topk_logits paths pass. Threading them here " + "attach_media_token_validity_mask, which the train / " + "get_logprobs / get_topk_logits paths pass. Threading it here " "needs a SingleController VLM recipe to verify against; until " "then use train_presharded, which delegates to train." ) diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index e8488f5d1b7..5922fa9016d 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -181,7 +181,7 @@ run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller # Token-capture (gate-authoritative) path: same SC+Gym smoke with the gate # custodying token lineage and the finalizer publishing training rows. -run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh token_capture.enabled=true +run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh ++token_capture.enabled=true cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 0e41d9e8ce5..c603b188915 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -1600,7 +1600,8 @@ def _postprocess_nemo_gym_to_nemo_rl_result( assert postprocess_calls == [(nemo_gym_row, nemo_gym_result, tokenizer, True)] assert streamed_results[0][0] == 7 - assert streamed_results[0][1] == {"message_log": []} + assert streamed_results[0][1] == nemo_gym_row["agent_ref"] + assert streamed_results[0][2] == {"message_log": []} asyncio.run(_run()) @@ -1697,8 +1698,9 @@ def _require_spinup(self): ): streamed.append(item) - row_index, result, _metrics = streamed[0] + row_index, agent_ref, result, _metrics = streamed[0] assert row_index == 3 + assert agent_ref == row["agent_ref"] assert [message["role"] for message in result["message_log"]] == [ "user", "assistant", @@ -1804,7 +1806,7 @@ def test_nemo_gym_sanity( for result_ref in nemo_gym.run_rollouts.options(num_returns="streaming").remote( nemo_gym_sanity_test_data["input"], "" ): - rowidx, result, _ = ray.get(result_ref) + rowidx, _agent_ref, result, _ = ray.get(result_ref) actual_result[rowidx] = result expected_result = nemo_gym_sanity_test_data["expected_output"] diff --git a/tests/unit/environments/test_nemo_gym_health.py b/tests/unit/environments/test_nemo_gym_health.py index 41f5d0c4cc8..4d495483499 100644 --- a/tests/unit/environments/test_nemo_gym_health.py +++ b/tests/unit/environments/test_nemo_gym_health.py @@ -24,6 +24,10 @@ state and previously surfaced it as an AttributeError from deep inside a rollout. """ +import asyncio +from collections.abc import AsyncIterator +from typing import Any + import pytest from nemo_rl.environments.nemo_gym import NemoGym @@ -54,6 +58,39 @@ def shutdown(self) -> None: self.shutdowns += 1 +class _TaskSourceResolvingRolloutHelper: + """Mimic Gym's synchronous task_source-to-agent_ref resolution.""" + + def run_examples( + self, examples: list[dict[str, Any]], head_server_config: str + ) -> list[Any]: + assert head_server_config == "head-server" + assert all("agent_ref" not in example for example in examples) + for example in examples: + example["agent_ref"] = { + "type": "responses_api_agents", + "name": "workplace_assistant_simple_agent", + } + return [] + + +class _TaskSourceResolvingRolloutHelperWithResult(_TaskSourceResolvingRolloutHelper): + def run_examples( + self, examples: list[dict[str, Any]], head_server_config: str + ) -> list[Any]: + super().run_examples(examples, head_server_config) + + async def completed(example): + return example, {} + + return [completed(example) for example in examples] + + +async def _drain(async_generator: AsyncIterator[Any]) -> None: + async for _ in async_generator: + pass + + class TestHealthCheck: def test_a_healthy_gym_polls_the_run_helper(self): env = _unspun() @@ -96,3 +133,44 @@ def test_shutdown_still_forwards_when_spun_up(self): env.shutdown() assert run_helper.shutdowns == 1 assert env.rh is None + + +def test_run_rollouts_resolves_task_source_before_reading_agent_ref(): + """Gym 0.15 collated rows are task_source-routed until run_examples.""" + env = _unspun() + env.rh = _FakeRunHelper() + env._tokenizer = object() + env.head_server_config = "head-server" + env.rch = _TaskSourceResolvingRolloutHelper() + + rows = [{"task_source": "workplace_assistant"}] + asyncio.run(_drain(env.run_rollouts(rows, "timing/test"))) + + assert rows[0]["agent_ref"]["name"] == "workplace_assistant_simple_agent" + + +def test_run_rollouts_echoes_resolved_agent_ref_with_streamed_result(): + """The caller's serialized row copy cannot observe actor-local mutation.""" + env = _unspun() + env.rh = _FakeRunHelper() + env._tokenizer = object() + env.head_server_config = "head-server" + env.rch = _TaskSourceResolvingRolloutHelperWithResult() + env._postprocess_nemo_gym_to_nemo_rl_result = lambda *_args, **_kwargs: { + "message_log": [] + } + + async def collect(): + return [ + item + async for item in env.run_rollouts( + [{"task_source": "workplace_assistant", "_rowidx": 0}], + "timing/test", + ) + ] + + streamed = asyncio.run(collect()) + assert streamed[0][1] == { + "type": "responses_api_agents", + "name": "workplace_assistant_simple_agent", + } diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index c545e73cf8a..f6317d3a8ec 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -440,8 +440,8 @@ async def _body(): def _gym_rows(count: int) -> list[dict]: - """Rows shaped the way _build_inputs stamps them: each carries its own index.""" - return [{"_rowidx": i, "agent_ref": {"name": "agent"}} for i in range(count)] + """Gym 0.15 rows carry task_source until the remote actor resolves an agent.""" + return [{"_rowidx": i, "task_source": "workplace_assistant"} for i in range(count)] class _PartialGymMethod: @@ -480,6 +480,7 @@ async def _row_result(rowidx: int): """A minimally complete NeMo-Gym result, enough to build a Completion.""" return ( rowidx, + {"name": "agent"}, { "input_message_log": [{"role": "user", "token_ids": [1]}], "message_log": [{"role": "assistant", "token_ids": [2]}], @@ -507,7 +508,12 @@ def remote(self, inputs, timer_prefix): async def _stream(self, num_inputs): async def _result(rowidx): - return rowidx, {"input_message_log": [], "message_log": []}, None + return ( + rowidx, + {"name": "agent"}, + {"input_message_log": [], "message_log": []}, + None, + ) for rowidx in range(min(self._rows_to_yield, num_inputs)): yield _result(rowidx) diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 56f03815736..3891e2564d4 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1877,7 +1877,7 @@ def remote( ], } timing = {"timing/remote": 1.0} if position == 3 else None - values.append(_ReadyRef((rowidx, result, timing))) + values.append(_ReadyRef((rowidx, {"name": "agent"}, result, timing))) return _AsyncOnlyStream(values) class _PolicyGeneration: @@ -1887,7 +1887,7 @@ class _PolicyGeneration: for task_index in (10, 10, 11, 11): rows.append( { - "agent_ref": {"name": "agent"}, + "task_source": "workplace_assistant", "responses_create_params": {}, "_ng_task_index": task_index, } @@ -1921,6 +1921,9 @@ class _PolicyGeneration: def _postprocess_group(**kwargs): assert kwargs["log_full_result_tables"] is False + assert all( + row["agent_ref"] == {"name": "agent"} for row in kwargs["nemo_gym_rows"] + ) task_index = int(kwargs["nemo_gym_rows"][0]["_ng_task_index"]) for result, original_log in zip( kwargs["results"], kwargs["input_batch"]["message_log"] @@ -2013,7 +2016,8 @@ async def _collect(): assert boundary == "nemo_gym_return" assert enabled is True assert payload[0] == expected_rowidx - assert payload[1]["rowidx"] == expected_rowidx + assert payload[1] == {"name": "agent"} + assert payload[2]["rowidx"] == expected_rowidx assert all( captured_video_payloads[rowidx] is video_payloads[rowidx] for rowidx in range(4) ) @@ -2024,8 +2028,8 @@ async def _collect(): def test_nemo_gym_stream_accumulator_validates_rows_and_completion(): rows = [ - {"agent_ref": {"name": "agent"}}, - {"agent_ref": {"name": "agent"}}, + {"task_source": "workplace_assistant"}, + {"task_source": "workplace_assistant"}, ] accumulator = rollouts_mod._NemoGymStreamAccumulator( rows=rows, @@ -2033,33 +2037,46 @@ def test_nemo_gym_stream_accumulator_validates_rows_and_completion(): allow_mixed_agents=False, ) - assert accumulator.add(0, {"row": 0}) is None + resolved_agent_ref = { + "type": "responses_api_agents", + "name": "workplace_assistant_simple_agent", + } + assert accumulator.add(0, {"row": 0}, resolved_agent_ref=resolved_agent_ref) is None with pytest.raises(ValueError, match="duplicate row index 0"): - accumulator.add(0, {"row": 0}) + accumulator.add(0, {"row": 0}, resolved_agent_ref=resolved_agent_ref) with pytest.raises(RuntimeError, match=r"missing row indices \[1\]"): accumulator.finish() + completed = accumulator.add(1, {"row": 1}, resolved_agent_ref=resolved_agent_ref) + assert completed is not None + assert [row["agent_ref"] for row in completed.rows] == [ + resolved_agent_ref, + resolved_agent_ref, + ] + with pytest.raises(ValueError, match="outside the expected range"): rollouts_mod._NemoGymStreamAccumulator( rows=rows, num_generations=2, allow_mixed_agents=False, - ).add(2, {"row": 2}) + ).add(2, {"row": 2}, resolved_agent_ref=resolved_agent_ref) def test_nemo_gym_stream_accumulator_rejects_mixed_agent_group(): accumulator = rollouts_mod._NemoGymStreamAccumulator( rows=[ - {"agent_ref": {"name": "agent-a"}}, - {"agent_ref": {"name": "agent-b"}}, + {"task_source": "task"}, + {"task_source": "task"}, ], num_generations=2, allow_mixed_agents=False, ) - assert accumulator.add(0, {"row": 0}) is None + assert ( + accumulator.add(0, {"row": 0}, resolved_agent_ref={"name": "agent-a"}) is None + ) with pytest.raises(ValueError, match="one NeMo-Gym agent"): - accumulator.add(1, {"row": 1}) + accumulator.add(1, {"row": 1}, resolved_agent_ref={"name": "agent-b"}) @pytest.mark.parametrize("log_full_result_tables", [False, True]) @@ -2164,6 +2181,7 @@ def __init__(self): _ReadyRef( ( 1, + {"name": "agent"}, { "value": "second", "input_message_log": [ @@ -2176,6 +2194,7 @@ def __init__(self): _ReadyRef( ( 0, + {"name": "agent"}, { "value": "first", "input_message_log": [ @@ -2228,8 +2247,8 @@ def remote(self, inputs, timer_prefix): completions, prompt_message_log, metrics = asyncio.run( manager._run_rollouts( inputs=[ - {"_rowidx": 0, "agent_ref": {"name": "agent"}}, - {"_rowidx": 1, "agent_ref": {"name": "agent"}}, + {"_rowidx": 0, "task_source": "workplace_assistant"}, + {"_rowidx": 1, "task_source": "workplace_assistant"}, ], timer=rollouts_mod.Timer(), timer_prefix="timing/test", @@ -2246,6 +2265,44 @@ def remote(self, inputs, timer_prefix): } +def test_nemo_gym_rollout_record_persists_runtime_resolved_agent_ref(): + manager = object.__new__(AsyncNemoGymRolloutImpl) + manager._num_generations_per_prompt = 2 + manager._generation_config = { + "temperature": 1.0, + "top_p": 1.0, + "max_new_tokens": 32, + } + + resolved_agent_ref = { + "type": "responses_api_agents", + "name": "workplace_assistant_simple_agent", + } + + async def _run_rollouts(inputs, timer, timer_prefix): + del timer, timer_prefix + for row in inputs: + row["agent_ref"] = resolved_agent_ref + receipt_completion = SimpleNamespace(env_extras={"ng_receipt": {}}) + return [receipt_completion, receipt_completion], [], {} + + manager._run_rollouts = _run_rollouts + input_sample = { + "idx": 4, + "message_log": [], + "extra_env_info": { + "task_source": "workplace_assistant", + "responses_create_params": {}, + }, + } + + record = asyncio.run(manager.run_rollout(input_sample)) + + assert "agent_ref" not in input_sample["extra_env_info"] + assert record.extra_env_info["task_source"] == "workplace_assistant" + assert record.extra_env_info["agent_ref"] == resolved_agent_ref + + def test_rollout_manager_rejects_duplicate_stream_rows(): class _ReadyRef: def __init__(self, value): @@ -2261,8 +2318,8 @@ class _DuplicateStream: def __init__(self): self.values = iter( [ - _ReadyRef((0, {"value": "first"}, None)), - _ReadyRef((0, {"value": "duplicate"}, None)), + _ReadyRef((0, {"name": "agent"}, {"value": "first"}, None)), + _ReadyRef((0, {"name": "agent"}, {"value": "duplicate"}, None)), ] ) diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index f996690deef..f5ccdcb4b6f 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -513,6 +513,12 @@ def test_megatron_policy_generation( f"tp={tensor_parallel_size} pp={pipeline_parallel_size}" ) + if pipeline_parallel_size > 1: + pytest.xfail( + "FIXME(@cspades/@tdene): MCore async-scheduled generation segfaults with PP>1 " + "in dynamic_context.calculate_log_probs_tensors when slicing log_probs." + ) + config = deepcopy(basic_megatron_test_config) config["megatron_cfg"]["tensor_model_parallel_size"] = tensor_parallel_size config["megatron_cfg"]["pipeline_model_parallel_size"] = pipeline_parallel_size diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py index 062f9b2edb6..e177cf74914 100644 --- a/tests/unit/models/generation/test_megatron_generation_parse.py +++ b/tests/unit/models/generation/test_megatron_generation_parse.py @@ -38,10 +38,7 @@ ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.distributed.held_port import ( - HeldPortReservation, - receive_held_socket, -) +from nemo_rl.distributed.held_port import HeldPortReservation from nemo_rl.models.generation.megatron.megatron_worker import ( MegatronGenerationMixin, ) @@ -196,60 +193,59 @@ def test_http_server_port_reservation(monkeypatch): with socket.create_connection(("127.0.0.1", port), timeout=5): pass - # Worker-side adoption: the same live socket, duplicated across the - # process boundary. The holder confirms that its original descriptor is - # closed before this returns, preventing MCore's SO_REUSEPORT listeners - # from racing the old non-reusable socket. - reserved = receive_held_socket(port) - try: - assert holder._sock.fileno() == -1 - assert reserved.getsockname()[1] == port - with socket.create_connection(("127.0.0.1", port), timeout=5): - pass - - # MCore closes the handed-off fd and gives every frontend replica its - # own SO_REUSEPORT listener. Verify that such a listener can join the - # reservation's reuse group even before this duplicate is closed. - replica = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - try: - replica.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - replica.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - replica.bind(("0.0.0.0", port)) - finally: - replica.close() - - # Server start with the network and MLM server stubbed out. - started = {} - monkeypatch.setattr( - mlm_text_gen_server, - "start_text_gen_server", - lambda **kwargs: started.update(kwargs), - ) - monkeypatch.setattr( - "nemo_rl.distributed.virtual_cluster._get_free_port_local", - lambda: 12345, - ) - monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) - requests_mock = MagicMock() - health_get = requests_mock.Session.return_value.__enter__.return_value.get - health_get.return_value.status_code = 200 - monkeypatch.setattr( - "nemo_rl.models.generation.megatron.megatron_worker.requests", - requests_mock, - ) + # Server start with the network and MLM server stubbed out. + started = {} + monkeypatch.setattr( + mlm_text_gen_server, + "start_text_gen_server", + lambda **kwargs: started.update(kwargs), + ) + monkeypatch.setattr( + "nemo_rl.distributed.virtual_cluster._get_free_port_local", + lambda: 12345, + ) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + requests_mock = MagicMock() + health_get = requests_mock.Session.return_value.__enter__.return_value.get + health_get.return_value.status_code = 200 + monkeypatch.setattr( + "nemo_rl.models.generation.megatron.megatron_worker.requests", + requests_mock, + ) - for reserved_socket, expected_port in ((reserved, port), (None, 12345)): - worker = SimpleNamespace( - coordinator_addr="tcp://127.0.0.1:5555", - megatron_tokenizer=object(), - rank=0, - cfg={"generation": {"mcore_generation_config": {"parsers": []}}}, - _reserved_http_server_socket=reserved_socket, - inference_wrapped_model=SimpleNamespace(multimodal_prompt_config=None), - ) - base_url = MegatronGenerationMixin._setup_openai_api_server(worker) - assert started["sock"] is reserved_socket + for reserved_port, expected_port in ((port, port), (None, 12345)): + worker = SimpleNamespace( + coordinator_addr="tcp://127.0.0.1:5555", + megatron_tokenizer=object(), + rank=0, + cfg={"generation": {"mcore_generation_config": {"parsers": []}}}, + _reserved_http_server_port=reserved_port, + inference_wrapped_model=SimpleNamespace(multimodal_prompt_config=None), + ) + base_url = MegatronGenerationMixin._setup_openai_api_server(worker) + reserved_socket = started["sock"] + try: assert started["server_port"] == expected_port assert base_url == f"http://10.0.0.5:{expected_port}/v1" - finally: - reserved.close() + if reserved_port is None: + assert reserved_socket is None + continue + + # Adoption occurs only at HTTP startup, after model initialization + # can no longer leak the listener into long-lived child processes. + assert holder._sock.fileno() == -1 + assert reserved_socket.getsockname()[1] == port + + # MCore closes the handed-off fd and gives every frontend replica + # its own SO_REUSEPORT listener. Such a listener can join the reuse + # group while this test stub still holds the adopted duplicate. + replica = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + replica.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + replica.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + replica.bind(("0.0.0.0", port)) + finally: + replica.close() + finally: + if reserved_socket is not None: + reserved_socket.close() diff --git a/tests/unit/models/megatron/test_nemotron_omni_model.py b/tests/unit/models/megatron/test_nemotron_omni_model.py index adaae050e79..b89e5fd3b70 100644 --- a/tests/unit/models/megatron/test_nemotron_omni_model.py +++ b/tests/unit/models/megatron/test_nemotron_omni_model.py @@ -79,7 +79,7 @@ class _TinyOmniProvider(NemotronOmniModelProvider): use_vision_backbone_fp8_arch: bool = False vision_proj_ffn_hidden_size: int = 256 pipeline_model_parallel_size: int = 1 - use_cpu_initialization: bool = True + use_cpu_initialization: bool = False gradient_accumulation_fusion: bool = False nemotron_omni_contract: str = NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT @@ -151,6 +151,7 @@ def _build_distributed_model( use_distributed_optimizer=False, check_for_nan_in_grad=True, ), + use_cpu_initialization=False, wrap_with_ddp=True, mixed_precision_wrapper=None, ) diff --git a/tests/unit/test_effort_shaping.py b/tests/unit/test_effort_shaping.py index af14f7c0b82..3be2efe74f7 100644 --- a/tests/unit/test_effort_shaping.py +++ b/tests/unit/test_effort_shaping.py @@ -268,7 +268,12 @@ def options(self, *, num_returns): def remote(self, inputs, timer_prefix): del inputs, timer_prefix - return _Stream([_ReadyRef((i, r, None)) for i, r in enumerate(self._results)]) + return _Stream( + [ + _ReadyRef((i, {"name": "agent"}, result, None)) + for i, result in enumerate(self._results) + ] + ) def _gym_result(reward: float, response_tokens: int) -> dict: diff --git a/uv.lock b/uv.lock index 9e758eb35ce..3b5adff49d3 100644 --- a/uv.lock +++ b/uv.lock @@ -3159,7 +3159,7 @@ requires-dist = [ { name = "timm" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", specifier = ">=4.67.1" }, - { name = "transformers", specifier = ">=5.8,<=5.12.1" }, + { name = "transformers", specifier = ">=5.12.1,<=5.15.0" }, { name = "typing-extensions" }, { name = "wandb", specifier = ">=0.25.0" }, ] @@ -3303,7 +3303,7 @@ requires-dist = [ { name = "torch", specifier = ">=2.6.0" }, { name = "torch-memory-saver", marker = "extra == 'dev'", git = "https://github.com/fzyzcjy/torch_memory_saver.git?rev=9bc9a442e6d108c7b7903def199896a005143aaf" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=27486e03cfc1fa41f6932dcecdc47c71c47eac3e" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" },