Simplify model call record by bxyu-nvidia · Pull Request #1992 · NVIDIA-NeMo/Gym

bxyu-nvidia · 2026-07-10T17:55:25Z

TODOs as of Fri Jul 10 11:39PM

Key design choices

Decouple model call record logging from the route-specific logic
Logging occurs in the route handler which can capture the status code and generically capture stream responses.
Responses API models can override their relevant route function in order to change the route-specific logic e.g. anthropic_model proxying to an actual Anthropic model vs converting to Responses can handle their model call record creation logic differently
Re-use the FastAPI routing logic rather than making a bespoke regex
Clear future path for opt-in tool call timing by retrieving the current model call record from request.state
Re-use the FastAPI thread pool for running sync tasks for the logging to disk
Once we upgrade to truly asynchronous file ops, we can simply swap the logging function from sync to async and FastAPI will run the background task async rather than using a thread pool

Other consequences

The logging to disk is set to a background task which runs after the http response has been sent back so they don't slow down the actual response

…1483) Opt-in (off by default) FastAPI middleware on every Gym model server that records each /v1/responses, /v1/chat/completions, and /v1/messages exchange -- including failed calls -- into a per-rollout, durable CaptureStore. Best-effort; never alters the response. Delivers the #1483 per-step contract as a typed StepRecord (JSONSchema): token accounting (in/out/reasoning/total), request/response/tool_calls, structured reasoning_content, cache hit/miss, latency, normalized error_category/status_code, and trial/turn/step indices. assemble_step_records builds the per-trial trajectory; aggregate_rollout_metrics the per-rollout totals; assemble_rollout the eval-only NeMoGym trajectory (token-ids / log-probs are intentionally not surfaced -- on-policy RL assembly stays on the RL side and consumes this same capture). Correlation is OpenAI-compatible: callers set the x-nemo-gym-rollout-id header or point their base_url at a /ng-rollout/<rollout_id>/v1 prefix (stripped before routing). Wired uniformly via the base agent (resolve_model_base_url + rollout_call_kwargs) so agents need no correlation code; hermes adopts it as the reference. Tested: unit coverage of the contract, classifiers, capture/assembly, and correlation; an e2e through the real SimpleModelServer install path; validated live against the inference gateway. Signed-off-by: Michal Bien <mbien@nvidia.com>

…ation Rewrite the capture middleware as a pure ASGI middleware so it composes with streaming (SSE) responses: it forwards send chunks unchanged while buffering the request and non-streaming response. This unblocks the Claude Code CLI path, which always streams /v1/messages (the previous BaseHTTPMiddleware could not both read the request body and return a streaming response). Wire claude_code_agent correlation: run() forwards the rollout id to its /v1/responses self-call, and _run_claude_code applies the /ng-rollout/<id> prefix to ANTHROPIC_BASE_URL so the CLI's streaming /v1/messages calls key to the rollout. Tests: streaming /v1/messages e2e through the real model server (forwarded intact and captured + correlated); claude_code prefix unit tests; scope-based header helpers. Docs: streaming-response note. Signed-off-by: Michal Bien <mbien@nvidia.com>

…d + attempt id Fold each rollout's captured model-call trajectory into its rollout_collection record as ng_trajectory_capture = {rollout_id, metrics, steps} -- the same StepRecord shape for every agent harness (raw request/response stay in the capture store). The capture is authoritative for per-step model-call stats; the harness output and reward on the record are untouched. Add a resume attempt identifier: rollout_id_from_run_body appends -a<n> when _ng_attempt_index > 0 (attempt 0 keeps the bare key, backward-compatible), and _load_from_cache stamps the prior failure count on re-dispatched rows, so a retry's captured trajectory stays separable from the prior attempt's. This gives CLU the NEL-Next-equivalent problem_id (task_index), seed (rollout_index), session (rollout_id), and retry identifiers on one record. capture_dirs_from_config resolves capture dirs from $NEMO_GYM_TRAJECTORY_DIR + observability -enabled model-server configs; peek_global_config_dict reads the cached config without forcing a CLI parse. Best-effort: the merge never alters or breaks rollout collection. Tests for attempt-aware rollout id, capture-dir resolution, and uniform merge shape across agents/dialects. HTTP-level retry_count remains a follow-up. Signed-off-by: Michal Bien <mbien@nvidia.com>

… tests Review: merge_capture_into_record read and parsed each rollout's capture file twice (once via assemble_step_records, again inside aggregate_rollout_metrics). Extract a pure aggregate_step_records(steps) and reuse the steps already read, so the merge reads the capture once per rollout; aggregate_rollout_metrics now delegates to it (behavior unchanged). Also fix the pre-existing Python 3.12 argparse failures in the CLI "did you mean" helper: 3.12 no longer quotes invalid-choice options, so parse both quoted and unquoted choice lists. Full unit suite green (620); openai capture e2e green (5). Signed-off-by: Michal Bien <mbien@nvidia.com>

…pture dir by config key Address P0 review feedback: - The /ng-rollout/<id> correlation prefix is now always stripped before routing (the capture middleware is always installed), so a default `gym eval` with capture off no longer 404s on every prefixed model call. Capture still only records when observability_enabled. - capture_dirs_from_config resolves a server's default capture dir off its config key (which equals the producer's config.name), instead of the "model_server" fallback, so the merge reads the same directory the producer wrote to. Warn once when capture is enabled but no dir resolves. Regression tests added for both (prefix-stripped-when-disabled, resolve-by-config-key, warn-once). Signed-off-by: Michal Bien <mbien@nvidia.com>

