Repository navigation
feat(experimentalist): support ATIF agent traces - #1127
Conversation
benchmarks/run.py writes each run under artifacts/experimentalist-benchmarks/. That is generated output, but it was not ignored: besides adding tens of megabytes to git status, it broke `ruff check plugins/nemo-experimentalist/`, because the optimizer writes candidate agents there and a candidate whose generated code does not compile is a normal outcome rather than a lint failure. Follows the convention already used by nemo-guardrails/benchmarks/artifacts. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Christian Schüller <cschueller@nvidia.com>
…ces dir When the agent dies before writing any traces, /app/traces never exists and the environment refuses to copy it, raising RuntimeError. That escaped the metrics guard and replaced the agent's own error with a Docker copy failure, so a candidate rejected for an IndentationError reported "Could not find the file /app/traces/. in container" instead of the traceback that caused it. Catching RuntimeError lets the run reach the existing check below, which raises with the agent's exit code and stderr. A missing traces directory is a symptom of the agent dying, never the cause. Independent of ATIF support; surfaced while diagnosing a rejected candidate. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Christian Schüller <cschueller@nvidia.com>
The Experimentalist read agent traces as OTLP only. An agent that emits ATIF produced trials with no trace at all, silently: discovery found no *.jsonl, the trial was skipped, and nothing said why. Traces now carry their format from discovery through analysis: - Discovery recognises *.atif.json alongside *.jsonl and stamps trace_format on the ResourceRef, so the format is decided once and carried forward rather than re-derived downstream. HarborEvaluatorConfig.trace_format selects which artifact becomes the trial's trace; both stay exposed in resources. - Upload posts the trajectory to Intake's ATIF ingest with evaluation_context carrying experiment and test-case identity. Trial identity rides on the trajectory's own session_id, which Intake adopts as the trace id. Intake maps ATIF to spans server-side. - The trace explorer gains a parser for Intake rows. Intake promotes known attributes into typed columns and strips those keys from raw_attributes, so a reader that knows only the OpenInference names finds nothing -- token counts and error fields were missing on every span read back from storage, OTLP included. from_intake now reads the columns directly; _load_trace_models keeps a per-record shape check, because it deliberately accepts Intake-shaped payloads from disk. An ATIF branch keyed on the source column covers fields that have no column at all. Three silent failures become diagnosable: discovery warns when the configured format matches nothing but the other format's artifacts are present, persist_evaluation logs the trial it drops, and from_ref rejects a local .atif.json with a message naming the real cause rather than reporting an empty trace. extra.verifier_result is deliberately not injected. persist_evaluation already creates evaluator_results per trial metric, so deriving a second set from the trajectory would double-count. Verified end to end against a live platform: ATIF trajectories accepted by Intake's validators, spans read back with token counts and tool call ids populated, and experiment grouping intact. Per-step durations are absent, since Intake takes a step's end time only from a producer-supplied invocation block; trajectory-level duration is unaffected. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Christian Schüller <cschueller@nvidia.com>
📝 WalkthroughWalkthroughThis PR adds selectable ATIF trace support to Harbor evaluation, backend ingestion, and TraceExplorer. It adds ATIF payload construction and upload, Intake-row span parsing, local ATIF-file validation, wrapper error handling, tests, artifact filtering, and README guidance. ChangesATIF trace support
Sequence Diagram(s)sequenceDiagram
participant HarborEvaluator
participant ExperimentalistBackend
participant IntakeAPI
participant TraceExplorer
HarborEvaluator->>ExperimentalistBackend: persist selected trace format
ExperimentalistBackend->>IntakeAPI: upload OTLP protobuf or ATIF JSON
IntakeAPI-->>ExperimentalistBackend: return trace identifier
TraceExplorer->>IntakeAPI: load ingested trace rows
IntakeAPI-->>TraceExplorer: return Intake rows
TraceExplorer->>TraceExplorer: convert rows into spans
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py (2)
36-48: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueThe trajectory file is parsed twice per trial.
_persist_trialcallsread_session_id, then_upload_trace_atifcallsbuild_ingest_payload, and each call re-reads and re-parses the file. Trajectories can be large. Consider loading once and passing the parsed dict, or returning both the session id and the payload from one call.Also applies to: 76-88
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py` around lines 36 - 48, The trial persistence flow currently parses the trajectory twice through read_session_id and build_ingest_payload. Update _persist_trial and the related _upload_trace_atif/build_ingest_payload flow to load the ATIF file once, reuse the parsed dictionary for session_id validation and payload construction, and preserve the existing error behavior for missing or invalid session IDs.
27-33: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
local_path_from_uriin_load. The helper inentities.pydecodes escaped paths and rejects non-local schemes. ReplacePath(urlparse(ref.uri).path)withlocal_path_from_uri(ref.uri, context="ATIF trajectory").🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py` around lines 27 - 33, Update _load to obtain the trajectory path through local_path_from_uri(ref.uri, context="ATIF trajectory") instead of constructing a Path from urlparse(ref.uri).path, preserving the existing JSON parsing and object validation.plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/evaluator/harbor.py (1)
657-659: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueType
trace_formatasLiteral["otlp", "atif"].
HarborEvaluatorConfig.trace_formatis aLiteral, but these helpers accept anystr. An unexpected value silently selects OTLP. Narrowing the annotation makes the contract checkable byty.♻️ Proposed change
def _trial_resources( - trial_dir: Path, *, trace_format: str = "otlp" + trial_dir: Path, *, trace_format: Literal["otlp", "atif"] = "otlp" ) -> tuple[dict[str, ResourceRef], ResourceRef | None]:async def _trials_from_dir( - self, job_dir: Path, tasks: Sequence[Task], *, trace_format: str = "otlp" + self, job_dir: Path, tasks: Sequence[Task], *, trace_format: Literal["otlp", "atif"] = "otlp" ) -> Sequence[TrialResult]:Note that
_runreadstrace_formatfromoptions_dict.pop(...), which isAny, so a cast or re-read fromoptions.trace_formatkeeps the types clean.Also applies to: 1321-1323
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/evaluator/harbor.py` around lines 657 - 659, Update the trace_format annotations for _trial_resources and the other helper at the referenced location to Literal["otlp", "atif"], importing Literal as needed. In _run, avoid passing the Any result from options_dict.pop directly; cast it to the literal type or reuse options.trace_format so the narrowed contract remains type-safe.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@plugins/nemo-experimentalist/README.md`:
- Around line 65-69: Update the YAML configuration example in the
EvolutionaryOptimizerConfig documentation to place evaluator.trace_format at the
top level, removing the incorrect optimizer nesting while preserving the
existing trace_format values and comment.
In
`@plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_explorer.py`:
- Around line 842-933: The non-ATIF TOOL branch in _span_from_intake_row
currently discards tool call identifiers by setting tool_call_id to an empty
string. Restore the OTLP-compatible attribute fallback used by
_span_from_otlp_record, while preserving the existing ATIF-specific
_atif_tool_call_id(attrs) behavior.
---
Nitpick comments:
In
`@plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py`:
- Around line 36-48: The trial persistence flow currently parses the trajectory
twice through read_session_id and build_ingest_payload. Update _persist_trial
and the related _upload_trace_atif/build_ingest_payload flow to load the ATIF
file once, reuse the parsed dictionary for session_id validation and payload
construction, and preserve the existing error behavior for missing or invalid
session IDs.
- Around line 27-33: Update _load to obtain the trajectory path through
local_path_from_uri(ref.uri, context="ATIF trajectory") instead of constructing
a Path from urlparse(ref.uri).path, preserving the existing JSON parsing and
object validation.
In
`@plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/evaluator/harbor.py`:
- Around line 657-659: Update the trace_format annotations for _trial_resources
and the other helper at the referenced location to Literal["otlp", "atif"],
importing Literal as needed. In _run, avoid passing the Any result from
options_dict.pop directly; cast it to the literal type or reuse
options.trace_format so the narrowed contract remains type-safe.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1ef1ba8c-a85e-4df9-ab91-8ba275ef4a71
📒 Files selected for processing (12)
plugins/nemo-experimentalist/README.mdplugins/nemo-experimentalist/artifacts/.gitignoreplugins/nemo-experimentalist/examples/tau3-nooa-agent/harbor_wrapper.pyplugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.pyplugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/evaluator/harbor.pyplugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_explorer.pyplugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.pyplugins/nemo-experimentalist/tests/experimentalist/test_evaluator_harbor.pyplugins/nemo-experimentalist/tests/experimentalist/test_tau3_nooa_wrapper.pyplugins/nemo-experimentalist/tests/experimentalist/test_trace_explorer_from_ref.pyplugins/nemo-experimentalist/tests/test_atif.pyplugins/nemo-experimentalist/tests/test_experimentalist_backend.py
…path Two findings from review. The Intake parser hardcoded an empty tool_call_id for non-ATIF rows, but tool_call_id has no Intake column, so tool_call.id / tool.id / tool_call_id survive in raw_attributes. OTLP traces read back from Intake therefore lost their tool call ids, which the shared parser had preserved. Restores the attribute lookup and falls back to ATIF's preserved payload only for ATIF rows. The README documented trace_format under an optimizer key. That is the benchmark config's shape; an agent profile roots the same model at experiment_config, where evaluator is a top-level field. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Christian Schüller <cschueller@nvidia.com>
|
Reverts the ignore file added earlier in this branch. The directory it covered is transient benchmark output that a local run creates; with no such directory present there is nothing to ignore, and ruff check over the plugin passes on its own. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Christian Schüller <cschueller@nvidia.com>
Summary
The Experimentalist read agent traces as OTLP only, so an agent emitting ATIF produced trials with no trace at all — silently. Discovery found no
*.jsonl, the trial was skipped, and nothing said why. This teaches the trace path to carry its format from discovery through analysis, so an ATIF-emitting agent can be optimized. OTLP remains the default and its behavior is unchanged.Changes
ATIF trace support (
featcommit)_trial_resourcesrecognises*.atif.jsonalongside*.jsonland stampstrace_formaton theResourceRef, so the format is decided once and carried forward rather than re-derived downstream. NewHarborEvaluatorConfig.trace_formatselects which artifact becomes the trial's trace; both stay exposed inresources.experimentalist/atif.pyreads the trajectory and builds the ingest payload;_persist_trialdispatches on the stamped format and posts to Intake's ATIF ingest withevaluation_contextcarrying experiment and test-case identity. Trial identity rides on the trajectory's ownsession_id, which Intake adopts as the trace id.raw_attributes, so a reader that knows only the OpenInference names finds nothing. Token counts and error fields were missing on every span read back from storage, OTLP included.from_intakenow parses the columns directly;_load_trace_modelskeeps a per-record shape check because it deliberately accepts Intake-shaped payloads from disk. An ATIF branch keyed on thesourcecolumn covers fields that have no column at all.persist_evaluationlogs the trial it drops;from_refrejects a local.atif.jsonwith a message naming the real cause instead of reporting an empty trace.trace_formatsetting, the artifact contract, the Intake requirement, and the two diagnostics.Independent of ATIF, kept as separate commits so they can be reviewed or cherry-picked on their own:
fix(tau3-nooa-agent)— when the agent dies before writing traces,/app/tracesnever exists and the environment refuses to copy it, raisingRuntimeError. That escaped the metrics guard and replaced the agent's own error with a Docker copy failure. Catching it lets the run reach the existing check that raises with the agent's exit code and stderr.chore(experimentalist)— ignoreartifacts/experimentalist-benchmarks/. It is generated output, and because the optimizer writes candidate agents there, a candidate whose generated code does not compile was failingruff check plugins/nemo-experimentalist/.Review fixes (
fix(experimentalist)commit, from CodeRabbit)tool_call_idhas no Intake column, sotool_call.id/tool.id/tool_call_idsurvive inraw_attributes. The new Intake parser hardcoded an empty value for non-ATIF rows, so OTLP traces read back from Intake lost their tool call ids — a regression against the shared parser it replaced. Restores the attribute lookup, with ATIF's preserved payload as the fallback only for ATIF rows, and adds a regression test.trace_formatunder anoptimizerkey. That is the benchmark config's shape; an agent profile roots the same model atexperiment_config, whereevaluatoris a top-level field. Corrected.Type of Change
Quality Gates
31 tests added across four files:
test_atif.py(payload construction),test_evaluator_harbor.py(format discovery and selection),test_experimentalist_backend.py(ATIF upload),test_trace_explorer_from_ref.py(Intake-row parsing, both backends). Thefixcommit's 3 tests were confirmed to fail without the fix and pass with it.Verification
Signed-off-by:traileruv run pre-commit run -apasses, or any blocked checks are identified belowTargeted validation:
End-to-end verification against a live platform. Beyond the unit tests, the ATIF path was exercised with a real agent and a running Intake: trajectories were accepted by Intake's validators, 7 traces / 86 spans landed with correct
evaluation_idandtest_case_idgrouping, and spans read back with token counts and tool call ids populated. The agent-side scaffolding used for that run was reverted before this PR; only the feature remains.Known limitation. Per-step durations are absent on ATIF-derived spans: Intake takes a step's end time only from a producer-supplied invocation block, and infers nothing from adjacent steps. Trajectory-level duration and the Studio rollup are unaffected. This is documented in the README.