feat(profiler): ATIF-native profiler

[afourniernv]

NAT-366: ATIF-native profiler

Summary

How it works

┌─────────────────────────────────────────────────────────────────────────────┐
│                         ProfilerRunner.run(all_steps)                         │
└─────────────────────────────────────────────────────────────────────────────┘
                                        │
                                        ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│  INPUT (either format)                                                       │
│  ┌─────────────────────────────┐    ┌─────────────────────────────┐        │
│  │ list[list[IntermediateStep]] │    │     list[Trajectory]         │        │
│  │ (from nat eval workflow)     │    │     (ATIF, from file/API)    │        │
│  └─────────────────────────────┘    └─────────────────────────────┘        │
└─────────────────────────────────────────────────────────────────────────────┘
                                        │
                                        ▼
                    ┌───────────────────────────────────────┐
                    │  _normalize_to_trajectories()         │
                    │  IntermediateStep? → convert to ATIF  │
                    │  Trajectory? → pass through           │
                    └───────────────────────────────────────┘
                                        │
                                        ▼
                    ┌───────────────────────────────────────┐
                    │  list[Trajectory]  (unified format)   │
                    └───────────────────────────────────────┘
                                        │
                                        ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│  create_dataframe_from_atif(trajectories)                                    │
│  ┌─────────────────────────────────────────────────────────────────────┐    │
│  │  For each Trajectory step:                                           │    │
│  │  • user → WORKFLOW_START row                                        │    │
│  │  • agent + tool_calls + observation → TOOL_END row per tool          │    │
│  │  • agent + metrics (nonzero tokens) → LLM_START + LLM_END rows       │    │
│  │  • agent + message, no (tool_calls and observation) → WORKFLOW_END   │    │
│  │  Uses step.extra for function_id, parent_function_id, etc.            │    │
│  └─────────────────────────────────────────────────────────────────────┘    │
└─────────────────────────────────────────────────────────────────────────────┘
                                        │
                                        ▼
                    ┌───────────────────────────────────────┐
                    │  pd.DataFrame  (profiler schema)      │
                    └───────────────────────────────────────┘
                                        │
            ┌───────────────────────────┼───────────────────────────┐
            ▼                           ▼                           ▼
┌───────────────────┐   ┌───────────────────────────┐   ┌───────────────────┐
│ dataframe_to_      │   │  Downstream modules        │   │  JSON / CSV        │
│ profiler_traces()  │   │  • LLMMetrics              │   │  output            │
│ → traces_dict      │   │  • bottleneck_analysis     │   │                    │
│ → traces_obj       │   │  • token_uniqueness       │   │  all_requests_     │
│ (JSON, Trie)       │   │  • workflow_runtimes       │   │  profiler_traces   │
└───────────────────┘   │  • prompt_caching          │   │  standardized_     │
                        │  • ModelTrainer (forecast) │   │  data_all.csv      │
                        └───────────────────────────┘   └───────────────────┘

Stage	Input	Output
Boundary	IntermediateStep or Trajectory	list[Trajectory]
Conversion	list[Trajectory]	pd.DataFrame
Downstream	pd.DataFrame	Metrics, JSON, CSV, forecasts

Implements NAT-366: the profiler now accepts ATIF Trajectories and uses ATIF internally. Both IntermediateStep (existing) and Trajectory (ATIF) inputs are supported at the boundary.

Requirements (per JIRA / Anuradha doc)

• Input boundary: Profiler accepts list[list[IntermediateStep]] and list[Trajectory]. IntermediateStep input is converted to ATIF at the boundary when needed.
• Internal format: Profiler uses ATIF internally. create_dataframe_from_atif(trajectories) produces the same DataFrame schema as create_standardized_dataframe. Downstream modules consume DataFrame.
• step.extra: Uses step.extra for profiling metadata (function_id, parent_function_id, span_event_timestamp, nat_tool_ancestry for per-tool ancestry).
• Backward compatibility: IntermediateStep path remains supported.

Changes

New

• atif_dataframe.py: create_dataframe_from_atif(trajectories) — builds profiler DataFrame from ATIF trajectories. Maps user → WORKFLOW_START; agent with metrics → LLM_START/LLM_END; agent with tool_calls+observation → TOOL_END; agent message-only → WORKFLOW_END. Includes cached_tokens in total_tokens. dataframe_to_profiler_traces(df) for JSON output and PredictionTrieBuilder.

Modified

• profile_runner.py: Uses create_dataframe_from_atif; _normalize_to_trajectories() converts IntermediateStep → ATIF when needed.
• Downstream modules: Accept pd.DataFrame | list[list[IntermediateStep]]; use DataFrame path when available.
• RandomForestModel / ModelTrainer: Accept DataFrame from create_dataframe_from_atif for training.
• atif_converter.py: Minor updates for profiling metadata in step.extra.

Tests

• 6 new tests: TOOL_END path, cached_tokens, WORKFLOW_END message-only, dataframe_to_profiler_traces, empty DataFrame, ModelTrainer with DataFrame.
• 252 profiler tests pass.
• Smoke-tested via nat eval with simple calculator example (profiler enabled).

Testing

• Unit: uv run pytest packages/nvidia_nat_profiler/tests/ --ignore=.../test_profiler_e2e.py
• Smoke: nat eval --config_file examples/evaluation_and_profiling/simple_calculator_eval/configs/config-atif-smoke.yml

By Submitting this PR I confirm:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 946eb237-c4dd-4acf-9d8e-a62c64691a39

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[AnuradhaKaruppiah]

Thx @afourniernv. I like the guiding principles of this change.

Is there a way to look for drift in behavior before and after? For example for evaluators, I added a script to provide a summary of diff in scores when using IntermediateStep vs. AtifTrajectory - PR-1747

[AnuradhaKaruppiah]

With the changes in this #1747 you should be able to always pass ATIF Trajectory list to ProfilerRunner.run in EvaluationRun.profile_workflow.

That way we can drop the IntermediateStepToATIFConverter and have nvidia-nat-profiler only deal with ATIF (and drop all IntermediateStep references)