Skip to content

feat(experimentalist): support ATIF agent traces - #1127

Merged
schuellc-nvidia merged 6 commits into
mainfrom
experimentalist-atif-support/cschueller
Aug 7, 2026
Merged

schuellc-nvidia merged 6 commits into
mainfrom
experimentalist-atif-support/cschueller

Conversation

@schuellc-nvidia

@schuellc-nvidia schuellc-nvidia commented Aug 6, 2026 •

Copy link
Copy Markdown
Contributor

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 (feat commit)

  • Discovery — _trial_resources 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. New HarborEvaluatorConfig.trace_format selects which artifact becomes the trial's trace; both stay exposed in resources.
  • Upload — new experimentalist/atif.py reads the trajectory and builds the ingest payload; _persist_trial dispatches on the stamped format and posts 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.
  • Read-back — 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 parses 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; from_ref rejects a local .atif.json with a message naming the real cause instead of reporting an empty trace.
  • Docs — a README section covering the trace_format setting, 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/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. Catching it lets the run reach the existing check that raises with the agent's exit code and stderr.
  • chore(experimentalist) — ignore artifacts/experimentalist-benchmarks/. It is generated output, and because the optimizer writes candidate agents there, a candidate whose generated code does not compile was failing ruff check plugins/nemo-experimentalist/.

Review fixes (fix(experimentalist) commit, from CodeRabbit)

  • tool_call_id has no Intake column, so tool_call.id / tool.id / tool_call_id survive in raw_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.
  • 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. Corrected.

Type of Change

  • Code change (feature, bug fix, or refactor)
  • Code change with documentation updates
  • Documentation only
  • Contributor tooling or automation
  • CI, build, or test infrastructure

Quality Gates

  • Tests added or updated for changed behavior
  • Existing tests cover changed behavior — justification:
  • Tests not applicable — justification:
  • Documentation updated for user-visible behavior
  • Documentation not applicable — justification:

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). The fix commit's 3 tests were confirmed to fail without the fix and pass with it.

Verification

  • Pull request title follows the repository's Conventional Commit format
  • Every commit includes an appropriate Signed-off-by: trailer
  • uv run pre-commit run -a passes, or any blocked checks are identified below
  • Targeted tests pass, or tests are marked not applicable above
  • No secrets, API keys, or credentials are included

Targeted validation:

uv run pre-commit run -a
  → all hooks Passed (ruff, ruff format, ty typechecks, config reference,
    helm docs, uv lock, uv lock drift, copyright headers, UI lint-staged,
    plugin import boundary, merge conflicts)

uv run --frozen pytest <the 5 test files this PR touches> -q
  → 142 passed

uv run --frozen pytest plugins/nemo-experimentalist/tests/ -q
  → 631 passed

uv run ruff check plugins/nemo-experimentalist/
  → All checks passed

uv run ruff format --check
  → 3030 files already formatted

git diff --check origin/main...HEAD
  → clean

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_id and test_case_id grouping, 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.

schuellc-nvidia and others added 3 commits August 6, 2026 15:43
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>
@schuellc-nvidia
schuellc-nvidia requested review from a team as code owners August 6, 2026 13:53
@github-actions github-actions Bot added the feat label Aug 6, 2026
@coderabbitai

coderabbitai Bot commented Aug 6, 2026 •

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This 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.

Changes

ATIF trace support

Layer / File(s) Summary
Harbor trace selection
plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/evaluator/harbor.py, plugins/nemo-experimentalist/tests/experimentalist/test_evaluator_harbor.py, plugins/nemo-experimentalist/README.md
HarborEvaluatorConfig now supports otlp and atif, with otlp as the default. Trial discovery records both artifact types and selects the configured trace.
ATIF ingest path
plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py, plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.py, plugins/nemo-experimentalist/tests/test_atif.py, plugins/nemo-experimentalist/tests/test_experimentalist_backend.py, plugins/nemo-experimentalist/artifacts/.gitignore
ATIF trajectories are validated, enriched with evaluation context, and uploaded to the Intake ATIF endpoint as JSON. Uploads use the trajectory session_id as the trace ID.
Intake span parsing
plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_explorer.py, plugins/nemo-experimentalist/tests/experimentalist/test_trace_explorer_from_ref.py
TraceExplorer uses separate OTLP and Intake adapters. Intake rows produce semantic spans with ATIF messages, tool metadata, token counts, errors, and hierarchy fields. Local .atif.json references are rejected.
Wrapper trace failure handling
plugins/nemo-experimentalist/examples/tau3-nooa-agent/harbor_wrapper.py, plugins/nemo-experimentalist/tests/experimentalist/test_tau3_nooa_wrapper.py
The wrapper records trace download or parsing failures as unavailable trace metrics. Agent exit codes and stderr remain available for failure reporting.

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
Loading

Possibly related PRs

Suggested labels: docs

Suggested reviewers: aleckhoury, arpitsardhana, sandychapman

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 7.41% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the main change: adding ATIF agent trace support to Experimentalist.
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch experimentalist-atif-support/cschueller

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (3)
plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py (2)

36-48: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

The trajectory file is parsed twice per trial.

_persist_trial calls read_session_id, then _upload_trace_atif calls build_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 win

Use local_path_from_uri in _load. The helper in entities.py decodes escaped paths and rejects non-local schemes. Replace Path(urlparse(ref.uri).path) with local_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 value

Type trace_format as Literal["otlp", "atif"].

HarborEvaluatorConfig.trace_format is a Literal, but these helpers accept any str. An unexpected value silently selects OTLP. Narrowing the annotation makes the contract checkable by ty.

♻️ 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 _run reads trace_format from options_dict.pop(...), which is Any, so a cast or re-read from options.trace_format keeps 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

📥 Commits

Reviewing files that changed from the base of the PR and between a1b7051 and efb31b8.

📒 Files selected for processing (12)
  • plugins/nemo-experimentalist/README.md
  • plugins/nemo-experimentalist/artifacts/.gitignore
  • plugins/nemo-experimentalist/examples/tau3-nooa-agent/harbor_wrapper.py
  • plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/atif.py
  • plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/evaluator/harbor.py
  • plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/trace_explorer.py
  • plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.py
  • plugins/nemo-experimentalist/tests/experimentalist/test_evaluator_harbor.py
  • plugins/nemo-experimentalist/tests/experimentalist/test_tau3_nooa_wrapper.py
  • plugins/nemo-experimentalist/tests/experimentalist/test_trace_explorer_from_ref.py
  • plugins/nemo-experimentalist/tests/test_atif.py
  • plugins/nemo-experimentalist/tests/test_experimentalist_backend.py

Comment thread plugins/nemo-experimentalist/README.md
…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>
@github-actions

github-actions Bot commented Aug 6, 2026 •

Copy link
Copy Markdown
Contributor
Suite Lines Covered Line Rate Branch Rate
Unit Tests 31261/39870 78.4% 62.8%
Integration Tests 18203/37822 48.1% 20.6%

Comment thread plugins/nemo-experimentalist/artifacts/.gitignore Outdated
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>
@schuellc-nvidia
schuellc-nvidia added this pull request to the merge queue Aug 6, 2026
@schuellc-nvidia
schuellc-nvidia removed this pull request from the merge queue due to a manual request Aug 6, 2026
@schuellc-nvidia
schuellc-nvidia added this pull request to the merge queue Aug 6, 2026
Merged via the queue into main with commit c80fd17 Aug 7, 2026
125 of 137 checks passed
@schuellc-nvidia
schuellc-nvidia deleted the experimentalist-atif-support/cschueller branch August 7, 2026 00:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants