feat(token-id-capture): deliver rebuilt trajectories safely

nemo_gym/base_responses_api_model.py

@@ -72,6 +72,7 @@
     CaptureContext,
     capture_tokens,
     installed_token_sink,
+    register_call_intent,
     reset_token_sink,
     set_token_sink,
 )
@@ -209,6 +210,7 @@ async def _invoke_chat_completions(
         # chat_completions() signatures vary across servers: some take a leading `request`, some
         # only `body`. Dispatch on whichever this server declares so the shared dispatch works for
         # all of them.
+        await register_call_intent()
         if "request" in inspect.signature(self.chat_completions).parameters:
             completion = await self.chat_completions(request=request, body=params)
         else:
@@ -242,6 +244,7 @@ async def _invoke_responses(
         # responses() signatures vary across servers: some take a leading `request`, some only
         # `body`. Dispatch on whichever this server declares so the default messages() works for
         # all of them.
+        await register_call_intent()
         if "request" in inspect.signature(self.responses).parameters:
             response = await self.responses(request=request, body=params)
         else:
@@ -1107,6 +1110,19 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None
         # Installed sinks are resolved for each request.
         token_sink = self._configured_sink or installed_token_sink() or self._token_store
         capture_wanted = token_capture_requested and (token_sink is not None or self._token_capture_enabled)
+        if token_capture_requested and dialect is None and token_sink is not None:
+            # This call cannot produce a capture record.
+            # Its output may still feed a later prompt.
+            # Mark the rollout incomplete before forwarding the request.
+            try:
+                await token_sink.mark_incomplete(rollout_from_path, "")
+            except Exception:
+                logger.warning(
+                    "Could not mark rollout %s incomplete for unobserved path %s.",
+                    rollout_from_path,
+                    path,
+                    exc_info=True,
+                )
         if (self._store is None and not capture_wanted) or rollout_from_path is None or dialect is None:
             await self._app(scope, receive, send)
             return

nemo_gym/rollout_collection.py

@@ -85,6 +85,20 @@
     setup_server_client as setup_server_client_utils,
 )
 from nemo_gym.skills import SkillsConfig, load_skill_directory
+from nemo_gym.token_id_capture import (
+    TokenCaptureStore,
+    TokenIdCaptureConfig,
+    clear_token_captures_for_rollouts,
+    installed_token_source,
+    token_id_capture_dirs_from_config,
+)
+from nemo_gym.token_id_capture.config import token_id_capture_enabled_for_agent
+from nemo_gym.token_id_capture.delivery import (
+    MASK_SAMPLE_KEY,
+    capture_build_can_retire,
+    finalize_rollout_token_capture,
+    retire_rollout_token_capture,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -790,13 +804,53 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D
 
         # Resolve capture dirs once so each rollout's captured model calls can be folded
         # into its record below (uniform across agents; no-op when capture is off / dirs absent).
-        capture_dirs = model_call_capture_dirs_from_config(get_global_config_dict())
+        global_config = get_global_config_dict()
+        capture_dirs = model_call_capture_dirs_from_config(global_config)
+        # Resolve the training-token store directory once.
+        # Training capture is independent of evaluation capture.
+        # An empty result disables training-token capture.
+        token_capture_dirs = token_id_capture_dirs_from_config(global_config)
+        # The finalizer reads and freezes records through this source.
+        # The source is absent when capture or response rebuilding is disabled.
+        # A framework-owned transport may rebuild through its own source.
+        # The sink still records captures when Gym does not rebuild.
+        # Reruns still clear deterministic rollout ids before dispatch.
+        token_source = None
+        owned_token_source = None
+        token_capture_config = TokenIdCaptureConfig.model_validate(global_config)
+        if token_capture_config.enabled and token_capture_config.token_id_capture.rebuild_response:
+            token_source = installed_token_source()
+            if token_source is None and token_capture_dirs:
+                token_source = TokenCaptureStore(token_capture_dirs[0])
+            if isinstance(token_source, TokenCaptureStore):
+                owned_token_source = token_source
 
         # Clear only rows about to be dispatched, after resume has assigned retry suffixes. This also
         # removes a kill-shaped attempt's partial capture when its rollout-attempt id is reused.
         if capture_dirs:
             print("Clearing existing model-call captures for rollouts being dispatched")
             clear_model_call_captures_for_rollouts(input_rows, capture_dirs)
+        token_capture_rows = [
+            row
+            for row in input_rows
+            if token_id_capture_enabled_for_agent(global_config, (row.get(AGENT_REF_KEY_NAME) or {}).get("name"))
+        ]
+        if token_capture_config.token_id_capture.rebuild_response and token_capture_rows and token_source is None:
+            raise ValueError(
+                "Token capture response rebuilding requires a TokenSource in the rollout-collector process. "
+                "Call install_token_source before starting collection or configure token_id_capture.dir."
+            )
+        if token_capture_dirs and token_capture_rows:
+            # Token stores append under deterministic rollout ids.
+            # Clear stale records to avoid merging different attempts.
+            print("Clearing existing token captures for rollouts being dispatched")
+            clear_token_captures_for_rollouts(token_capture_rows, token_capture_dirs)
+
+        # Stop a run that produces mostly masked captures.
+        finalized_count = 0
+        masked_count = 0
+        mask_reasons: Counter = Counter()
+        warned_malformed_rollout_id = False
 
         # Intermediate status printing
         pcts_to_print = list(range(1, 100)) + [99.5]
@@ -834,6 +888,43 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D
             if "ng_model_call_capture" in result or "ng_agent_observations" in result or NG_TRAJECTORY_KEY in result:
                 _attach_trajectory_record(row, result)
 
+            # Freeze and rebuild tokens only for participating agents.
+            # This step does not retire the frozen snapshot.
+            # It leaves harness output and reward unchanged.
+            # Direct callers of run_examples finalize each record themselves.
+            token_capture_build = None
+            if token_id_capture_enabled_for_agent(
+                global_config,
+                (row.get(AGENT_REF_KEY_NAME) or {}).get("name"),
+            ):
+                token_capture_build = await finalize_rollout_token_capture(result, token_source)
+                if token_capture_build is not None:
+                    finalized_count += 1
+                    if token_capture_build.get(MASK_SAMPLE_KEY):
+                        masked_count += 1
+                        # Aggregate available reasons for the abort message.
+                        build_metrics = token_capture_build.get("metrics") or {}
+                        if build_metrics.get("capture_incomplete"):
+                            mask_reasons["capture_incomplete"] += 1
+                        if build_metrics.get("unresolved_parent_calls"):
+                            mask_reasons["unresolved_parent_calls"] += 1
+                        build_error = token_capture_build.get("error") or build_metrics.get("error")
+                        if build_error:
+                            mask_reasons[str(build_error)] += 1
+                    settings = token_capture_config.token_id_capture
+                    if (
+                        settings.max_mask_fraction is not None
+                        and finalized_count >= settings.mask_fraction_min_samples
+                        and masked_count / finalized_count > settings.max_mask_fraction
+                    ):
+                        raise RuntimeError(
+                            f"{masked_count}/{finalized_count} finalized rollouts "
+                            f"({masked_count / finalized_count:.1%}) are masked, exceeding "
+                            f"token_id_capture.max_mask_fraction={settings.max_mask_fraction}. "
+                            f"Mask reasons: {dict(mask_reasons)}. Aborting instead of collecting "
+                            "mostly token-less data."
+                        )
+
             no_persist = bool(result.get(NG_NO_PERSIST_KEY))
             failure_class = result.get(NG_FAILURE_CLASS_KEY)
 
@@ -856,6 +947,21 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D
                 results_file.flush()
                 persisted_rows.append(row)
                 persisted_results.append(result)
+                try:
+                    rollout_id = maybe_rollout_id_from_run_body(result)
+                except (TypeError, ValueError) as error:
+                    # Preserve capture evidence when the rollout id is invalid.
+                    rollout_id = None
+                    if not warned_malformed_rollout_id:
+                        warned_malformed_rollout_id = True
+                        warnings.warn(
+                            f"a result carries a malformed rollout id ({error}); "
+                            "its token capture will not be retired.",
+                            stacklevel=2,
+                        )
+                if rollout_id is not None and capture_build_can_retire(token_capture_build):
+                    os.fsync(results_file.fileno())
+                    await retire_rollout_token_capture(rollout_id, token_source, token_capture_build)
 
             counts_left[row[AGENT_REF_KEY_NAME]["name"]] -= 1
             if counts_left[row[AGENT_REF_KEY_NAME]["name"]] <= 0:
@@ -893,6 +999,8 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D
 
         results_file.close()
         failures_file.close()
+        if owned_token_source is not None:
+            await owned_token_source.close()
 
         if config.upload_rollouts and get_exporters():  # pragma: no cover
             print("Uploading rollouts. This may take a few minutes if your data is large.")

nemo_gym/token_id_capture/__init__.py

@@ -25,6 +25,15 @@
 Framework transports may provide their own sink and source.
 There is no HTTP token reader.
 This leaf package avoids imports from Gym's server stack.
+The rollout-record finalizer needs Gym's server stack.
+It is deliberately not re-exported here.
+Import ``nemo_gym.token_id_capture.delivery`` from server-side code.
+The incomplete state prevents training on a rollout that lost a model call.
+Finalization freezes and rebuilds the rollout.
+Finalization does not retire the snapshot.
+The caller retires it only after durable handoff.
+Retirement uses the frozen ``snapshot_id`` and version.
+Failed or masked builds retain their capture evidence.
 """
 
 from nemo_gym.token_id_capture.builder import (
@@ -63,6 +72,7 @@
     capture_tokens,
     commit_entry,
     current_capture_context,
+    register_call_intent,
     reset_token_sink,
     set_token_sink,
 )
@@ -89,6 +99,7 @@
     "CaptureContext",
     "set_token_sink",
     "reset_token_sink",
+    "register_call_intent",
     "capture_tokens",
     "commit_entry",
     "current_capture_context",

nemo_gym/token_id_capture/config.py

@@ -48,11 +48,24 @@
 They do not inherit a sink installed by a launcher.
 Configure the sink here so each worker builds its own.
 Programmatic installation must occur inside the serving process.
+
+Choosing who reads them back
+----------------------------
+``rebuild_response`` controls whether Gym rebuilds a finished rollout.
+Gym freezes captured records before rebuilding ``response.output``.
+Rebuilding does not retire the frozen snapshot.
+Gym retires a successful build only after durable handoff.
+Retirement uses the frozen ``snapshot_id`` and version.
+Failed or masked builds retain their capture evidence.
+Set it to false when a framework reads through its own ``TokenSource``.
+Gym then stops after the write.
+Read ownership is independent of write ownership.
 """
 
 from __future__ import annotations
 
 import logging
+from collections.abc import Mapping
 from importlib import import_module
 from pathlib import Path
 from typing import Any
@@ -88,8 +101,14 @@ class TokenIdCaptureSettings(BaseModel):
     # A real transport needs explicit endpoint, client, or credential wiring.
     # Use ``${oc.env:VAR}`` for secrets instead of writing them here.
     sink_kwargs: dict[str, Any] = Field(default_factory=dict)
-    # Rebuild opaque-harness responses from captured records after the run.
+    # Whether Gym freezes capture records and rebuilds the response.
+    # Finalization does not retire the frozen snapshot.
+    # Durable delivery permits retirement by snapshot id and version.
     rebuild_response: bool = True
+    # Abort once enough finalized rollouts exceed this masked fraction.
+    # ``None`` disables the limit.
+    max_mask_fraction: float | None = None
+    mask_fraction_min_samples: int = 50
 
 
 class TokenIdCaptureConfig(BaseModel):
@@ -179,3 +198,26 @@ def _build_endpoint(target: str, kwargs: dict[str, Any], protocol: type, kind: s
 def token_id_capture_config(global_config_dict: Any) -> TokenIdCaptureConfig:
     """Read the capture settings out of a global config dict."""
     return TokenIdCaptureConfig.model_validate(global_config_dict or {})
+
+
+def token_id_capture_enabled_for_agent(global_config_dict: Any, agent_name: str | None) -> bool:
+    """Return whether one configured agent participates in capture."""
+    config = token_id_capture_config(global_config_dict)
+    settings = config.token_id_capture
+    if not settings.enabled:
+        return False
+    if settings.all_agents:
+        return True
+    if not agent_name or not isinstance(global_config_dict, Mapping):
+        return False
+    server_entry = global_config_dict.get(agent_name)
+    if not isinstance(server_entry, Mapping):
+        return False
+    agents = server_entry.get("responses_api_agents")
+    if not isinstance(agents, Mapping):
+        return False
+    return any(
+        bool(agent_config.get("token_id_capture", False))
+        for agent_config in agents.values()
+        if isinstance(agent_config, Mapping)
+    )