ci: Bump Megatron-Bridge to 5ed9799

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

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

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

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]

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(

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."
             )

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

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"]