diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index 8f2a5c0ec3..fb8618a9f2 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -41,6 +41,8 @@ Use `/ng-rollout/` as the model-server URL prefix. SDKs append their the rollout prefix before routing. Correlation is caller-supplied. An unprefixed call is forwarded normally but is not captured. +SDK-based harnesses must configure their `model_server` field for capture; calls sent directly to an +external provider do not cross a Gym model server. Agents built on `SimpleResponsesAPIAgent` can use these helpers: `url_path_for_run(url_path, body)` prefixes a downstream call from the run request's task/rollout indices (only when @@ -77,9 +79,10 @@ totals = aggregate_model_call_metrics(store, rollout_id) ``` `ModelCallRecord` is an observability serialization model derived from captured HTTP exchanges. It -contains a unique server-generated `model_call_id`, typed `model_ref`, wall-clock `started_at` and -`completed_at`, a `call_index`, API dialect, token and cache usage, latency, error details, tool -calls, reasoning content, and the captured request and response. `started_at` is recorded immediately +contains a unique server-generated `model_call_id`, the protocol `response_id` when present, typed +`model_ref`, wall-clock `started_at` and `completed_at`, a `call_index`, API dialect, token and cache +usage, latency, error details, tool calls, reasoning content, and the captured request and response. +`started_at` is recorded immediately before invoking the downstream ASGI application; `completed_at` is recorded when that invocation returns or raises, before capture parsing and persistence. Both are UTC Unix seconds for external trace correlation; durations use the monotonic latency fields. `call_index` reflects durable append @@ -111,7 +114,32 @@ data does not mix with an earlier attempt. Before dispatch, the collector clears for that exact rollout-attempt id, including a kill-shaped attempt being redispatched. The attachment is additive: it does not replace or rewrite the existing response, reward, -`NeMoGymResponse`, token-id, or log-prob fields. Downstream consumers can choose whether to read it. +`NeMoGymResponse`, token-id, or log-prob fields. Downstream consumers can choose whether to read it, +and aggregate-metrics requests exclude it. + +## Agent observations + +Supported Agent Servers may also attach `ng_agent_observations` when observability is enabled. It +contains an unordered `records` list of typed agent invocations, tool-call intervals, explicit +context-compaction events, and sandbox observations. `gaps` reports unavailable evidence. An +invocation's `conversation` contains the ordered, normalized items exposed by that integration. + +Agent observations and model-call capture are separate evidence. Join an invocation's model-call +references by `model_call_id`, or by the exact `(model_ref, response_id)` pair when the harness sees +the protocol response ID. Do not infer ownership from timestamps, text, or list position. The full +model request and response remain in `CaptureStore`; rollout attachments intentionally omit them. + +Tool timestamps are UTC Unix seconds when the source provides wall-clock time. `duration_ms` is the +measured interval, and `timing_source` identifies executor, harness, or artifact-derived timing. + +Model-visible tool calls and results remain in `AgentInvocation.conversation` as `function_call` and +`function_call_output` items. Execution timing and outcome, when observable, live in +`ToolCallObservation` and join through `(invocation_id, tool_call_id)`. A tool observation may also +reference its enclosing `sandbox_id`; concurrent calls retain independent timing and outcome. + +Each `SandboxObservation` covers one sandbox execution. Usage fields contain measured values only; +configured limits are never reported as usage. Integrations emit only facts available at their +execution boundary or in retained artifacts and report unavailable evidence in `gaps`. ## Limitations @@ -119,3 +147,7 @@ The model-server boundary observes model HTTP requests and responses. It can rec reasoning, and tool results present in those payloads, but it does not observe the actual tool execution boundary, environment events, context compaction, semantic turns, or subagent structure. Those require instrumentation at the agent, tool, environment, or rollout layer. + +Sandbox CPU and memory observations describe the sandbox as a whole. They cannot be attributed to +individual overlapping tool calls unless each call runs in a separately measured execution scope. +Gym therefore records the shared sandbox relationship without estimating per-tool resource usage. diff --git a/nemo_gym/base_resources_server.py b/nemo_gym/base_resources_server.py index 1c4b6be35c..2eba9f0605 100644 --- a/nemo_gym/base_resources_server.py +++ b/nemo_gym/base_resources_server.py @@ -30,6 +30,7 @@ NeMoGymResponseCreateParamsNonStreaming, ) from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics +from nemo_gym.rollout_correlation import RolloutContextMiddleware from nemo_gym.server_utils import BaseRunServerInstanceConfig, BaseServer, SimpleServer @@ -110,6 +111,7 @@ def setup_webserver(self) -> FastAPI: app = FastAPI() self.setup_session_middleware(app) + app.add_middleware(RolloutContextMiddleware) app.post("/seed_session")(self.seed_session) app.post("/verify")(self.verify) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 18fe0816ab..c6e76cd89f 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -14,6 +14,7 @@ # limitations under the License. from abc import abstractmethod from collections.abc import Mapping +from functools import wraps from typing import Any, Optional from fastapi import Body, FastAPI, Request @@ -24,7 +25,6 @@ BaseRunRequest, BaseVerifyResponse, ) -from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body from nemo_gym.config_types import ROLLOUT_PATH_PREFIX from nemo_gym.global_config import OBSERVABILITY_ENABLED_KEY_NAME, get_first_server_config_dict from nemo_gym.openai_utils import ( @@ -32,6 +32,7 @@ NeMoGymResponseCreateParamsNonStreaming, ) from nemo_gym.reward_profile import AggregateMetricsMixin, compute_aggregate_metrics +from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body, rollout_context from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, BaseServer, @@ -62,7 +63,18 @@ def setup_webserver(self) -> FastAPI: # responses() recovers the rollout id from the path (see url_path_for_request) to correlate # its model calls. Same handler, so unprefixed calls are unaffected. app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses")(self.responses) - app.post("/run")(self.run) + + run = self.run + + @wraps(run) + async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyResponse: + body = kwargs.get("body") + if body is None: + body = next((arg for arg in args if isinstance(arg, BaseRunRequest)), None) + with rollout_context(self.rollout_id_from_run(body)): + return await run(*args, **kwargs) + + app.post("/run")(run_with_rollout_context) app.post("/aggregate_metrics")(self.aggregate_metrics) return app diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 7d29549ad9..377de77daa 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -33,7 +33,6 @@ import logging import os import re -import threading import time from abc import abstractmethod from pathlib import Path @@ -48,11 +47,6 @@ from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, ModelServerRef -from nemo_gym.global_config import ( - ATTEMPT_INDEX_KEY_NAME, - ROLLOUT_INDEX_KEY_NAME, - TASK_INDEX_KEY_NAME, -) from nemo_gym.openai_utils import ( NeMoGymChatCompletion, NeMoGymChatCompletionCreateParamsNonStreaming, @@ -65,6 +59,8 @@ synthesize_responses_sse, validate_streaming_responses_params, ) +from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body +from nemo_gym.rollout_observability import AgentObservationBundle, ObservationGap, join_model_call_observations from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, BaseServer, @@ -228,7 +224,6 @@ class CaptureStore: def __init__(self, root: str | Path) -> None: self._root = Path(root) self._root.mkdir(parents=True, exist_ok=True) - self._lock = threading.Lock() @property def root(self) -> Path: @@ -237,25 +232,32 @@ def root(self) -> Path: def path_for(self, rollout_id: str) -> Path: return self._root / f"{_validate_rollout_id(rollout_id)}.capture.jsonl" + def incomplete_path_for(self, rollout_id: str) -> Path: + return self._root / f"{_validate_rollout_id(rollout_id)}.capture.incomplete" + + def mark_incomplete(self, rollout_id: str) -> None: + self.incomplete_path_for(rollout_id).touch(exist_ok=True) + + def is_incomplete(self, rollout_id: str) -> bool: + return self.incomplete_path_for(rollout_id).exists() + def record(self, rollout_id: str, exchange: dict[str, Any]) -> None: """Append one exchange and fsync (durable across a killed box). - ``flock`` serializes appends across worker processes (a model server may run with - ``num_workers > 1``, where the in-process lock can't coordinate); the in-process lock - serializes threads. This does blocking file IO + fsync, so callers run it off the event - loop (the capture middleware offloads it via ``asyncio.to_thread``). + ``flock`` serializes appends to the same rollout across worker processes and threads while + allowing independent rollouts to write concurrently. This does blocking file IO + fsync, + so callers run it off the event loop (the capture middleware uses ``asyncio.to_thread``). """ line = orjson.dumps(exchange, default=str, option=orjson.OPT_APPEND_NEWLINE) path = self.path_for(rollout_id) - with self._lock: - with path.open("ab") as handle: - fcntl.flock(handle.fileno(), fcntl.LOCK_EX) - try: - handle.write(line) - handle.flush() - os.fsync(handle.fileno()) - finally: - fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + with path.open("ab") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + try: + handle.write(line) + handle.flush() + os.fsync(handle.fileno()) + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) def read(self, rollout_id: str) -> list[dict[str, Any]]: path = self.path_for(rollout_id) @@ -263,49 +265,53 @@ def read(self, rollout_id: str) -> list[dict[str, Any]]: return [] exchanges: list[dict[str, Any]] = [] # Stream line-by-line; a capture can be large (token-ids / logprobs). - with self._lock: - with path.open("rb") as handle: - fcntl.flock(handle.fileno(), fcntl.LOCK_SH) - try: - for line in handle: - stripped = line.strip() - if not stripped: - continue - exchanges.append(orjson.loads(stripped)) - finally: - fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + with path.open("rb") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_SH) + try: + for line in handle: + stripped = line.strip() + if not stripped: + continue + exchanges.append(orjson.loads(stripped)) + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) return exchanges + def read_available(self, rollout_id: str) -> tuple[list[tuple[int, dict[str, Any]]], int]: + """Read valid exchanges without letting one damaged line hide the rest.""" + path = self.path_for(rollout_id) + if not path.exists(): + return [], 0 + exchanges: list[tuple[int, dict[str, Any]]] = [] + invalid_count = 0 + capture_index = 0 + with path.open("rb") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_SH) + try: + for line in handle: + stripped = line.strip() + if not stripped: + continue + try: + exchange = orjson.loads(stripped) + except orjson.JSONDecodeError: + invalid_count += 1 + else: + if isinstance(exchange, dict): + exchanges.append((capture_index, exchange)) + else: + invalid_count += 1 + capture_index += 1 + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + return exchanges, invalid_count -def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: - """Per-rollout model-call capture id from a run-request's task/rollout indices. - Reads the canonical row keys (``_ng_task_index`` / ``_ng_rollout_index``) that - rollout_collection ships to an agent's ``/run``. When a resume re-dispatch attempt is present - (``_ng_attempt_index`` > 0), an ``-a`` suffix is appended so a retry's captured model calls - stay separable from the prior attempt; the first attempt (0) keeps the bare ``-`` - key for backward compatibility. - """ - if isinstance(body, BaseModel): - data = body.model_dump() - elif isinstance(body, Mapping): - data = body - else: - return None - task = data.get(TASK_INDEX_KEY_NAME) - rollout = data.get(ROLLOUT_INDEX_KEY_NAME) - if task is None or rollout is None: - return None - rollout_id = f"{task}-{rollout}" - attempt = data.get(ATTEMPT_INDEX_KEY_NAME) - if attempt is not None: - attempt_index = int(attempt) - if attempt_index > 0: - rollout_id = f"{rollout_id}-a{attempt_index}" - return rollout_id +# --- Observability records derived from captured exchanges --- -# --- Observability records derived from captured exchanges --- +def _token_count(value: Any) -> Optional[int]: + return value if type(value) is int and value >= 0 else None def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: @@ -321,7 +327,7 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: cache-creation (~1.25x) differently from base input, so cost-accurate consumers should weight ``cached_tokens`` and ``cache_creation_tokens`` separately rather than summing ``tokens_in``. """ - if not usage: + if not isinstance(usage, Mapping): return { "tokens_in": None, "tokens_out": None, @@ -329,28 +335,30 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: "tokens_total": None, "cache_creation_tokens": None, } - tokens_in = usage.get("input_tokens") + tokens_in = _token_count(usage.get("input_tokens")) if tokens_in is None: - tokens_in = usage.get("prompt_tokens") - tokens_out = usage.get("output_tokens") + tokens_in = _token_count(usage.get("prompt_tokens")) + tokens_out = _token_count(usage.get("output_tokens")) if tokens_out is None: - tokens_out = usage.get("completion_tokens") + tokens_out = _token_count(usage.get("completion_tokens")) # Anthropic-native shape: top-level cache_* keys mean input_tokens excludes cached tokens. - cache_read = usage.get("cache_read_input_tokens") - cache_creation = usage.get("cache_creation_input_tokens") + cache_read = _token_count(usage.get("cache_read_input_tokens")) + cache_creation = _token_count(usage.get("cache_creation_input_tokens")) if cache_read is not None or cache_creation is not None: # A fully-cached response can omit input_tokens; use a 0 base so the folded prompt size is # preserved rather than dropped to null. (Top-level cache_* keys are Anthropic-only, so the # OpenAI/Responses path -- nested prompt_tokens_details.cached_tokens -- never enters here.) tokens_in = (tokens_in or 0) + (cache_read or 0) + (cache_creation or 0) - tokens_total = usage.get("total_tokens") + tokens_total = _token_count(usage.get("total_tokens")) if tokens_total is None and tokens_in is not None and tokens_out is not None: tokens_total = tokens_in + tokens_out details = usage.get("output_tokens_details") or usage.get("completion_tokens_details") or {} + if not isinstance(details, Mapping): + details = {} return { "tokens_in": tokens_in, "tokens_out": tokens_out, - "tokens_reasoning": details.get("reasoning_tokens"), + "tokens_reasoning": _token_count(details.get("reasoning_tokens")), "tokens_total": tokens_total, "cache_creation_tokens": cache_creation, } @@ -358,12 +366,14 @@ def extract_token_stats(usage: Any) -> dict[str, Optional[int]]: def _cache_signal(usage: Any) -> tuple[Optional[bool], Optional[int]]: """Cache hit/miss + cached-token count, from usage cache fields (OpenAI / Anthropic).""" - if not usage: + if not isinstance(usage, Mapping): return None, None details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") or {} - cached = details.get("cached_tokens") + if not isinstance(details, Mapping): + details = {} + cached = _token_count(details.get("cached_tokens")) if cached is None: - cached = usage.get("cache_read_input_tokens") # Anthropic + cached = _token_count(usage.get("cache_read_input_tokens")) # Anthropic if cached is None: return None, None return cached > 0, cached @@ -401,7 +411,7 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str, elif item.get("type") == "reasoning": for summary in item.get("summary") or []: text = summary.get("text") if isinstance(summary, dict) else None - if text: + if isinstance(text, str) and text: reasoning.append(text) return tool_calls, ("\n".join(reasoning) or None) @@ -409,19 +419,21 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str, if isinstance(choices, list): # Chat Completions for choice in choices: message = choice.get("message") if isinstance(choice, dict) else None - if not message: + if not isinstance(message, Mapping): continue for tc in message.get("tool_calls") or []: if not isinstance(tc, dict): continue fn = tc.get("function") or {} + if not isinstance(fn, Mapping): + fn = {} tool_calls.append( {"call_id": tc.get("id"), "name": fn.get("name"), "arguments": _as_arguments(fn.get("arguments"))} ) # vLLM and newer OpenAI-compatible servers emit `reasoning`; `reasoning_content` is the # older field. Accept either (reasoning_content wins when both are present). reasoning_text = message.get("reasoning_content") or message.get("reasoning") - if reasoning_text: + if isinstance(reasoning_text, str) and reasoning_text: reasoning.append(reasoning_text) return tool_calls, ("\n".join(reasoning) or None) @@ -434,7 +446,7 @@ def _tool_calls_and_reasoning(response: dict[str, Any]) -> tuple[list[dict[str, tool_calls.append( {"call_id": block.get("id"), "name": block.get("name"), "arguments": block.get("input") or {}} ) - elif block.get("type") in ("thinking", "redacted_thinking") and block.get("thinking"): + elif block.get("type") in ("thinking", "redacted_thinking") and isinstance(block.get("thinking"), str): reasoning.append(block["thinking"]) return tool_calls, ("\n".join(reasoning) or None) @@ -446,12 +458,15 @@ class ModelCallRecord(BaseModel): # Unique server-generated identity for each persisted call. model_call_id: Optional[str] = None + response_id: Optional[str] = None # Durable append order, not a causal or semantic order for concurrent calls. call_index: int model_ref: Optional[ModelServerRef] = None + model: Optional[str] = None dialect: Optional[str] = None status_code: Optional[int] = None + finish_reason: Optional[str] = None # Wall-clock bounds around the downstream ASGI invocation, as UTC Unix timestamps. These are # for external trace correlation; durations use the monotonic latency fields below. @@ -469,6 +484,8 @@ class ModelCallRecord(BaseModel): # Model-call record. request: Optional[dict[str, Any]] = None response: Optional[dict[str, Any]] = None + request_raw: Optional[str] = None + response_raw: Optional[str] = None tool_calls: list[dict[str, Any]] = Field(default_factory=list) # Structured reasoning (not flattened into the response text). @@ -490,20 +507,46 @@ class ModelCallRecord(BaseModel): def build_model_call_record(exchange: dict[str, Any], *, call_index: int) -> ModelCallRecord: """Map one captured exchange and its transport metadata into an observability record.""" - response = exchange.get("response") or {} + raw_response = exchange.get("response") + response = raw_response if isinstance(raw_response, dict) else {} tokens = extract_token_stats(response.get("usage")) cache_hit, cached_tokens = _cache_signal(response.get("usage")) tool_calls, reasoning_content = _tool_calls_and_reasoning(response) + raw_request = exchange.get("request") + request = raw_request if isinstance(raw_request, dict) else {} + choices = response.get("choices") + first_choice = choices[0] if isinstance(choices, list) and choices and isinstance(choices[0], dict) else {} + incomplete_details = response.get("incomplete_details") + if not isinstance(incomplete_details, dict): + incomplete_details = {} + finish_reason = next( + ( + value + for value in ( + response.get("stop_reason"), + first_choice.get("finish_reason"), + incomplete_details.get("reason"), + ) + if isinstance(value, str) + ), + None, + ) + model = response.get("model") or request.get("model") return ModelCallRecord( model_call_id=exchange.get("model_call_id"), + response_id=response.get("id") if isinstance(response.get("id"), str) else None, call_index=call_index, model_ref=exchange.get("model_ref"), + model=model if isinstance(model, str) else None, dialect=exchange.get("dialect"), status_code=exchange.get("status_code"), + finish_reason=finish_reason, started_at=exchange.get("started_at"), completed_at=exchange.get("completed_at"), - request=exchange.get("request"), - response=response or None, + request=raw_request if isinstance(raw_request, dict) else None, + response=raw_response if isinstance(raw_response, dict) else None, + request_raw=exchange.get("request_raw") if isinstance(exchange.get("request_raw"), str) else None, + response_raw=exchange.get("response_raw") if isinstance(exchange.get("response_raw"), str) else None, tool_calls=tool_calls, reasoning_content=reasoning_content, cache_hit=cache_hit, @@ -522,6 +565,18 @@ def read_model_call_records(store: CaptureStore, rollout_id: str) -> list[ModelC ] +def read_available_model_call_records(store: CaptureStore, rollout_id: str) -> tuple[list[ModelCallRecord], int]: + """Read valid call records and count damaged records.""" + exchanges, invalid_count = store.read_available(rollout_id) + calls = [] + for index, exchange in exchanges: + try: + calls.append(build_model_call_record(exchange, call_index=index)) + except Exception: + invalid_count += 1 + return calls, invalid_count + + def aggregate_model_call_records(calls: list[ModelCallRecord]) -> dict[str, Any]: """Aggregate token and latency values from model-call records.""" @@ -643,6 +698,24 @@ def _classify_exception(exc: BaseException) -> str: return "exception" +def _exception_http_details(exc: BaseException) -> tuple[Optional[int], bytes]: + response = getattr(exc, "response", None) + status = getattr(exc, "status", None) + if not isinstance(status, int): + status = getattr(exc, "status_code", None) + if not isinstance(status, int) and response is not None: + status = getattr(response, "status_code", None) + + body = getattr(exc, "response_content", None) + if body is None and response is not None: + body = getattr(response, "content", None) + if body is None: + body = getattr(response, "text", None) + if isinstance(body, str): + body = body.encode() + return (status if isinstance(status, int) else None, bytes(body) if isinstance(body, (bytes, bytearray)) else b"") + + # --- SSE reconstruction: rebuild a final response object from a streamed body --- def _parse_sse_events(raw: bytes) -> list[dict[str, Any]]: """Parse an SSE byte stream into its JSON ``data:`` payloads (best-effort; non-JSON skipped).""" @@ -719,11 +792,14 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An tool_calls: dict[int, dict[str, Any]] = {} usage: Optional[dict[str, Any]] = None model: Optional[str] = None + response_id: Optional[str] = None role = "assistant" finish_reason: Optional[str] = None saw_choice = False for chunk in events: model = chunk.get("model") or model + if isinstance(chunk.get("id"), str): + response_id = chunk["id"] if chunk.get("usage"): usage = chunk["usage"] for choice in chunk.get("choices") or []: @@ -762,6 +838,8 @@ def _reconstruct_chat_sse(events: list[dict[str, Any]]) -> Optional[dict[str, An "model": model, "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], } + if response_id is not None: + result["id"] = response_id if usage: result["usage"] = usage return result @@ -807,6 +885,7 @@ def _record( error_category: Optional[str], latency_ms: float, ttft_ms: Optional[float] = None, + response_raw: Optional[str] = None, ) -> None: """Append one exchange (success or failure). Best-effort: never raises.""" request_body = None @@ -837,9 +916,15 @@ def _record( } if request_raw is not None: exchange["request_raw"] = request_raw + if response_raw is not None: + exchange["response_raw"] = response_raw store.record(rollout_id, exchange) except Exception: logger.warning("Model-call capture failed for one %s call.", dialect, exc_info=True) + try: + store.mark_incomplete(rollout_id) + except Exception: + logger.warning("Could not mark rollout %s capture as incomplete.", rollout_id, exc_info=True) class _CaptureMiddleware: @@ -943,6 +1028,11 @@ async def _flush_deferred_response() -> None: await self._app(scope, _receive, _send) except Exception as exc: completed_at = time.time() + exception_status, exception_body = _exception_http_details(exc) + upstream_status = state["status"] or exception_status + upstream_body = bytes(state["body"]) or exception_body + error_category = _classify_status(upstream_status) if isinstance(upstream_status, int) else None + error_category = error_category or _classify_exception(exc) # Offload the blocking write+fsync so it never stalls the event loop. try: await asyncio.to_thread( @@ -956,10 +1046,11 @@ async def _flush_deferred_response() -> None: started_at=started_at, completed_at=completed_at, response_body=None, - status_code=None, - error_category=_classify_exception(exc), + status_code=upstream_status, + error_category=error_category, latency_ms=(time.perf_counter() - start) * 1000.0, ttft_ms=state["ttft_ms"], + response_raw=upstream_body.decode("utf-8", errors="replace") if upstream_body else None, ) except Exception: logger.warning("Model-call capture finalization failed.", exc_info=True) @@ -986,6 +1077,8 @@ def _parse_and_record() -> None: response_body = ( _reconstruct_streamed_response(body_bytes, dialect) if streaming else json.loads(body_bytes) ) + if not isinstance(response_body, dict): + response_body = None except Exception: response_body = None error_category = _classify_status(status) if status is not None else None @@ -1000,6 +1093,11 @@ def _parse_and_record() -> None: # doesn't silently count as a success with null tokens in reliability/cost sums. if error_category is None and body_bytes and response_body is None: error_category = "capture_parse_error" + response_raw = ( + body_bytes.decode("utf-8", errors="replace") + if body_bytes and (streaming or response_body is None) + else None + ) _record( store, dialect, @@ -1014,6 +1112,7 @@ def _parse_and_record() -> None: error_category=error_category, latency_ms=latency_ms, ttft_ms=ttft_ms, + response_raw=response_raw, ) try: @@ -1058,7 +1157,7 @@ def model_call_capture_dirs_from_config(global_config_dict: Any) -> list[Path]: def _store_for_rollout(rollout_id: str, capture_dirs: list[Path]) -> Optional[CaptureStore]: for directory in capture_dirs: store = CaptureStore(directory) - if store.path_for(rollout_id).exists(): + if store.path_for(rollout_id).exists() or store.is_incomplete(rollout_id): return store return None @@ -1078,34 +1177,60 @@ def clear_model_call_captures_for_rollouts(records: list[Any], capture_dirs: lis rollout_id = maybe_rollout_id_from_run_body(record) if rollout_id: store.path_for(rollout_id).unlink(missing_ok=True) + store.incomplete_path_for(rollout_id).unlink(missing_ok=True) -def merge_model_call_capture_into_record( - record: dict[str, Any], capture_dirs: list[Path], *, include_payloads: bool = False -) -> dict[str, Any]: +def merge_model_call_capture_into_record(record: dict[str, Any], capture_dirs: list[Path]) -> dict[str, Any]: """Attach captured model-call observability data to a rollout record in place. Keyed by the rollout id derived from the record's task/rollout/attempt indices, so the attached shape is identical for every agent harness. Adds ``ng_model_call_capture = {rollout_id, metrics, calls}`` where ``calls`` are derived observability - records. Raw request and response payloads remain in the capture store unless ``include_payloads`` - is true. No-op when no capture exists. The harness output and reward are not modified. + records. Raw request and response payloads remain in the capture store. Capture/read/join + failures are attached as ``gaps``. The harness output and reward are not modified. """ if not capture_dirs: return record rollout_id = maybe_rollout_id_from_run_body(record) if rollout_id is None: return record + gaps: list[ObservationGap] = [] store = _store_for_rollout(rollout_id, capture_dirs) if store is None: - return record - calls = read_model_call_records(store, rollout_id) - if not calls: - return record - exclude = None if include_payloads else {"request", "response"} - record["ng_model_call_capture"] = { + calls = [] + gaps.append(ObservationGap(code="model_call_capture_no_records")) + else: + try: + calls, invalid_count = read_available_model_call_records(store, rollout_id) + if invalid_count: + gaps.append( + ObservationGap( + code="model_call_capture_records_unreadable", + detail=f"count={invalid_count}", + ) + ) + if store.is_incomplete(rollout_id): + gaps.append(ObservationGap(code="model_call_capture_incomplete")) + elif not calls and not invalid_count: + gaps.append(ObservationGap(code="model_call_capture_no_records")) + except Exception: + logger.warning("Could not read model-call capture for rollout %s.", rollout_id, exc_info=True) + calls = [] + gaps.append(ObservationGap(code="model_call_capture_unreadable")) + observations = record.get("ng_agent_observations") + if observations is not None and calls: + try: + bundle = AgentObservationBundle.model_validate(observations) + record["ng_agent_observations"] = join_model_call_observations(bundle, calls).model_dump(mode="json") + except Exception: + logger.warning("Could not join agent observations for rollout %s.", rollout_id, exc_info=True) + gaps.append(ObservationGap(code="agent_observation_join_failed")) + capture = { "rollout_id": rollout_id, "metrics": aggregate_model_call_records(calls), - "calls": [call.model_dump(exclude=exclude) for call in calls], + "calls": [call.model_dump(exclude={"request", "response", "request_raw", "response_raw"}) for call in calls], } + if gaps: + capture["gaps"] = [gap.model_dump(mode="json", exclude_none=True) for gap in gaps] + record["ng_model_call_capture"] = capture return record diff --git a/nemo_gym/mcp_auto_exposure.py b/nemo_gym/mcp_auto_exposure.py index 45de0aaaef..0e6aab4f6d 100644 --- a/nemo_gym/mcp_auto_exposure.py +++ b/nemo_gym/mcp_auto_exposure.py @@ -438,6 +438,8 @@ def harvest_tools(app: FastAPI, server: Any) -> dict[str, MCPTool]: cls = m.cls if f"{cls.__module__}.{cls.__name__}" == "starlette.middleware.sessions.SessionMiddleware": continue # Gym's SessionMiddleware — replaced by a materialized session on direct dispatch + if f"{cls.__module__}.{cls.__name__}" == "nemo_gym.rollout_correlation.RolloutContextMiddleware": + continue # Correlation prefixes are handled before resource routes; direct MCP dispatch has no prefix. dispatch = m.kwargs.get("dispatch") if dispatch is not None and getattr(dispatch, "__module__", None) in _GYM_MIDDLEWARE_MODULES: continue # Gym's add_session_id / exception middleware diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index 2fb0bd15b1..d503687980 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -653,7 +653,17 @@ async def _fetch_agent_metrics(agent_name: str, agent_result_list: List[Dict]) - # Strip heavyweight fields before sending, but preserve response.usage stripped = [] for r in agent_result_list: - entry = {k: v for k, v in r.items() if k not in ("response", "responses_create_params")} + entry = { + k: v + for k, v in r.items() + if k + not in ( + "response", + "responses_create_params", + "ng_agent_observations", + "ng_model_call_capture", + ) + } usage = (r.get("response") or {}).get("usage") if usage: entry["response"] = {"usage": usage} diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py new file mode 100644 index 0000000000..c1d58a2694 --- /dev/null +++ b/nemo_gym/rollout_correlation.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 re +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Optional + +from pydantic import BaseModel + +from nemo_gym.config_types import ROLLOUT_PATH_PREFIX +from nemo_gym.global_config import ( + ATTEMPT_INDEX_KEY_NAME, + ROLLOUT_INDEX_KEY_NAME, + TASK_INDEX_KEY_NAME, +) + + +_ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_rollout_id", default=None) + + +def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: + """Build the capture key stamped by rollout collection.""" + if not isinstance(body, (BaseModel, Mapping)): + return None + + def field(key: str) -> Any: + return body.get(key) if isinstance(body, Mapping) else getattr(body, key, None) + + task = field(TASK_INDEX_KEY_NAME) + rollout = field(ROLLOUT_INDEX_KEY_NAME) + if task is None or rollout is None: + return None + + rollout_id = f"{task}-{rollout}" + attempt = field(ATTEMPT_INDEX_KEY_NAME) + if attempt is not None and int(attempt) > 0: + rollout_id = f"{rollout_id}-a{int(attempt)}" + return rollout_id + + +def current_rollout_id() -> Optional[str]: + return _ROLLOUT_ID.get() + + +@contextmanager +def rollout_context(rollout_id: Optional[str]) -> Iterator[None]: + token = _ROLLOUT_ID.set(rollout_id) + try: + yield + finally: + _ROLLOUT_ID.reset(token) + + +class RolloutContextMiddleware: + """Strip a rollout prefix and expose it to downstream Gym calls for this request.""" + + _PREFIX = re.compile( + rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P[A-Za-z0-9][A-Za-z0-9._-]*)(?P/.*)$" + ) + + def __init__(self, app: Any) -> None: + self._app = app + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + match = self._PREFIX.match(scope.get("path", "")) if scope.get("type") == "http" else None + if match is None: + await self._app(scope, receive, send) + return + + path = match.group("rest") + scope = {**scope, "path": path, "raw_path": path.encode()} + with rollout_context(match.group("rollout_id")): + await self._app(scope, receive, send) diff --git a/nemo_gym/rollout_observability.py b/nemo_gym/rollout_observability.py new file mode 100644 index 0000000000..162359e059 --- /dev/null +++ b/nemo_gym/rollout_observability.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small shared contract for observations exposed by Agent integrations.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Annotated, Literal, Optional + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from nemo_gym.config_types import ModelServerRef +from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseInputItem + + +if TYPE_CHECKING: + from nemo_gym.base_responses_api_model import ModelCallRecord + + +class ObservationModel(BaseModel): + model_config = ConfigDict(extra="forbid", validate_assignment=True) + + +class ModelCallRef(ObservationModel): + """Stable identifiers an Agent integration can observe for one model call.""" + + model_call_id: Optional[str] = None + model_ref: Optional[ModelServerRef] = None + response_id: Optional[str] = None + + @model_validator(mode="after") + def validate_join_key(self) -> "ModelCallRef": + if not self.model_call_id and not (self.model_ref is not None and self.response_id): + raise ValueError("model_call_id or both model_ref and response_id are required") + return self + + +class AgentInvocation(ObservationModel): + """One root Agent or subagent conversation observed by a harness.""" + + kind: Literal["agent_invocation"] = "agent_invocation" + invocation_id: str + parent_invocation_id: Optional[str] = None + spawned_by_tool_call_id: Optional[str] = None + status: Literal["completed", "failed", "incomplete", "unknown"] = Field( + default="unknown", description="Harness-reported invocation outcome; unknown when not explicit." + ) + duration_ms: Optional[float] = Field(default=None, ge=0) + error_type: Optional[str] = None + model_calls: list[ModelCallRef] = Field(default_factory=list) + conversation: list[NeMoGymResponseInputItem] = Field( + default_factory=list, + description="Normalized conversation items supported by this producer; gaps describe unavailable evidence.", + ) + + +class ToolCallObservation(ObservationModel): + """Timing observed for one tool call at an Agent-owned boundary.""" + + kind: Literal["tool_call"] = "tool_call" + invocation_id: str + tool_call_id: str + sandbox_id: Optional[str] = Field( + default=None, + description="Enclosing sandbox instance, shared by concurrent calls; not per-call resource attribution.", + ) + tool_name: Optional[str] = None + started_at: Optional[float] = None + completed_at: Optional[float] = None + duration_ms: Optional[float] = Field(default=None, ge=0) + timing_source: Optional[Literal["executor", "artifact", "harness"]] = None + status: Literal["completed", "failed", "timeout", "cancelled", "incomplete", "unknown"] = "unknown" + error_type: Optional[str] = None + + @model_validator(mode="after") + def validate_timing(self) -> "ToolCallObservation": + if self.started_at is not None and self.completed_at is not None and self.completed_at < self.started_at: + raise ValueError("completed_at must not precede started_at") + return self + + +class SandboxObservation(ObservationModel): + """Outcome and lifetime resource usage reported by a sandbox-owning harness.""" + + kind: Literal["sandbox"] = "sandbox" + role: Literal["agent", "verifier", "environment"] + provider: Optional[str] = None + sandbox_id: Optional[str] = None + outcome: Literal["completed", "failed", "timeout", "oom", "sandbox_error", "cancelled", "unknown"] = "unknown" + exit_code: Optional[int] = None + wall_time_s: Optional[float] = Field(default=None, ge=0) + cpu_time_s: Optional[float] = Field( + default=None, + ge=0, + description="Cumulative CPU time for the sandbox, never an allocation or per-tool estimate.", + ) + peak_memory_mib: Optional[float] = Field( + default=None, + ge=0, + description="Measured sandbox high-water mark, never its configured memory limit.", + ) + resource_usage_source: Optional[str] = None + error_type: Optional[str] = None + + +class ContextCompactionObservation(ObservationModel): + """An explicit context-compaction event reported by the Agent harness.""" + + kind: Literal["context_compaction"] = "context_compaction" + invocation_id: str + observed_at: Optional[float] = None + trigger: Optional[str] = None + tokens_before: Optional[int] = None + tokens_after: Optional[int] = None + outcome: Literal["completed", "failed", "aborted", "unknown"] = "unknown" + summary: Optional[str] = None + first_kept_item_id: Optional[str] = None + before_model_call: Optional[ModelCallRef] = None + after_model_call: Optional[ModelCallRef] = None + + +class ObservationGap(ObservationModel): + """A fact that the selected integration could not observe or join exactly.""" + + code: str + invocation_id: Optional[str] = None + detail: Optional[str] = None + + +AgentObservationRecord = Annotated[ + AgentInvocation | ToolCallObservation | ContextCompactionObservation | SandboxObservation, + Field(discriminator="kind"), +] + + +class AgentObservationBundle(ObservationModel): + """Normalized observations returned by one Agent Server for one rollout.""" + + source: str + records: list[AgentObservationRecord] = Field( + default_factory=list, + description="Unordered typed records; list position does not imply execution order.", + ) + gaps: list[ObservationGap] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_identity(self) -> "AgentObservationBundle": + invocation_ids = [record.invocation_id for record in self.records if isinstance(record, AgentInvocation)] + if len(invocation_ids) != len(set(invocation_ids)): + raise ValueError("invocation_id must be unique within an observation bundle") + return self + + +def join_model_call_observations( + bundle: AgentObservationBundle, + calls: Iterable[ModelCallRecord], +) -> AgentObservationBundle: + """Resolve harness call references against captured model calls without guessing ownership.""" + + result = bundle.model_copy() + result.records = [ + record.model_copy() if isinstance(record, (AgentInvocation, ContextCompactionObservation)) else record + for record in bundle.records + ] + invocations = [record for record in result.records if isinstance(record, AgentInvocation)] + compactions = [record for record in result.records if isinstance(record, ContextCompactionObservation)] + captured = list(calls) + by_call_id: dict[str, list[ModelCallRecord]] = {} + by_response: dict[tuple[str, str, str], list[ModelCallRecord]] = {} + for call in captured: + if call.model_call_id: + by_call_id.setdefault(call.model_call_id, []).append(call) + if call.model_ref is not None and call.response_id: + key = (call.model_ref.type, call.model_ref.name, call.response_id) + by_response.setdefault(key, []).append(call) + + def matches(ref: ModelCallRef) -> list[ModelCallRecord]: + if ref.model_call_id: + candidates = by_call_id.get(ref.model_call_id, []) + return [ + call + for call in candidates + if (ref.model_ref is None or ref.model_ref == call.model_ref) + and (ref.response_id is None or ref.response_id == call.response_id) + ] + assert ref.model_ref is not None and ref.response_id is not None + return by_response.get((ref.model_ref.type, ref.model_ref.name, ref.response_id), []) + + def canonical(ref: ModelCallRef, call: ModelCallRecord) -> ModelCallRef: + return ModelCallRef.model_validate( + { + "model_call_id": call.model_call_id, + "model_ref": call.model_ref, + "response_id": call.response_id, + } + ) + + join_codes = { + "model_call_reference_ambiguous", + "model_call_reference_conflict", + "model_call_reference_unmatched", + } + result.gaps = [ + gap + for gap in bundle.gaps + if gap.code not in join_codes + and not ( + gap.code == "model_call_ownership_unavailable" + and gap.invocation_id is None + and (gap.detail is None or gap.detail.startswith("capture:")) + ) + ] + + claimed: set[int] = set() + join_gaps: list[ObservationGap] = [] + for invocation in invocations: + resolved: list[ModelCallRef] = [] + for ref in invocation.model_calls: + candidates = matches(ref) + detail = ref.model_call_id or ref.response_id + if len(candidates) != 1: + join_gaps.append( + ObservationGap( + code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"), + invocation_id=invocation.invocation_id, + detail=detail, + ) + ) + resolved.append(ref) + continue + + call = candidates[0] + identity = id(call) + if identity in claimed: + join_gaps.append( + ObservationGap( + code="model_call_reference_conflict", + invocation_id=invocation.invocation_id, + detail=call.model_call_id or call.response_id, + ) + ) + resolved.append(ref) + continue + claimed.add(identity) + resolved.append(canonical(ref, call)) + invocation.model_calls = resolved + + for compaction in compactions: + for field_name in ("before_model_call", "after_model_call"): + ref = getattr(compaction, field_name) + if ref is None: + continue + candidates = matches(ref) + if len(candidates) == 1: + setattr(compaction, field_name, canonical(ref, candidates[0])) + else: + join_gaps.append( + ObservationGap( + code=("model_call_reference_ambiguous" if candidates else "model_call_reference_unmatched"), + invocation_id=compaction.invocation_id, + detail=f"{field_name}:{ref.model_call_id or ref.response_id}", + ) + ) + + result.gaps.extend(join_gaps) + for call in captured: + if id(call) not in claimed: + result.gaps.append( + ObservationGap( + code="model_call_ownership_unavailable", + detail=f"capture:{call.model_call_id or call.response_id or 'unknown'}:call_index={call.call_index}", + ) + ) + result.gaps = list({(gap.code, gap.invocation_id, gap.detail): gap for gap in result.gaps}.values()) + return result + + +@dataclass(frozen=True, slots=True) +class AgentEpisode: + """An Agent response and the observations available at its execution boundary.""" + + response: NeMoGymResponse + observations: AgentObservationBundle diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index 03070f237f..b676c8f49a 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -64,6 +64,7 @@ DRY_RUN_KEY_NAME, HEAD_SERVER_KEY_NAME, NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, + OBSERVABILITY_ENABLED_KEY_NAME, RAY_HEAD_NODE_ADDRESS_KEY_NAME, GlobalConfigDictParser, GlobalConfigDictParserConfig, @@ -71,6 +72,7 @@ get_global_config_dict, ) from nemo_gym.profiling import Profiler +from nemo_gym.rollout_correlation import current_rollout_id, maybe_rollout_id_from_run_body _GLOBAL_AIOHTTP_CLIENT: Union[None, ClientSession] = None @@ -330,10 +332,30 @@ async def request( server_config_dict = get_first_server_config_dict(self.global_config_dict, server_name) base_url = self._build_server_base_url(server_config_dict) + json_obj = kwargs.get("json") if "json" in kwargs: - json_obj = kwargs["json"] if isinstance(json_obj, BaseModel): - kwargs["json"] = json_obj.model_dump(exclude_unset=True) + json_obj = json_obj.model_dump(exclude_unset=True) + kwargs["json"] = json_obj + + observability_enabled = self.global_config_dict.get(OBSERVABILITY_ENABLED_KEY_NAME, False) + server_entry = self.global_config_dict.get(server_name) + rollout_id = current_rollout_id() + if observability_enabled and server_entry is not None and "resources_servers" in server_entry: + if url_path == "/verify": + rollout_id = rollout_id or maybe_rollout_id_from_run_body(json_obj) + if rollout_id is not None and not url_path.startswith(f"/{ROLLOUT_PATH_PREFIX}/"): + url_path = f"{rollout_path_prefix(rollout_id)}{url_path}" + + if ( + rollout_id is not None + and observability_enabled + and server_entry is not None + and "responses_api_models" in server_entry + and url_path.partition("?")[0] in {"/v1/responses", "/v1/chat/completions", "/v1/messages"} + and not url_path.startswith(f"/{ROLLOUT_PATH_PREFIX}/") + ): + url_path = f"{rollout_path_prefix(rollout_id)}{url_path}" return await request(method=method, url=f"{base_url}{url_path}", _internal=True, **kwargs) diff --git a/resources_servers/gdpval/app.py b/resources_servers/gdpval/app.py index 47f36db923..65fdf53ae7 100644 --- a/resources_servers/gdpval/app.py +++ b/resources_servers/gdpval/app.py @@ -50,7 +50,8 @@ SimpleResourcesServer, ) from nemo_gym.config_types import AggregateMetrics, AggregateMetricsRequest, ModelServerRef -from nemo_gym.server_utils import get_server_url +from nemo_gym.rollout_correlation import current_rollout_id +from nemo_gym.server_utils import apply_rollout_prefix, get_server_url from resources_servers.gdpval.judge_panel import ( ResolvedJudge, dir_contains_audio_video, @@ -337,7 +338,7 @@ def _resolve_judges(self) -> List[ResolvedJudge]: legacy_overrides = dict(self.config.judge_responses_create_params_overrides or {}) def _url(server: ModelServerRef) -> str: - return get_server_url(server.name) + "/v1" + return apply_rollout_prefix(get_server_url(server.name), current_rollout_id()) + "/v1" judges: List[ResolvedJudge] = [] for i, member in enumerate(self._effective_panel()): diff --git a/resources_servers/gdpval/tests/test_app.py b/resources_servers/gdpval/tests/test_app.py index dfbaa08475..a97bc99d48 100644 --- a/resources_servers/gdpval/tests/test_app.py +++ b/resources_servers/gdpval/tests/test_app.py @@ -22,6 +22,7 @@ NeMoGymResponseOutputMessage, NeMoGymResponseOutputText, ) +from nemo_gym.rollout_correlation import rollout_context from nemo_gym.server_utils import ServerClient from resources_servers.gdpval.app import ( GDPValResourcesServer, @@ -244,6 +245,7 @@ async def fake_score_with_rubric(**kwargs): body = _verify_request(rubric_json=[{"criterion": "clarity", "score": 1}]) with ( + rollout_context("7-3"), patch("resources_servers.gdpval.scoring.score_with_rubric", side_effect=fake_score_with_rubric), patch("resources_servers.gdpval.app.get_server_url", return_value="http://localhost:9999"), ): @@ -259,7 +261,7 @@ async def fake_score_with_rubric(**kwargs): assert judges[0].create_overrides == {"reasoning_effort": "medium"} assert judges[2].weight == 2.0 # All share the single proxy base_url. - assert {j.base_url for j in judges} == {"http://localhost:9999/v1"} + assert {j.base_url for j in judges} == {"http://localhost:9999/ng-rollout/7-3/v1"} # A seeded rng is threaded through for reproducible sampling. assert captured["rng"] is not None diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py index 6d08dfaa86..1f40bd0450 100644 --- a/responses_api_agents/stirrup_agent/app.py +++ b/responses_api_agents/stirrup_agent/app.py @@ -49,6 +49,7 @@ NeMoGymResponseOutputMessage, NeMoGymResponseOutputText, ) +from nemo_gym.rollout_correlation import current_rollout_id from nemo_gym.server_utils import get_response_json, raise_for_status from responses_api_agents.stirrup_agent.task_strategy import TaskSampleSkipError, TaskStrategy @@ -1037,12 +1038,7 @@ def model_post_init(self, __context: Any) -> None: # -- helpers ---------------------------------------------------------- def _get_model_base_url(self) -> str: - from nemo_gym.global_config import get_first_server_config_dict - from nemo_gym.server_utils import ServerClient - - global_config_dict = ServerClient.load_from_global_config().global_config_dict - model_server_config = get_first_server_config_dict(global_config_dict, self.config.model_server.name) - return f"http://{model_server_config['host']}:{model_server_config['port']}/v1" + return self.resolve_model_base_url(self.config.model_server.name, current_rollout_id()) # -- /v1/responses ---------------------------------------------------- diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py index ad8ff6913a..02441098bd 100644 --- a/responses_api_agents/stirrup_agent/tests/test_app.py +++ b/responses_api_agents/stirrup_agent/tests/test_app.py @@ -25,6 +25,7 @@ NeMoGymResponseOutputMessage, NeMoGymResponseOutputText, ) +from nemo_gym.rollout_correlation import rollout_context from nemo_gym.server_utils import ServerClient from responses_api_agents.stirrup_agent.app import ( NG_FAILURE_CLASS_KEY, @@ -125,6 +126,20 @@ def test_sanity(self) -> None: ) StirrupAgentWrapper(config=config, server_client=MagicMock(spec=ServerClient)) + def test_model_base_url_accepts_rollout_correlation(self) -> None: + wrapper = StirrupAgentWrapper(config=_make_config(), server_client=MagicMock(spec=ServerClient)) + + with patch.object( + StirrupAgentWrapper, + "resolve_model_base_url", + side_effect=lambda _, rollout_id: ( + f"http://model-host:8000/ng-rollout/{rollout_id}/v1" if rollout_id else "http://model-host:8000/v1" + ), + ): + with rollout_context("7-3"): + assert wrapper._get_model_base_url() == "http://model-host:8000/ng-rollout/7-3/v1" + assert wrapper._get_model_base_url() == "http://model-host:8000/v1" + def test_output_history_preserves_nemo_user_tool_results(self) -> None: """Run-history export should keep NeMo user-role tool results as tool outputs.""" history = [ diff --git a/responses_api_agents/tau2/tests/test_app.py b/responses_api_agents/tau2/tests/test_app.py index 53dfd9afb0..620b83c29f 100644 --- a/responses_api_agents/tau2/tests/test_app.py +++ b/responses_api_agents/tau2/tests/test_app.py @@ -17,6 +17,7 @@ from typing import Tuple from unittest.mock import AsyncMock, MagicMock, patch +import pytest from fastapi.testclient import TestClient from nemo_gym.base_responses_api_agent import AggregateMetricsRequest @@ -114,6 +115,45 @@ def _clean(d): assert _clean(expected_response_dict) == _clean(actual_response_dict) + @pytest.mark.parametrize( + ("observability_enabled", "url_suffix"), + [(True, "/ng-rollout/7-2/v1"), (False, "/v1")], + ) + def test_policy_and_user_model_calls_share_rollout_correlation( + self, observability_enabled: bool, url_suffix: str + ) -> None: + example_jsonl = Path(__file__).parent.parent / "data" / "example.jsonl" + request_body = json.loads(example_jsonl.read_text().splitlines()[0]) + request_body |= {"_ng_task_index": 7, "_ng_rollout_index": 2} + + config, server = self._dummy_server() + config.model_server.name = "policy" + config.user_model_server.name = "user" + server.server_client.global_config_dict = {"observability_enabled": observability_enabled} + with patch("responses_api_agents.tau2.app.ensure_tau2_data_dir"): + client = TestClient(server.setup_webserver()) + + class StopRun(Exception): + pass + + captured = {} + + async def stop_after_config(**kwargs): + captured.update(kwargs) + raise StopRun + + model_urls = {"policy": "http://policy:8000", "user": "http://user:8001"} + with ( + patch("responses_api_agents.tau2.app.get_server_url", side_effect=model_urls.__getitem__), + patch("responses_api_agents.tau2.app.run_single_task", stop_after_config), + pytest.raises(StopRun), + ): + client.post("/run", json=request_body) + + config = captured["config"] + assert config.llm_args_agent["api_base"] == model_urls["policy"] + url_suffix + assert config.llm_args_user["api_base"] == model_urls["user"] + url_suffix + async def test_compute_metrics(self) -> None: example_rollouts_fpath = Path(__file__).parent.parent / "data" / "example_rollouts.jsonl" with example_rollouts_fpath.open() as f: diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index c1fd00304d..0e1775ff3f 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -135,6 +135,7 @@ def test_build_model_call_record_from_exchange(): "latency_ms": 18.4, "request": {"input": "hi"}, "response": { + "id": "resp-1", "model": "m", "usage": { "input_tokens": 10, @@ -152,8 +153,10 @@ def test_build_model_call_record_from_exchange(): } rec = build_model_call_record(exchange, call_index=3) assert rec.model_call_id == "call-1" + assert rec.response_id == "resp-1" assert rec.call_index == 3 assert rec.model_ref is not None and rec.model_ref.name == "srv" + assert rec.model == "m" assert rec.dialect == "responses" assert rec.started_at == 100.0 and rec.completed_at == 100.02 assert (rec.tokens_in, rec.tokens_out, rec.tokens_total, rec.tokens_reasoning) == (10, 5, 15, 3) @@ -161,12 +164,19 @@ def test_build_model_call_record_from_exchange(): assert rec.reasoning_content == "thinking..." assert rec.tool_calls == [{"call_id": "c1", "name": "calc", "arguments": {"x": 1}}] assert rec.latency_total_ms == 18.4 + assert build_model_call_record({"response": {"id": 123}}, call_index=0).response_id is None + empty = build_model_call_record({"request": {}, "response": {}}, call_index=0) + assert empty.request == {} + assert empty.response == {} assert { "model_call_id", + "response_id", "call_index", "model_ref", + "model", "dialect", "status_code", + "finish_reason", "started_at", "completed_at", "tokens_in", @@ -175,6 +185,8 @@ def test_build_model_call_record_from_exchange(): "tokens_total", "request", "response", + "request_raw", + "response_raw", "tool_calls", "reasoning_content", "cache_hit", @@ -186,6 +198,36 @@ def test_build_model_call_record_from_exchange(): } <= type(rec).model_json_schema()["properties"].keys() +@pytest.mark.parametrize( + "response,expected", + [ + ({"choices": [{"finish_reason": "tool_calls"}]}, "tool_calls"), + ({"stop_reason": "end_turn"}, "end_turn"), + ({"incomplete_details": {"reason": "max_output_tokens"}}, "max_output_tokens"), + ({"status": "completed"}, None), + ({"choices": {"finish_reason": "invalid"}, "incomplete_details": []}, None), + ], +) +def test_build_model_call_record_normalizes_finish_reason(response, expected): + assert build_model_call_record({"response": response}, call_index=0).finish_reason == expected + + +def test_build_model_call_record_tolerates_malformed_nested_shapes(): + record = build_model_call_record( + { + "response": { + "usage": {"input_tokens": "invalid", "prompt_tokens_details": {"cached_tokens": []}}, + "choices": [{"message": "invalid"}], + "output": [{"type": "reasoning", "summary": [{"text": {}}]}], + } + }, + call_index=0, + ) + + assert record.tokens_total is None + assert record.tool_calls == [] + + def test_capture_is_durable_before_stream_terminal_event_is_sent(tmp_path): import asyncio @@ -232,6 +274,48 @@ async def send(message): assert durable_call_counts == [0, 1, 1] +def test_capture_retains_partial_stream_when_downstream_raises(tmp_path): + import asyncio + + from nemo_gym.base_responses_api_model import _CaptureMiddleware + + store = CaptureStore(tmp_path) + partial = b'data: {"type":"response.output_text.delta","delta":"partial"}\n\n' + + async def app(_scope, receive, send): + await receive() + await send( + {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]} + ) + await send({"type": "http.response.body", "body": partial, "more_body": True}) + raise RuntimeError("stream failed") + + async def receive(): + return {"type": "http.request", "body": b'{"input":"hi"}', "more_body": False} + + async def send(_message): + pass + + with pytest.raises(RuntimeError, match="stream failed"): + asyncio.run( + _CaptureMiddleware(app, store=store, model_server_name="srv")( + { + "type": "http", + "path": "/ng-rollout/partial/v1/responses", + "raw_path": b"/ng-rollout/partial/v1/responses", + "headers": [], + }, + receive, + send, + ) + ) + + [call] = read_model_call_records(store, "partial") + assert call.status_code == 200 + assert call.error_category == "exception" + assert call.response_raw == partial.decode() + + def test_stream_error_events_are_terminal(): from nemo_gym.base_responses_api_model import _consume_terminal_sse_event @@ -318,6 +402,29 @@ async def _boom(body: dict = Body()) -> dict: assert calls[0].latency_ttft_ms is None # nothing streamed before the raise +def test_raised_upstream_error_preserves_status_and_body(tmp_path): + class UpstreamError(RuntimeError): + response = SimpleNamespace(status_code=403, content=b'{"error":{"message":"denied"}}') + + app = FastAPI() + + @app.post("/v1/responses") + async def _boom(body: dict = Body()) -> dict: + raise UpstreamError + + _install_capture(app, tmp_path) + response = TestClient(app, raise_server_exceptions=False).post( + "/ng-rollout/r-upstream/v1/responses", + json={"input": "x"}, + ) + + assert response.status_code == 500 + [exchange] = CaptureStore(tmp_path).read("r-upstream") + assert exchange["status_code"] == 403 + assert exchange["error_category"] == "auth" + assert json.loads(exchange["response_raw"]) == {"error": {"message": "denied"}} + + @pytest.mark.parametrize("request_bytes", [b"{not-json", b"[]"]) def test_invalid_request_body_does_not_drop_capture(tmp_path, request_bytes): app = FastAPI() @@ -343,7 +450,7 @@ async def _responses(body: dict = Body()) -> dict: [record] = read_model_call_records(CaptureStore(tmp_path), "invalid-request") assert record.request is None - assert "request_raw" not in record.model_dump() + assert record.request_raw == request_bytes.decode() def test_per_rollout_url_prefix_correlates_and_is_openai_compatible(tmp_path): @@ -490,16 +597,15 @@ def _boom(_root): assert obs.make_capture_store(config) is None -def test_record_swallows_store_failure(): +def test_record_swallows_store_failure_and_marks_capture_incomplete(tmp_path, monkeypatch): from nemo_gym.base_responses_api_model import _record - class _BadStore: - def record(self, *args, **kwargs): - raise RuntimeError("disk full") + store = CaptureStore(tmp_path) + monkeypatch.setattr(store, "record", MagicMock(side_effect=RuntimeError("disk full"))) # Best-effort: a failing store must not raise out of _record. _record( - _BadStore(), + store, "chat", "srv", b"{}", @@ -512,6 +618,7 @@ def record(self, *args, **kwargs): error_category=None, latency_ms=1.0, ) + assert store.is_incomplete("r") def test_record_falls_back_to_raw_when_request_parser_raises(tmp_path, monkeypatch): @@ -555,8 +662,11 @@ async def _r(body: dict = Body()) -> PlainTextResponse: assert r.status_code == 200 and r.text == "not json" # response passed through unaltered records = CaptureStore(tmp_path).read("rnj") assert len(records) == 1 and records[0]["response"] is None # non-JSON body -> None + assert records[0]["response_raw"] == "not json" # a 2xx whose body we couldn't parse is flagged, not silently counted as a clean success assert records[0]["error_category"] == "capture_parse_error" + [record] = read_model_call_records(CaptureStore(tmp_path), "rnj") + assert record.response_raw == "not json" def test_as_arguments(): @@ -835,11 +945,13 @@ async def gen(): records = CaptureStore(tmp_path).read("3-0") assert len(records) == 1 and records[0]["response"] is not None # reassembled, not dropped + assert records[0]["response_raw"] == r.text # lossless SSE evidence is retained alongside the projection calls = read_model_call_records(CaptureStore(tmp_path), "3-0") assert len(calls) == 1 call = calls[0] assert call.model_call_id + assert call.response_id == "msg_1" assert call.model_ref is not None and call.model_ref.name == "srv" assert call.started_at is not None and call.completed_at is not None assert call.started_at <= call.completed_at @@ -850,13 +962,18 @@ async def gen(): assert call.tool_calls == [{"call_id": "t1", "name": "calc", "arguments": {"x": 1}}] assert call.latency_ttft_ms is not None assert call.error_category is None + assert call.response_raw == r.text def test_reconstruct_chat_sse(): from nemo_gym.base_responses_api_model import _reconstruct_streamed_response chunks = [ - {"model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}]}, + { + "id": "chatcmpl-1", + "model": "m", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}], + }, {"choices": [{"index": 0, "delta": {"content": "lo", "reasoning": "hmm"}}]}, # vLLM `reasoning` alias { "choices": [ @@ -882,6 +999,7 @@ def test_reconstruct_chat_sse(): raw = (b"".join(_sse("", c) for c in chunks) + b"data: [DONE]\n\n").replace(b"\n", b"\r\n") resp = _reconstruct_streamed_response(raw, "chat") msg = resp["choices"][0]["message"] + assert resp["id"] == "chatcmpl-1" assert msg["content"] == "Hello" and msg["reasoning_content"] == "hmm" assert msg["tool_calls"][0]["function"] == {"name": "f", "arguments": '{"a":1}'} assert resp["usage"]["total_tokens"] == 8 @@ -957,17 +1075,41 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path): from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record store = CaptureStore(tmp_path) - store.record( - "0-0", - _capture_exchange( - "responses", - "A", - {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, - {"output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}]}, - ), + exchange = _capture_exchange( + "responses", + "A", + {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, + { + "id": "resp-A", + "output": [{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "ok"}]}], + }, ) - - record = {"_ng_task_index": 0, "_ng_rollout_index": 0, "reward": 1.0, "response": {"harness": "A"}} + exchange["request_raw"] = "malformed request" + exchange["response_raw"] = "malformed response" + store.record("0-0", exchange) + + record = { + "_ng_task_index": 0, + "_ng_rollout_index": 0, + "reward": 1.0, + "response": {"harness": "A"}, + "ng_agent_observations": { + "source": "test", + "records": [ + { + "kind": "agent_invocation", + "invocation_id": "root", + "model_calls": [ + { + "model_ref": {"type": "responses_api_models", "name": "A"}, + "response_id": "resp-A", + } + ], + } + ], + "gaps": [{"code": "model_call_ownership_unavailable"}], + }, + } merge_model_call_capture_into_record(record, [tmp_path]) capture = record["ng_model_call_capture"] @@ -976,33 +1118,74 @@ def test_merge_capture_attaches_metrics_without_raw_payloads(tmp_path): assert capture["metrics"]["num_calls"] == 1 attached_call = capture["calls"][0] assert attached_call["model_call_id"] == "call-A" + assert attached_call["response_id"] == "resp-A" assert attached_call["model_ref"] == {"type": "responses_api_models", "name": "A"} + assert attached_call["model"] == "m" assert attached_call["started_at"] == 100.0 and attached_call["completed_at"] == 100.01 assert attached_call["tokens_in"] == 3 - assert "request" not in attached_call and "response" not in attached_call + assert {"request", "response", "request_raw", "response_raw"}.isdisjoint(attached_call) assert record["response"] == {"harness": "A"} and record["reward"] == 1.0 + [joined_ref] = record["ng_agent_observations"]["records"][0]["model_calls"] + assert joined_ref["model_call_id"] == "call-A" + assert record["ng_agent_observations"]["gaps"] == [] -def test_merge_capture_noop_without_capture(tmp_path): - from nemo_gym.base_responses_api_model import merge_model_call_capture_into_record +def test_merge_capture_reports_missing_capture(tmp_path): + from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record - rec = {"_ng_task_index": 9, "_ng_rollout_index": 9, "reward": 1.0} + rec = { + "_ng_task_index": 9, + "_ng_rollout_index": 9, + "reward": 1.0, + "ng_agent_observations": { + "source": "test", + "gaps": [{"code": "model_call_ownership_unavailable"}], + }, + } merge_model_call_capture_into_record(rec, [tmp_path]) # no capture file for 9-9 - assert "ng_model_call_capture" not in rec - merge_model_call_capture_into_record(rec, []) # no dirs - assert "ng_model_call_capture" not in rec + assert rec["ng_model_call_capture"]["calls"] == [] + assert [gap["code"] for gap in rec["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"] + assert rec["ng_agent_observations"]["gaps"] == [{"code": "model_call_ownership_unavailable"}] + + CaptureStore(tmp_path).path_for("8-8").touch() + empty = {"_ng_task_index": 8, "_ng_rollout_index": 8} + merge_model_call_capture_into_record(empty, [tmp_path]) + assert [gap["code"] for gap in empty["ng_model_call_capture"]["gaps"]] == ["model_call_capture_no_records"] -def test_merge_capture_surfaces_malformed_data_only_when_active(tmp_path): +def test_merge_capture_preserves_valid_records_around_malformed_data(tmp_path): from nemo_gym.base_responses_api_model import merge_model_call_capture_into_record store = CaptureStore(tmp_path) - store.path_for("9-9").write_bytes(b"{not-json}\n") - record = {"_ng_task_index": 9, "_ng_rollout_index": 9} + first = _capture_exchange("responses", "A", {}, {"id": "resp-A"}) + third = _capture_exchange("responses", "B", {}, {"id": "resp-B"}) + store.path_for("9-8").write_bytes(orjson.dumps(first) + b"\n{not-json}\n" + orjson.dumps(third) + b"\n") + record = {"_ng_task_index": 9, "_ng_rollout_index": 8} - merge_model_call_capture_into_record(record, []) - with pytest.raises(orjson.JSONDecodeError): - merge_model_call_capture_into_record(record, [tmp_path]) + merge_model_call_capture_into_record(record, [tmp_path]) + + assert [call["call_index"] for call in record["ng_model_call_capture"]["calls"]] == [0, 2] + assert [call["response_id"] for call in record["ng_model_call_capture"]["calls"]] == ["resp-A", "resp-B"] + assert record["ng_model_call_capture"]["gaps"] == [ + { + "code": "model_call_capture_records_unreadable", + "detail": "count=1", + } + ] + + +def test_merge_capture_reports_partial_write_loss(tmp_path): + from nemo_gym.base_responses_api_model import CaptureStore, merge_model_call_capture_into_record + + store = CaptureStore(tmp_path) + store.record("4-2", _capture_exchange("responses", "A", {}, {"id": "resp-A"})) + store.mark_incomplete("4-2") + record = {"_ng_task_index": 4, "_ng_rollout_index": 2} + + merge_model_call_capture_into_record(record, [tmp_path]) + + assert len(record["ng_model_call_capture"]["calls"]) == 1 + assert [gap["code"] for gap in record["ng_model_call_capture"]["gaps"]] == ["model_call_capture_incomplete"] def test_clear_model_call_captures_for_rollouts_run_scoping(tmp_path, monkeypatch): @@ -1012,11 +1195,12 @@ def test_clear_model_call_captures_for_rollouts_run_scoping(tmp_path, monkeypatc store = CaptureStore(tmp_path) store.record("0-0", {"dialect": "chat", "request": {}, "response": {}}) store.record("1-0", {"dialect": "chat", "request": {}, "response": {}}) + store.mark_incomplete("0-0") assert store.read("0-0") and store.read("1-0") # Clears only the rollout ids about to be (re)run; rows without indices are skipped, others stay. clear_model_call_captures_for_rollouts([{"_ng_task_index": 0, "_ng_rollout_index": 0}, {"no": "id"}], [tmp_path]) - assert store.read("0-0") == [] and store.read("1-0") + assert store.read("0-0") == [] and not store.is_incomplete("0-0") and store.read("1-0") clear_model_call_captures_for_rollouts([{"_ng_task_index": 1, "_ng_rollout_index": 0}], []) # no dirs -> no-op assert store.read("1-0") @@ -1126,10 +1310,40 @@ def _write(i: int) -> None: for t in threads: t.join() rows = store.read("0-0") - assert len(rows) == 20 # flock + in-process lock: no lost or corrupted appends + assert len(rows) == 20 # flock prevents lost or corrupted appends assert sorted(r["request"]["i"] for r in rows) == list(range(20)) +def test_capture_store_does_not_serialize_independent_rollouts(tmp_path, monkeypatch): + import os + import threading + + store = CaptureStore(tmp_path) + first_fsync = threading.Event() + release_first = threading.Event() + original_fsync = os.fsync + + def _fsync(fd): + if threading.current_thread().name == "first-rollout": + first_fsync.set() + assert release_first.wait(timeout=5) + original_fsync(fd) + + monkeypatch.setattr(os, "fsync", _fsync) + first = threading.Thread(name="first-rollout", target=store.record, args=("0-0", {"request": {}, "response": {}})) + second = threading.Thread(target=store.record, args=("1-0", {"request": {}, "response": {}})) + first.start() + assert first_fsync.wait(timeout=5) + second.start() + second.join(timeout=1) + try: + assert not second.is_alive() + assert store.read("1-0") == [{"request": {}, "response": {}}] + finally: + release_first.set() + first.join(timeout=5) + + def test_capture_store_read_waits_for_in_progress_append(tmp_path): import fcntl import threading diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index ece7bacb1b..d3145eb403 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -994,7 +994,14 @@ def setup_server_client(self): {AGENT_REF_KEY_NAME: {"name": "my_agent"}, TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 1}, ] results = [ - {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 10}}}, + { + TASK_INDEX_KEY_NAME: 0, + ROLLOUT_INDEX_KEY_NAME: 0, + "reward": 1.0, + "response": {"usage": {"tokens": 10}}, + "ng_agent_observations": {"invocations": [{"conversation": ["large"]}]}, + "ng_model_call_capture": {"calls": [{"request": "large"}]}, + }, {TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 12}}}, {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 0, "reward": 1.0, "response": {"usage": {"tokens": 8}}}, {TASK_INDEX_KEY_NAME: 1, ROLLOUT_INDEX_KEY_NAME: 1, "reward": 0.0, "response": {"usage": {"tokens": 15}}}, @@ -1023,6 +1030,8 @@ def setup_server_client(self): ) for item in sent_data: assert "responses_create_params" not in item + assert "ng_agent_observations" not in item + assert "ng_model_call_capture" not in item assert "usage" in item["response"] async def test_call_aggregate_metrics_multiple_agents(self, tmp_path: Path) -> None: diff --git a/tests/unit_tests/test_rollout_correlation.py b/tests/unit_tests/test_rollout_correlation.py new file mode 100644 index 0000000000..3954ece3fc --- /dev/null +++ b/tests/unit_tests/test_rollout_correlation.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +from urllib.parse import urlsplit + +import orjson +import pytest +from fastapi import Body, FastAPI +from fastapi.responses import JSONResponse +from omegaconf import OmegaConf +from pydantic import ConfigDict +from starlette.testclient import TestClient + +import nemo_gym.server_utils +from nemo_gym.base_resources_server import ( + BaseResourcesServerConfig, + BaseRunRequest, + BaseVerifyRequest, + BaseVerifyResponse, + SimpleResourcesServer, +) +from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig, SimpleResponsesAPIAgent +from nemo_gym.base_responses_api_model import ( + CaptureStore, + ModelCallCaptureConfig, + install_model_call_capture, + merge_model_call_capture_into_record, +) +from nemo_gym.config_types import BaseServerConfig +from nemo_gym.rollout_correlation import maybe_rollout_id_from_run_body +from nemo_gym.server_utils import ServerClient, get_response_json + + +def _model_response(model: str, text: str = "") -> dict: + return { + "id": f"resp-{model}", + "created_at": 0.0, + "model": model, + "object": "response", + "output": [ + { + "id": f"msg-{model}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + + +def _model_app(capture_dir, name: str) -> FastAPI: + app = FastAPI() + + @app.post("/v1/responses") + async def responses(body: dict = Body()) -> JSONResponse: + return JSONResponse(_model_response(name, "[[A=B]]" if name == "judge" else "answer")) + + install_model_call_capture( + app, + ModelCallCaptureConfig(observability_enabled=True, model_call_capture_dir=capture_dir), + model_server_name=name, + ) + return app + + +class _JudgeResourcesServer(SimpleResourcesServer): + def setup_webserver(self) -> FastAPI: + app = super().setup_webserver() + app.post("/tool")(self.tool) + return app + + async def tool(self, body: dict = Body()) -> dict: + tool_model = await self.server_client.post( + server_name="tool_model", + url_path="/v1/responses", + json={"input": body["input"]}, + ) + return await get_response_json(tool_model) + + async def verify(self, body: BaseVerifyRequest) -> BaseVerifyResponse: + judge = await self.server_client.post( + server_name="judge", + url_path="/v1/responses", + json={"input": "grade"}, + ) + await get_response_json(judge) + return BaseVerifyResponse(**body.model_dump(), reward=1.0) + + +class _AgentRunRequest(BaseRunRequest): + model_config = ConfigDict(extra="allow") + + +class _Agent(SimpleResponsesAPIAgent): + async def responses(self, body): + raise NotImplementedError + + async def run(self, body: _AgentRunRequest) -> BaseVerifyResponse: + policy = await self.server_client.post( + server_name="policy", + url_path="/v1/responses", + json=body.responses_create_params, + ) + response = orjson.loads(await policy.read()) + tool = await self.server_client.post( + server_name="resources", + url_path="/tool", + json={"input": "lookup"}, + ) + await get_response_json(tool) + verify = await self.server_client.post( + server_name="resources", + url_path="/verify", + json={ + "responses_create_params": body.responses_create_params.model_dump(), + "response": response, + }, + ) + return BaseVerifyResponse.model_validate(orjson.loads(await verify.read())) + + +class _Response: + def __init__(self, response) -> None: + self.status = response.status_code + self.ok = response.is_success + self.cookies = response.cookies + self._content = response.content + + async def read(self) -> bytes: + return self._content + + +@pytest.mark.asyncio +async def test_verify_correlates_policy_and_judge_calls_and_preserves_raw_capture(tmp_path, monkeypatch) -> None: + capture_dir = tmp_path / "captures" + config = OmegaConf.create( + { + "observability_enabled": True, + "policy": {"responses_api_models": {"model": {"host": "policy.test", "port": 80}}}, + "tool_model": {"responses_api_models": {"model": {"host": "tool-model.test", "port": 80}}}, + "judge": {"responses_api_models": {"model": {"host": "judge.test", "port": 80}}}, + "resources": {"resources_servers": {"judge": {"host": "resources.test", "port": 80}}}, + "agent": {"responses_api_agents": {"agent": {"host": "agent.test", "port": 80}}}, + } + ) + server_client = ServerClient( + head_server_config=BaseServerConfig(host="head.test", port=80), + global_config_dict=config, + ) + resources = _JudgeResourcesServer( + config=BaseResourcesServerConfig( + host="resources.test", + port=80, + entrypoint="app.py", + name="resources", + ), + server_client=server_client, + ) + agent = _Agent( + config=BaseResponsesAPIAgentConfig( + host="agent.test", + port=80, + entrypoint="app.py", + name="agent", + ), + server_client=server_client, + ) + clients = { + "policy.test": TestClient(_model_app(capture_dir, "policy")), + "tool-model.test": TestClient(_model_app(capture_dir, "tool_model")), + "judge.test": TestClient(_model_app(capture_dir, "judge")), + "resources.test": TestClient(resources.setup_webserver()), + "agent.test": TestClient(agent.setup_webserver()), + } + + async def dispatch(method: str, url: str, **kwargs): + parsed = urlsplit(url) + response = clients[parsed.hostname].request(method, parsed.path, json=kwargs.get("json")) + return _Response(response) + + monkeypatch.setattr(nemo_gym.server_utils, "request", dispatch) + + verify = await server_client.post( + server_name="agent", + url_path="/run", + json={ + "_ng_task_index": 4, + "_ng_rollout_index": 2, + "responses_create_params": {"input": "solve"}, + }, + ) + assert orjson.loads(await verify.read())["reward"] == 1.0 + + store = CaptureStore(capture_dir) + capture_path = store.path_for("4-2") + assert capture_path.is_file() + exchanges = store.read("4-2") + assert [exchange["model_ref"]["name"] for exchange in exchanges] == ["policy", "tool_model", "judge"] + assert all(exchange.get("request") is not None or exchange.get("request_raw") for exchange in exchanges) + assert all(exchange.get("response") is not None or exchange.get("response_raw") for exchange in exchanges) + + rollout = {"_ng_task_index": 4, "_ng_rollout_index": 2} + merge_model_call_capture_into_record(rollout, [capture_dir]) + assert capture_path.is_file() + assert len(capture_path.read_bytes()) > 0 + + +def test_rollout_id_does_not_serialize_run_body() -> None: + class UndumpableRunRequest(_AgentRunRequest): + def model_dump(self, *args, **kwargs): + raise AssertionError("run body must not be serialized") + + body = UndumpableRunRequest.model_validate( + { + "_ng_task_index": 4, + "_ng_rollout_index": 2, + "responses_create_params": {"input": "solve"}, + } + ) + + assert maybe_rollout_id_from_run_body(body) == "4-2" diff --git a/tests/unit_tests/test_rollout_observability.py b/tests/unit_tests/test_rollout_observability.py new file mode 100644 index 0000000000..2a09972cf5 --- /dev/null +++ b/tests/unit_tests/test_rollout_observability.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pydantic import ValidationError + +from nemo_gym.base_responses_api_model import ModelCallRecord +from nemo_gym.config_types import ModelServerRef +from nemo_gym.rollout_observability import ( + AgentInvocation, + AgentObservationBundle, + ContextCompactionObservation, + ModelCallRef, + ObservationGap, + SandboxObservation, + ToolCallObservation, + join_model_call_observations, +) + + +@pytest.mark.parametrize( + "value", + ({}, {"response_id": "resp-1"}, {"model_ref": {"name": "policy", "type": "responses_api_models"}}), +) +def test_model_call_ref_rejects_incomplete_join_keys(value: dict) -> None: + with pytest.raises(ValidationError, match="model_call_id or both model_ref and response_id"): + ModelCallRef.model_validate(value) + + +def test_observation_bundle_rejects_duplicate_invocation_ids() -> None: + with pytest.raises(ValidationError, match="invocation_id must be unique"): + AgentObservationBundle( + source="test", + records=[AgentInvocation(invocation_id="root"), AgentInvocation(invocation_id="root")], + ) + + +def test_observation_models_reject_unknown_fields() -> None: + with pytest.raises(ValidationError, match="producer_extension"): + ModelCallRef.model_validate({"model_call_id": "call-1", "producer_extension": "unexpected"}) + + +@pytest.mark.parametrize( + "timing", + ( + {"duration_ms": -1}, + {"started_at": 2.0, "completed_at": 1.0}, + ), +) +def test_tool_call_observation_rejects_invalid_timing(timing: dict) -> None: + with pytest.raises(ValidationError): + ToolCallObservation(invocation_id="root", tool_call_id="call-1", **timing) + + +def test_agent_invocation_rejects_negative_duration() -> None: + with pytest.raises(ValidationError): + AgentInvocation(invocation_id="root", duration_ms=-1) + + +def test_join_model_calls_resolves_exact_references_and_reports_unowned_calls() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], + ) + ], + gaps=[ObservationGap(code="model_call_ownership_unavailable")], + ) + calls = [ + ModelCallRecord( + model_call_id="call-1", + response_id="resp-1", + model_ref=model_ref, + call_index=0, + ), + ModelCallRecord(model_call_id="call-2", model_ref=model_ref, call_index=1), + ] + + joined = join_model_call_observations(bundle, calls) + + [invocation] = [record for record in joined.records if isinstance(record, AgentInvocation)] + [joined_call] = invocation.model_calls + assert joined_call.model_call_id == "call-1" + assert joined_call.model_ref == model_ref + assert joined_call.response_id == "resp-1" + ownership_gaps = [gap for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"] + assert [gap.detail for gap in ownership_gaps] == ["capture:call-2:call_index=1"] + + +def test_join_model_calls_does_not_guess_ambiguous_response_ids() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ModelCallRef(model_ref=model_ref, response_id="resp-1")], + ) + ], + ) + calls = [ + ModelCallRecord(model_call_id="call-1", response_id="resp-1", model_ref=model_ref, call_index=0), + ModelCallRecord(model_call_id="call-2", response_id="resp-1", model_ref=model_ref, call_index=1), + ] + + joined = join_model_call_observations(bundle, calls) + + [invocation] = [record for record in joined.records if isinstance(record, AgentInvocation)] + assert invocation.model_calls[0].model_call_id is None + assert "model_call_reference_ambiguous" in {gap.code for gap in joined.gaps} + assert [gap.detail for gap in joined.gaps if gap.code == "model_call_ownership_unavailable"] == [ + "capture:call-1:call_index=0", + "capture:call-2:call_index=1", + ] + + +def test_join_model_calls_reports_conflicting_and_unmatched_references() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + AgentInvocation( + invocation_id="root", + model_calls=[ + ModelCallRef(model_call_id="call-1"), + ModelCallRef(model_ref=model_ref, response_id="resp-1"), + ModelCallRef(model_call_id="missing"), + ], + ) + ], + gaps=[ + ObservationGap( + code="model_call_ownership_unavailable", + invocation_id="root", + detail="producer gap", + ) + ], + ) + calls = [ + ModelCallRecord( + model_call_id="call-1", + response_id="resp-1", + model_ref=model_ref, + call_index=0, + ) + ] + + joined = join_model_call_observations(bundle, calls) + joined_again = join_model_call_observations(joined, calls) + + assert joined_again.model_dump() == joined.model_dump() + assert {gap.code for gap in joined.gaps} == { + "model_call_ownership_unavailable", + "model_call_reference_conflict", + "model_call_reference_unmatched", + } + assert any(gap.detail == "producer gap" for gap in joined.gaps) + + +def test_join_model_calls_resolves_compaction_boundaries() -> None: + model_ref = ModelServerRef(name="policy", type="responses_api_models") + bundle = AgentObservationBundle( + source="test", + records=[ + ContextCompactionObservation( + invocation_id="root", + before_model_call=ModelCallRef(model_ref=model_ref, response_id="before"), + after_model_call=ModelCallRef(model_ref=model_ref, response_id="after"), + ) + ], + ) + calls = [ + ModelCallRecord( + model_call_id="call-before", + response_id="before", + model_ref=model_ref, + call_index=0, + ), + ModelCallRecord( + model_call_id="call-after", + response_id="after", + model_ref=model_ref, + call_index=1, + ), + ] + + joined = join_model_call_observations(bundle, calls) + [compaction] = [record for record in joined.records if isinstance(record, ContextCompactionObservation)] + + assert compaction.before_model_call.model_call_id == "call-before" + assert compaction.after_model_call.model_call_id == "call-after" + + +def test_sandbox_observation_rejects_negative_usage() -> None: + with pytest.raises(ValidationError): + SandboxObservation(role="agent", cpu_time_s=-1)