… cache tokens, off-loop writes - num_workers>1: flock(LOCK_EX) around the capture append+fsync so concurrent same-rollout writes from different worker processes can't interleave/corrupt (the in-process lock is kept). - Anthropic prompt caching: extract_token_stats folds cache_read + cache_creation into tokens_in for the native /v1/messages usage shape (true prompt size) and surfaces cache_creation_tokens; OpenAI usage (where cached_tokens is a subset of input_tokens) is left untouched to avoid double counting. - Offload the capture write+fsync to asyncio.to_thread so it never stalls the event loop; fsync kept for durability. - Document trial_index/turn_index as reserved (step_index is the authoritative per-call index). Regression tests: concurrent-append integrity, Anthropic fold, OpenAI no-double-count, cache_creation. Signed-off-by: Michal Bien <mbien@nvidia.com>

…mbly) Close the remaining (P2) review threads and the streaming gap so capture is uniform across harnesses and model providers: - Streaming (SSE): the capture middleware now reassembles the streamed response per dialect (Anthropic Messages, Chat Completions, Responses), so streamed /v1/messages calls -- the path Claude Code always uses -- capture response, tokens_*, tool_calls, and reasoning_content just like a non-streamed call, and record latency_ttft_ms (time to first byte). Best-effort: an unparseable stream leaves response null with request/correlation/status still recorded. - num_turns: report None when no turn_index markers exist (a turn is not a model step) instead of silently equating it to num_steps; num_steps stays authoritative. - reasoning: accept the vLLM `reasoning` field alongside `reasoning_content`. - Guard non-dict items in the Responses output loop. - hermes run(): serialize the run body once (reused for seed + verify + rollout correlation). - Remove the unused summarize_response/_usage helpers (and the stale "Compact CLU telemetry" note). - Docs: tighten the turn definition, note turn_index/num_turns are reserved/null, and document that streamed responses are now reassembled (+ latency_ttft_ms). Regression tests: SSE reassembly through the server (Anthropic) + per-dialect reconstruction units + best-effort-None, num_turns-None-without-markers, reasoning alias, non-dict output guard, raised-call capture. Affected suites green (observability/e2e 49, hermes 16, rollout 41, claude 44); ruff clean. Signed-off-by: Michal Bien <mbien@nvidia.com>

…e refs Make the capture code self-contained for review: trim a few narrative docstrings/comments and remove internal issue references. No behavior change. Signed-off-by: Michal Bien <mbien@nvidia.com>

…ure loops too The Responses output loop already skipped non-dict items; apply the same guard to the Chat Completions tool_calls and Anthropic content-block loops so a malformed payload can't crash _tool_calls_and_reasoning. Signed-off-by: Michal Bien <mbien@nvidia.com>

…ulated fields - turn_index is derived from user-message boundaries (a turn = control returned to the user) along a single linear thread; it is null for non-linear captures (sub-agent interleaving, history rewriting, or a delta wire's later user turn). num_turns aggregates the distinct derived turns. - Run-scope per-rollout capture: a fresh (non-resume) run clears its rollouts' stale capture files first, so a re-run does not append onto a previous run's capture. - Remove the unused trial/turn request-header path and the unpopulated trial_index / retry_count fields; step_index and num_steps stay the authoritative per-call / per-rollout counts. Signed-off-by: Michal Bien <mbien@nvidia.com>

…gate, SSE/output guards, error category; per-agent correlation; docs - capture_dirs_from_config resolves a server's dir off the top-level instance key (== producer config.name) instead of the leaf impl key, so a bare observability_enabled config no longer silently no-ops; rewrote the test to drive the real producer (make_capture_store) + consumer and assert they agree. - claude_code: gate the /ng-rollout prefix on model_server being set, so a real Anthropic endpoint no longer 404s on every call; added a test. - capture middleware: fold the body parse + SSE reassembly into the off-loop record and guard it (a malformed body can't surface as an ASGI error after the response was sent); flag a 2xx with an unparseable body as capture_parse_error instead of a silent success. - trajectory_capture: guard a non-list Responses output so a malformed payload can't drop the whole merge. - correlation: wire the rollout id into mini_swe_agent_2 and gymnasium_agent (which reach the model via a manual base_url / ServerClient.post rather than the base agent); added a gymnasium correlation test. - tests: cross-process flock regression for the multi-worker capture path. - docs: link Trial/Turn/Step to key-terminology; correct the streamed-body "buffered" wording and the per-agent correlation note. Signed-off-by: Michal Bien <mbien@nvidia.com>

…d chat choices Post-review follow-ups to 88cfb2c: - install_trajectory_capture docstring still claimed streamed SSE bodies are "not buffered"; they are buffered for reassembly (the class docstring was already corrected) -- align the wording. - _tool_calls_and_reasoning: guard the Chat Completions choices with isinstance(list) for parity with the Responses output and Anthropic content branches, so a malformed (non-list) choices can't raise and drop the rollout merge. Signed-off-by: Michal Bien <mbien@nvidia.com>

…ring) - extract_token_stats: use a 0 base for fully-cached Anthropic responses (only cache_read/creation, no input_tokens) so tokens_in is preserved instead of dropping to null; add regression test - doc tokens_reasoning as OpenAI/Responses-only (null != 0 for Anthropic) and tokens_in as a prompt-size (not cost) metric - doc num_turns as best-effort/often-null; num_steps authoritative - document the write-after-forward ordering: the durable capture JSONL is the merge's source of truth, num_steps recomputed from it Signed-off-by: Michal Bien <mbien@nvidia.com>

Signed-off-by: Michal Bien <mbien@nvidia.com>