Skip to content

Add experimental nvidia-nat-app Agent Performance Primitives subpackage#1636

Merged
rapids-bot[bot] merged 15 commits intoNVIDIA:developfrom
mpenn:mpenn_nat_app
Feb 24, 2026
Merged

Add experimental nvidia-nat-app Agent Performance Primitives subpackage#1636
rapids-bot[bot] merged 15 commits intoNVIDIA:developfrom
mpenn:mpenn_nat_app

Conversation

@mpenn
Copy link
Contributor

@mpenn mpenn commented Feb 24, 2026

Description

Closes

Summary

Adds nvidia-nat-app, a framework-agnostic library of Agent Performance Primitives (APP) for accelerating agentic applications. This PR focuses provides graph analysis, parallelization, and speculative execution primitives that can be embedded in any agent framework.

Features

  • Graph optimizationquick_optimize() takes nodes and edges and returns parallel execution stages via a 6-stage pipeline: extract → validate → topology → node analysis → edge classification → scheduling.
  • AST-based static analysisanalyze_function() and classify_edge() infer state reads/writes and edge necessity without executing code.
  • Speculative executionplan_speculation(), speculative_opportunities(), and run_speculation() support router-branch speculation with @speculation_unsafe and SpeculationSafetyConfig.
  • Benchmarkingbenchmark() compares sequential vs. parallel execution and custom strategies.
  • Framework integrationAbstractFrameworkAdapter and build_graph_and_adapter() for framework-specific integration; public API uses plain Python types.
  • Constraints@sequential, @depends_on, @has_side_effects, and OptimizationConfig for controlling optimization.

API

quick_optimize | analyze_function | classify_edge | find_parallel_stages | benchmark | speculative_opportunities | plan_speculation | partition_targets | run_speculation

Notes

  • Framework-agnostic, no external dependencies. Python 3.11–3.13. Experimental (API may change).

Example

from nat_app import quick_optimize

stages = quick_optimize(
    nodes={"parse": parse_fn, "research_a": fn_a, "research_b": fn_b, "synthesize": fn_c},
    edges=[("parse", "research_a"), ("parse", "research_b"),
           ("research_a", "synthesize"), ("research_b", "synthesize")],
)
# → [{"parse"}, {"research_a", "research_b"}, {"synthesize"}]

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.

Summary by CodeRabbit

  • New Features

    • Add a complete framework‑agnostic optimization toolkit: graph model, static AST analysis, topology & scheduling, pipelined compiler, optimization stages, and high‑level APIs (quick_optimize, analyze_function, classify_edge, find_parallel_stages, benchmark, speculative_opportunities).
    • Speculation planning and execution: router‑branch strategy, resolution policies, planner, and a speculative runner; runtime utilities for execution state, metrics, result handling, constraints, and priority assignment.
  • Documentation

    • Package metadata and PyPI documentation for the NVIDIA NAT‑APP subpackage.
  • Tests

    • Extensive unit and integration tests covering compiler, graph, analysis, scheduling, speculation, stages, executors, and utilities.

…ich includes a framework-agnostic set of performance primitives for agentic AI applications.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
Signed-off-by: Matthew Penn <mpenn@nvidia.com>
@mpenn mpenn requested a review from dnandakumar-nv February 24, 2026 00:46
@coderabbitai
Copy link

coderabbitai bot commented Feb 24, 2026

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

Adds the NVIDIA nvidia-nat-app package: a complete graph optimization pipeline (graph model, AST static analysis, scheduling, constraints, speculation planning/strategies, executors, pipeline stages), a DefaultGraphCompiler and optimizer, packaging metadata, and an extensive test suite. No existing public APIs were removed.

Changes

Cohort / File(s) Summary
Packaging & docs
packages/nvidia_nat_app/pyproject.toml, packages/nvidia_nat_app/src/meta/pypi.md, pyproject.toml
New package metadata, PyPI description, and workspace pyproject updates for nvidia-nat-app.
Top-level API
packages/nvidia_nat_app/src/nat_app/__init__.py, packages/nvidia_nat_app/src/nat_app/api.py
Public package entrypoint emits ExperimentalWarning and re-exports public API; adds framework-agnostic helper APIs (quick_optimize, analyze_function, classify_edge, find_parallel_stages, benchmark, speculative_opportunities).
Compiler core
packages/nvidia_nat_app/src/nat_app/compiler/...
Adds pipelined compiler framework: CompilationContext, CompilationStage protocol, AbstractCompiler, AbstractPipelinedCompiler, DefaultGraphCompiler, errors, optimizer, context_to_result and compilation utilities.
Pipeline stages
packages/nvidia_nat_app/src/nat_app/stages/*
New pipeline stages (Extract, Validate, Topology, NodeAnalysis, LLMAnalysis, EdgeClassification, PriorityAssignment, Scheduling) and central re-export for stages.
Graph model & types
packages/nvidia_nat_app/src/nat_app/graph/types.py, .../models.py
Introduces Graph, Edge, NodeInfo, EdgeType, Graph result types, cost/prioritization types and graph factory helpers.
Topology & scheduling
packages/nvidia_nat_app/src/nat_app/graph/{topology,scheduling,analysis,static_analysis,...}
Adds cycle/router detection, topology analysis, node analysis, dependency graph builders, parallel-group discovery, edge classification, and compute_optimized_order scheduling logic (dense algorithmic changes).
Access & static analysis
packages/nvidia_nat_app/src/nat_app/graph/access.py, .../static_analysis.py, llm_detection.py
Adds AccessSet multi-object path tracking, comprehensive AST-based static function analysis, and LLM call discovery/counting utilities.
Constraints
packages/nvidia_nat_app/src/nat_app/constraints/*
Adds constraint decorators, NodeConstraints/ResolvedConstraints/OptimizationConfig dataclasses, and resolution/merge logic combining decorators, config, and heuristics.
Speculation framework
packages/nvidia_nat_app/src/nat_app/speculation/*
Adds SpeculationPlan, SpeculationPlanner, Resolution protocol, safety config/decorators, RouterBranchStrategy/Resolution and strategy plumbing (identify/plan/partitioning).
Executors & runtime
packages/nvidia_nat_app/src/nat_app/executors/*
Adds ExecutionState, ExecutionMetrics, ResultHandler, speculative runner (run_speculation/SpeculativeResult) and executor re-exports.
Tests
packages/nvidia_nat_app/tests/**
Extensive unit and integration tests added across compiler, stages, graph, topology, static analysis, scheduling, constraints, speculation, executors, priority assignment, LLM detection, and API end-to-end tests — many new test files.
CI script tweak
ci/scripts/github/build_wheel.sh
Conditional wheel validation path added to skip CLI nat checks for nvidia_nat_app wheels and instead import nat_app.

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant API
  participant Compiler
  participant StageExtract as "Extract (stage)"
  participant StageNodeAnalysis as "NodeAnalysis (stage)"
  participant StageEdgeClass as "EdgeClassification (stage)"
  participant StageScheduling as "Scheduling (stage)"
  participant Result

  Client->>API: quick_optimize(nodes, edges, ...)
  API->>Compiler: compile_to_result(graph_source)
  Compiler->>StageExtract: apply(context)
  StageExtract-->>Compiler: context [graph, reducer_fields, node_funcs]
  Compiler->>StageNodeAnalysis: apply(context)
  StageNodeAnalysis-->>Compiler: context [node_analyses, resolved_constraints]
  Compiler->>StageEdgeClass: apply(context)
  StageEdgeClass-->>Compiler: context [edge_analyses, parallel_groups]
  Compiler->>StageScheduling: apply(context)
  StageScheduling-->>Compiler: context [optimized_order]
  Compiler->>Result: context_to_result(context)
  Result-->>API: TransformationResult
  API-->>Client: list[set[str]] (parallel stages)
Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 36.80% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed The pull request title 'Add experimental nvidia-nat-app Agent Performance Primitives subpackage' is concise (71 chars), uses imperative mood ('Add'), and accurately describes the main change introducing the new subpackage.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
  • 📝 Generate docstrings (stacked PR)
  • 📝 Generate docstrings (commit on current branch)
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

…re proper functionality.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
@mpenn mpenn added feature request New feature or request non-breaking Non-breaking change labels Feb 24, 2026
- Renamed `default` module to `default_graph_compiler` for clarity.
- Updated all relevant imports across the codebase to reflect the new module name.
- Introduced `DefaultGraphCompiler` class with a standard graph optimization pipeline.
- Added tests for the new `DefaultGraphCompiler` to ensure functionality and correctness.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
…ctPipelinedCompiler` description. Removed unnecessary line breaks for better readability.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
Copy link
Contributor

@dnandakumar-nv dnandakumar-nv left a comment

Choose a reason for hiding this comment

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

very elegant

- Updated `compute_optimized_order` to treat missing nodes in `node_analyses` as opaque, logging a warning and assigning them structural dependencies.
- Improved cycle detection in `detect_cycles` to handle single-node cycles and self-loops more effectively.
- Added tests to verify the handling of missing nodes and self-loops in scheduling and cycle detection.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
@mpenn mpenn marked this pull request as ready for review February 24, 2026 07:21
@mpenn mpenn requested review from a team as code owners February 24, 2026 07:21
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 9

Note

Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.

🧹 Nitpick comments (23)
packages/nvidia_nat_app/tests/conftest.py (2)

26-35: Fixture does not follow the naming convention.

Per coding guidelines, pytest fixtures should define the name argument and use the fixture_ prefix on the decorated function.

Proposed fix
-@pytest.fixture(autouse=True)
-def _suppress_experimental_warning():
+@pytest.fixture(name="_suppress_experimental_warning", autouse=True)
+def fixture_suppress_experimental_warning():
     """Suppress the package-level experimental warning during tests."""

As per coding guidelines, "Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture function being decorated should be named using the fixture_ prefix."

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/conftest.py` around lines 26 - 35, The fixture
function _suppress_experimental_warning should follow project naming and
decorator conventions: rename the function to
fixture_suppress_experimental_warning and update the pytest.fixture decorator to
include the name argument (e.g.,
pytest.fixture(name="suppress_experimental_warning", autouse=True)); keep the
existing behavior that imports ExperimentalWarning from nat_app and uses
warnings.catch_warnings with warnings.filterwarnings("ignore",
category=ExperimentalWarning) and yield to preserve autouse semantics.

38-45: Add type hints to MinimalAdapter methods.

extract and build are missing type annotations. Even for test helpers, the guidelines require type hints on all method parameters and return values.

Proposed fix
 class MinimalAdapter(AbstractFrameworkAdapter):
     """Minimal concrete adapter for tests that need an AbstractFrameworkAdapter."""

-    def extract(self, source):
+    def extract(self, source: Any) -> Any:
         return source

-    def build(self, original, result):
+    def build(self, original: Any, result: Any) -> Any:
         return result

Add the import at the top:

+from typing import Any

As per coding guidelines, "Python methods should use type hints for all parameters and return values."

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/conftest.py` around lines 38 - 45,
MinimalAdapter's methods lack type annotations; update the extract and build
signatures to include explicit parameter and return types (e.g., def
extract(self, source: Any) -> Any and def build(self, original: Any, result:
Any) -> Any) and add the necessary import (from typing import Any) at the top of
the file so the test helper complies with the project's type-hinting guidelines
while keeping method names MinimalAdapter, extract, and build unchanged.
packages/nvidia_nat_app/pyproject.toml (1)

71-73: Remove or uncomment the commented-out pytest configuration.

Commented-out configuration is a debug artifact. If these options are needed, uncomment them; otherwise remove them to keep the file clean.

Proposed fix
-#[tool.pytest.ini_options]
-#pythonpath = ["."]
-#asyncio_mode = "auto"
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/pyproject.toml` around lines 71 - 73, The
pyproject.toml contains a commented-out pytest configuration block (the
#[tool.pytest.ini_options] section along with the pythonpath and asyncio_mode
lines); either remove these commented lines entirely or uncomment them so they
become active configuration (uncomment the #[tool.pytest.ini_options],
pythonpath = ["."], and asyncio_mode = "auto" lines) depending on whether pytest
needs those settings, ensuring the file no longer contains leftover
debug/commented config.
packages/nvidia_nat_app/src/nat_app/speculation/safety.py (1)

122-126: Consider using Sequence[str] for possible_targets.

The coding guidelines prefer collections.abc abstractions (Sequence over list) for public API surfaces. If consumers should not mutate this field after construction, Sequence[str] communicates immutability intent more clearly.

Proposed change
+from collections.abc import Sequence
 ...
-    possible_targets: list[str]
+    possible_targets: Sequence[str]

As per coding guidelines, "Prefer 'collections.abc' / 'typing' abstractions ('Sequence' over 'list')."

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/speculation/safety.py` around lines 122 -
126, Change the public annotation for the router field possible_targets from
concrete list[str] to the abstract Sequence[str] to signal non-mutability in the
API; update the import to bring in Sequence (from collections.abc or typing) and
adjust any construction sites or type checks that assume list methods (e.g.,
.append) to either copy to a mutable list internally or treat as a Sequence in
functions that consume it, referencing the possible_targets attribute on the
router node class/structure.
packages/nvidia_nat_app/tests/executors/test_execution_state.py (1)

22-177: Consider a shared ExecutionState fixture to reduce repetition.

ExecutionState() is recreated in nearly every test; a fixture will make the tests shorter and follow repo test conventions.

♻️ Example fixture + usage
+@pytest.fixture(name="execution_state")
+def fixture_execution_state():
+    return ExecutionState()
+
 class TestExecutionStateDefaults:
-    def test_fresh_state_has_zeroed_counters(self):
-        state = ExecutionState()
+    def test_fresh_state_has_zeroed_counters(self, execution_state):
+        state = execution_state
         assert state.tools_launched == 0
As per coding guidelines, “Any frequently repeated code should be extracted into pytest fixtures. Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture function being decorated should be named using the fixture_ prefix, using snake_case.”
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/executors/test_execution_state.py` around lines
22 - 177, Tests recreate ExecutionState() repeatedly; add a shared pytest
fixture to reduce duplication: define a function named fixture_execution_state
that returns ExecutionState() and decorate it with
`@pytest.fixture`(name="execution_state") (ensure pytest is imported), then update
tests to accept an execution_state parameter and replace local state =
ExecutionState() with using the execution_state object (references:
ExecutionState, fixture_execution_state, and tests like
TestExecutionStateDefaults::test_fresh_state_has_zeroed_counters).
packages/nvidia_nat_app/tests/executors/test_runner.py (1)

64-68: Extract repeated ExecutionState setup into a pytest fixture.

The pattern state = ExecutionState(); state.execution_start_time = 0.0 is duplicated in 11 of 12 async test methods. Per the coding guidelines, frequently repeated code should be extracted into a pytest fixture.

♻️ Suggested fixture
+    `@pytest.fixture`(name="execution_state")
+    def fixture_execution_state(self):
+        state = ExecutionState()
+        state.execution_start_time = 0.0
+        return state
+
     async def test_basic_left_chosen(self):
         plan = _make_plan()
-        state = ExecutionState()
-        state.execution_start_time = 0.0
+        state = execution_state  # inject via parameter

Then update each test method to accept execution_state as a parameter instead of constructing it inline.

As per coding guidelines: "Any frequently repeated code should be extracted into pytest fixtures."

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py` around lines 64 - 68,
Extract the repeated setup by adding a pytest fixture (e.g., name it
execution_state) that constructs an ExecutionState instance, sets
execution_state.execution_start_time = 0.0, and returns it; then update each
async test (the async def test_* functions that currently call ExecutionState()
and set execution_start_time) to accept execution_state as a parameter and
remove the inline construction and assignment. Locate uses of ExecutionState in
the test file (functions like test_basic_left_chosen and the other 10 async
tests) and replace their two-line setup with the new fixture parameter to avoid
duplication.
packages/nvidia_nat_app/src/nat_app/stages/edge_classification.py (1)

59-66: Use direct enum comparison instead of .value string comparison.

Replace ea.edge_type.value == "necessary" with ea.edge_type == EdgeType.NECESSARY for idiomatic enum handling and better type safety.

♻️ Suggested fix
+from nat_app.graph.models import EdgeType
+
     necessary: set[tuple[str, str]] = set()
     unnecessary: set[tuple[str, str]] = set()
     for ea in edge_analyses:
         edge = (ea.source, ea.target)
-        if ea.edge_type.value == "necessary":
+        if ea.edge_type == EdgeType.NECESSARY:
             necessary.add(edge)
-        elif ea.edge_type.value == "unnecessary":
+        elif ea.edge_type == EdgeType.UNNECESSARY:
             unnecessary.add(edge)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/stages/edge_classification.py` around
lines 59 - 66, The code compares enum values by string (ea.edge_type.value ==
"necessary") which is fragile; update the comparisons in the loop over
edge_analyses to use direct enum members instead (compare ea.edge_type to
EdgeType.NECESSARY and EdgeType.UNNECESSARY), leaving the rest of the logic
(adding to necessary and unnecessary sets) unchanged so type-safety and
idiomatic enum usage are preserved.
packages/nvidia_nat_app/tests/speculation/test_plan.py (1)

58-275: Extract common graph setups into pytest fixtures to reduce test duplication.

Multiple test methods repeat the same plan_speculation() calls with identical or similar nodes, edges, and conditional_edges arguments. Consolidate these into reusable fixtures named with the fixture_ prefix and set the fixture name= argument per coding guidelines.

Examples from the test suite:

  • Basic two-node graph: nodes={"router": route_fn, "a": fn_a, "b": fn_b} with edges=[("router", "a"), ("router", "b")] and matching conditional_edges
  • Three-node graph with merge: same router and branches plus "merge": fn_merge with additional edges to merge node

This refactoring will improve readability and maintainability across the test classes.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/speculation/test_plan.py` around lines 58 -
275, Extract repeated plan_speculation setups into pytest fixtures (prefixed
with fixture_) to DRY the tests: create a fixture for the basic
router-two-branches graph that returns the arguments (nodes with route_fn, fn_a,
fn_b; edges list; conditional_edges mapping) and a second fixture for the merge
graph that also includes "merge": fn_merge and the extra edges; register each
fixture with the name= argument (e.g., name="fixture_basic_graph",
"fixture_merge_graph") and update tests in TestPlanSpeculationBasic,
TestTargetsAndCancellation, and TestSafetyFiltering to call plan_speculation
using the fixture-provided values instead of repeating the same
nodes/edges/conditional_edges and safety setups (refer to plan_speculation,
route_fn, fn_a, fn_b, fn_merge, conditional_edges, and SpeculationSafetyConfig
to locate where to replace duplicated literals).
packages/nvidia_nat_app/src/nat_app/stages/priority_assignment.py (1)

320-327: Consider adding strict=True to zip() for extra safety.

group.node_names and group.priorities should be the same length; using strict=True will raise ValueError if they unexpectedly differ, catching potential bugs at runtime.

Suggested improvement
-            for node_name, priority_level in zip(group.node_names, group.priorities):
+            for node_name, priority_level in zip(group.node_names, group.priorities, strict=True):

Note: While Python 3.11 supports this parameter, it is not enforced by the project's ruff configuration.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/stages/priority_assignment.py` around
lines 320 - 327, The loop using zip(group.node_names, group.priorities) should
use zip(..., strict=True) to ensure the two iterables are the same length and
fail fast on mismatch; update the zip call in the priority assignment loop
(where node_name and priority_level are produced and used to set
node_info.priority, assigned, and node_assigned_priority) to
zip(group.node_names, group.priorities, strict=True) so any unexpected length
discrepancy raises ValueError during testing/runtime, while keeping the existing
checks for graph.has_node and node_info.priority.
packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py (1)

15-28: Docstring formatting: avoid leading blank line.

This module docstring starts with a blank first line; consider moving the summary onto the opening quotes for consistency with repo docstring rules.

As per coding guidelines: “The first line of docstrings must be a concise description ending with a period.”

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py` around lines 15 -
28, Update the module docstring at the top of the file so it does not start with
a blank line: move the concise summary sentence to the very first line inside
the triple quotes and ensure it ends with a period (e.g., "Generic LLM call
detection engine."). Keep the subsequent explanatory paragraphs (which mention
discover_llm_names and count_llm_calls) on following lines, preserving existing
content but removing the initial empty line so the docstring follows the
repository's one-line-summary rule.
packages/nvidia_nat_app/tests/stages/test_priority_assignment.py (3)

153-171: Formatting: _nested_conditional_parallel_graph has a stray comma in the signature.

This is minor but noisy (and likely to be reformatted by yapf anyway).

Proposed cleanup
-def _nested_conditional_parallel_graph(llm_analysis: dict[str, LLMCallInfo], ) -> CompilationContext:
+def _nested_conditional_parallel_graph(llm_analysis: dict[str, LLMCallInfo]) -> CompilationContext:
As per coding guidelines: “Run yapf… for code formatting.”
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/stages/test_priority_assignment.py` around
lines 153 - 171, The function signature for _nested_conditional_parallel_graph
has an extraneous trailing comma after the parameter list; remove the stray
comma from the signature of _nested_conditional_parallel_graph(llm_analysis:
dict[str, LLMCallInfo]) so the parameter list is properly formatted and then run
yapf to normalize formatting.

880-881: Consider removing the __main__ pytest.main() runner.

It’s fine under a guard, but it’s atypical in repo tests and can confuse IDE/test discovery tooling.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/stages/test_priority_assignment.py` around
lines 880 - 881, Remove the test-module entrypoint that calls pytest.main under
the if __name__ == "__main__" guard: delete the block containing if __name__ ==
"__main__": and the pytest.main([__file__, "-v"]) invocation so test discovery
and IDE tooling are not confused; keep the rest of test_priority_assignment.py
unchanged.

317-340: Redundant graph mutation in test_fast_branch_gets_high_slow_gets_low().

_router_graph() already adds the target nodes and conditional edges for the passed branches, so the extra add_node("medium_path") and add_conditional_edges(...) reads like leftover churn and makes the test harder to follow.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/stages/test_priority_assignment.py` around
lines 317 - 340, The test_fast_branch_gets_high_slow_gets_low test contains
redundant graph mutations after creating ctx with _router_graph: remove the
extra ctx.metadata["graph"].add_node("medium_path") and the subsequent
ctx.metadata["graph"].add_conditional_edges("router", {"a": "fast", "b":
"medium_path", "c": "slow"}) since _router_graph already adds those target nodes
and conditional edges; keep the rest of the test (PriorityAssignmentStage
instantiation, stage.apply(ctx), and the three asserts) unchanged.
packages/nvidia_nat_app/src/nat_app/graph/types.py (3)

429-458: Determinism: adjacency APIs return arbitrary ordering (set → list).

If any downstream logic or tests rely on stable iteration order, consider returning sorted(...) in successors(), predecessors(), and to_adjacency().

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/types.py` around lines 429 - 458,
The adjacency APIs return lists created from sets which produce arbitrary
ordering; update successors(), predecessors(), and to_adjacency() to return
deterministically ordered lists by applying sorted() to the underlying sets (use
sorted(self._successors.get(node, set())) and
sorted(self._predecessors.get(node, set())) in successors and predecessors, and
in to_adjacency() return {n: sorted(succs) for n, succs in
self._successors.items()}), keeping the return types as list[str] and preserving
the same method names and attributes (_successors, _predecessors).

15-21: Docstring formatting: avoid leading blank first line.

This module docstring starts with a blank first line; consider moving the summary onto the opening quotes.

As per coding guidelines: “The first line of docstrings must be a concise description ending with a period.”

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/types.py` around lines 15 - 21, The
module docstring in packages/nvidia_nat_app/src/nat_app/graph/types.py currently
has a leading blank line; update the module-level docstring so the summary
starts immediately after the opening triple quotes and is a concise single-line
description ending with a period (e.g., "First-class Graph type — the central
interchange format for all nat_app algorithms."), followed by any additional
paragraphs separated by a blank line.

407-426: Mutation safety: get_conditional_targets() / conditional_edge_sources expose internal lists/dicts.

Returning internal mutable structures makes it easy for callers to mutate graph state without updating _edges / _edge_keys / adjacency. If external mutation is not intended, return deep-ish copies (copy dict + copy lists).

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/types.py` around lines 407 - 426,
get_conditional_targets and conditional_edge_sources currently return references
to the internal _conditional_targets so callers can mutate internal lists/dicts;
change both to return defensive copies: have get_conditional_targets(node)
return None or a new dict mapping branch -> new list(copy) for that node, and
have the conditional_edge_sources property return a new dict where each branch
mapping is a shallow-copied dict and each target list is copied (i.e. copy dicts
and lists, not the originals). Reference _conditional_targets when building
these copies and ensure callers that need to mutate the graph update
_edges/_edge_keys/adjacency rather than mutating returned structures.
packages/nvidia_nat_app/src/nat_app/graph/topology.py (2)

519-569: Determinism: iterate routers in a stable order to make chain output stable.

router_set is a set, so chain construction order can vary across runs. If chain ordering is used in plan priority or tests, prefer sorted iteration.

Proposed deterministic iteration
-    for rnode in router_set:
+    for rnode in sorted(router_set):
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/topology.py` around lines 519 -
569, find_router_chains produces nondeterministic chain order because it
iterates a set; make iteration stable by using sorted order. Update the function
so key iterations over router_set and branch_lookup (where successor is built)
use sorted(...) to ensure deterministic traversal (e.g., iterate
sorted(router_set) when starting chains and iterate
sorted(branch_lookup.items()) when constructing successor), leaving all logic
intact so behavior is unchanged except ordering.

15-31: Docstrings: remove leading blank first line (and keep code identifiers backticked).

Several docstrings start with a blank line (""" on its own line), which violates the “first line is the summary” rule and tends to trip doc tooling. Consider moving the summary up to the opening quotes for the module + a few public functions.

Proposed docstring formatting fix
-"""
-Cycle and router detection for graph optimization.
+"""Cycle and router detection for graph optimization.
@@
-def cycle_node_order(cycle: CycleInfo, edges: list[tuple[str, str]]) -> list[str]:
-    """
-    Return nodes in a cycle in execution order (entry first, excluding back edge).
+def cycle_node_order(cycle: CycleInfo, edges: list[tuple[str, str]]) -> list[str]:
+    """Return nodes in a cycle in execution order (entry first, excluding back edge).
@@
-def detect_routers(graph: Graph) -> list[RouterInfo]:
-    """
-    Detect router nodes (nodes with conditional edges).
+def detect_routers(graph: Graph) -> list[RouterInfo]:
+    """Detect router nodes (nodes with conditional edges).
@@
-def get_safe_parallelization_groups(
+def get_safe_parallelization_groups(
@@
-) -> list[set[str]]:
-    """
-    Get groups of nodes that can safely be parallelized.
+) -> list[set[str]]:
+    """Get groups of nodes that can safely be parallelized.

As per coding guidelines: “Provide Google-style docstrings… The first line of docstrings must be a concise description ending with a period.” and “Surround code entities with backticks in docstrings…”.

Also applies to: 391-401, 424-432, 575-586

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/topology.py` around lines 15 - 31,
The module and several function/class docstrings start with a leading blank line
and lack backticked code identifiers; update the top-level module docstring so
the first line is a concise summary sentence (no blank line after the opening
triple quotes) and ensure all code entities are wrapped in backticks (e.g.,
`CycleInfo`, `Graph`, `Tarjan`/“Tarjan's SCC”) and follow Google-style docstring
rules; apply the same change to the other flagged docstrings (the ones
describing cycle detection, router detection and optimization boundaries near
the `CycleInfo`/cycle-detection logic and `Graph`-related functions) so each
docstring’s first line is the short summary ending with a period and code names
are backticked.
packages/nvidia_nat_app/src/nat_app/api.py (1)

15-36: Naming/doc clarity: confirm intended public naming (nat_app / nvidia-nat-app) and fix a couple doc typos.

  • This module positions itself as embedding nvidia-nat-app, but repo guidelines prefer nvidia-nat and nat as namespace—please confirm the intentional exception for this experimental package.
  • quick_optimize() doc: conditional_edges description reads like {target_name: target_name}; should be {branch_name: target_node(s)}.
  • speculative_opportunities() doc: the branches bullet says {target: [...]} but code returns {label: [...]}.

Based on learnings: “Use 'nat' for the API namespace and CLI tool, 'nvidia-nat' for the package name…”.

Also applies to: 65-83, 407-429

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/api.py` around lines 15 - 36, Update the
module docstring and function docs to follow the repo naming guidelines (use
"nat" for the public API/CLI and "nvidia-nat" for the package name) and correct
the inaccurate parameter descriptions: in nat_app.api's module docstring and
quick_optimize() docstring adjust the package/namespace wording to "nat"
(API/CLI) vs "nvidia-nat" (package) and change quick_optimize's
conditional_edges description from the confusing "{target_name: target_name}" to
"{branch_name: target_node_or_nodes}" (i.e., branch label maps to target
node(s)); similarly update speculative_opportunities() doc to state that the
returned mapping keys are branch labels (e.g., "{label: [...]}") not targets,
and ensure DefaultGraphCompiler mention remains consistent with the naming
convention.
packages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.py (4)

181-260: Planning recomputes topology and trusts opp.metadata shape; consider tightening.

plan() recomputes topology even though identify() already received it, and it assumes opp.metadata["branches"] / ["merge_nodes"] exist with the right types. A small guard (or a typed metadata object) would make the strategy more robust to mixed-opportunity inputs.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.py`
around lines 181 - 260, plan() recomputes the graph topology and blindly trusts
opp.metadata["branches"] and ["merge_nodes"], which can break if identify()
already provided topology or metadata is malformed; modify plan() to prefer a
topology supplied by the opportunity (or accept topology as an extra argument
from identify()) instead of calling analyze_graph_topology(graph)
unconditionally, and add guards that validate opp.metadata keys and types
(ensure "branches" is a dict[str, set[str]] and "merge_nodes" is a set[str],
defaulting to {} / set() when missing or wrong type), plus null-safe lookups for
router (router_lookup.get(rnode)) and cond_targets so malformed opportunities
are skipped or handled safely.

42-110: RouterBranchResolution logic is clear; consider an explicit “unknown decision” policy.

Right now an unexpected decision_result defaults to “keep everything” (empty cancel set). If that is intentional, a short comment or explicit branch would help avoid ambiguous semantics, since this affects cancellation correctness.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.py`
around lines 42 - 110, The current behavior silently treats unknown decision
labels as "keep everything"; make this explicit by adding a clear
unknown-decision policy in RouterBranchResolution: update resolve(),
get_cancel_set(), and is_on_chosen_path() to check whether the resolved label
exists in cancel_map and, if not, apply a defined policy (e.g., treat as "cancel
all" by setting cancel = self.all_targets or alternatively raise a ValueError),
and document that choice in the class docstring; ensure you reference
RouterBranchResolution._resolve_label, RouterBranchResolution.resolve,
RouterBranchResolution.get_cancel_set, and
RouterBranchResolution.is_on_chosen_path when implementing the branch so the
behavior is unambiguous.

15-21: Docstring formatting nit: avoid leading blank line; backtick nat_app.

This module docstring starts with a blank first line and mentions nat_app without backticks. (Small, but aligns with repo docstring rules.)

As per coding guidelines: “The first line of docstrings must be a concise description ending with a period.” and “Surround code entities with backticks in docstrings…”.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.py`
around lines 15 - 21, Edit the module docstring at the top of router_branch.py
to remove the leading blank line, ensure the first line is a concise description
that ends with a period (e.g., "Router-branch speculation strategy."), and
surround the code entity nat_app with backticks (e.g., `nat_app`) in the rest of
the docstring; update the existing multi-line string so it follows the repo
docstring rules without changing behavior of any functions or classes in the
file.

112-118: Use TYPE_CHECKING instead of import-time assert isinstance() for protocol conformance checks.

While ResolutionPolicy is marked @runtime_checkable so the assertion will succeed at runtime, import-time assertions:

  • disappear under python -O (optimization flag),
  • create unnecessary side effects by instantiating objects at module load time.

A cleaner approach is to use a static-only check with TYPE_CHECKING:

Proposed static protocol conformance check
+from typing import TYPE_CHECKING
@@
-# -- Satisfy the ResolutionPolicy protocol check at import time --
-assert isinstance(RouterBranchResolution(
-    cancel_map={},
-    label_map=None,
-    all_targets=frozenset(),
-), ResolutionPolicy)
+if TYPE_CHECKING:
+    _: ResolutionPolicy = RouterBranchResolution(
+        cancel_map={},
+        label_map=None,
+        all_targets=frozenset(),
+    )
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.py`
around lines 112 - 118, Replace the import-time runtime assertion with a
static-only type check using typing.TYPE_CHECKING: remove the assert
isinstance(...) call and add a TYPE_CHECKING block that assigns a
RouterBranchResolution(...) instance to a variable annotated as ResolutionPolicy
(e.g., _check: ResolutionPolicy = RouterBranchResolution(cancel_map={},
label_map=None, all_targets=frozenset())), importing TYPE_CHECKING and
ResolutionPolicy for the type-check block so the check is only performed by
static type checkers and not at runtime.

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between bd6f4cf and dbfa73a.

⛔ Files ignored due to path filters (1)
  • packages/nvidia_nat_app/uv.lock is excluded by !**/*.lock
📒 Files selected for processing (94)
  • packages/nvidia_nat_app/pyproject.toml
  • packages/nvidia_nat_app/src/meta/pypi.md
  • packages/nvidia_nat_app/src/nat_app/__init__.py
  • packages/nvidia_nat_app/src/nat_app/api.py
  • packages/nvidia_nat_app/src/nat_app/compiler/__init__.py
  • packages/nvidia_nat_app/src/nat_app/compiler/compilation_context.py
  • packages/nvidia_nat_app/src/nat_app/compiler/compilation_stage.py
  • packages/nvidia_nat_app/src/nat_app/compiler/compiler.py
  • packages/nvidia_nat_app/src/nat_app/compiler/default_graph_compiler.py
  • packages/nvidia_nat_app/src/nat_app/compiler/errors.py
  • packages/nvidia_nat_app/src/nat_app/compiler/optimizer.py
  • packages/nvidia_nat_app/src/nat_app/compiler/pipelined_compiler.py
  • packages/nvidia_nat_app/src/nat_app/constraints/__init__.py
  • packages/nvidia_nat_app/src/nat_app/constraints/decorators.py
  • packages/nvidia_nat_app/src/nat_app/constraints/models.py
  • packages/nvidia_nat_app/src/nat_app/constraints/resolution.py
  • packages/nvidia_nat_app/src/nat_app/executors/__init__.py
  • packages/nvidia_nat_app/src/nat_app/executors/execution_state.py
  • packages/nvidia_nat_app/src/nat_app/executors/metrics.py
  • packages/nvidia_nat_app/src/nat_app/executors/result_handler.py
  • packages/nvidia_nat_app/src/nat_app/executors/runner.py
  • packages/nvidia_nat_app/src/nat_app/graph/__init__.py
  • packages/nvidia_nat_app/src/nat_app/graph/access.py
  • packages/nvidia_nat_app/src/nat_app/graph/adapter.py
  • packages/nvidia_nat_app/src/nat_app/graph/analysis.py
  • packages/nvidia_nat_app/src/nat_app/graph/factory.py
  • packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py
  • packages/nvidia_nat_app/src/nat_app/graph/models.py
  • packages/nvidia_nat_app/src/nat_app/graph/protocols.py
  • packages/nvidia_nat_app/src/nat_app/graph/scheduling.py
  • packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py
  • packages/nvidia_nat_app/src/nat_app/graph/topology.py
  • packages/nvidia_nat_app/src/nat_app/graph/types.py
  • packages/nvidia_nat_app/src/nat_app/speculation/__init__.py
  • packages/nvidia_nat_app/src/nat_app/speculation/plan.py
  • packages/nvidia_nat_app/src/nat_app/speculation/planner.py
  • packages/nvidia_nat_app/src/nat_app/speculation/resolution.py
  • packages/nvidia_nat_app/src/nat_app/speculation/safety.py
  • packages/nvidia_nat_app/src/nat_app/speculation/strategies/__init__.py
  • packages/nvidia_nat_app/src/nat_app/speculation/strategies/base.py
  • packages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.py
  • packages/nvidia_nat_app/src/nat_app/stages/__init__.py
  • packages/nvidia_nat_app/src/nat_app/stages/edge_classification.py
  • packages/nvidia_nat_app/src/nat_app/stages/extract.py
  • packages/nvidia_nat_app/src/nat_app/stages/llm_analysis.py
  • packages/nvidia_nat_app/src/nat_app/stages/node_analysis.py
  • packages/nvidia_nat_app/src/nat_app/stages/priority_assignment.py
  • packages/nvidia_nat_app/src/nat_app/stages/scheduling.py
  • packages/nvidia_nat_app/src/nat_app/stages/topology.py
  • packages/nvidia_nat_app/src/nat_app/stages/validate.py
  • packages/nvidia_nat_app/tests/__init__.py
  • packages/nvidia_nat_app/tests/compiler/__init__.py
  • packages/nvidia_nat_app/tests/compiler/test_compilation_context.py
  • packages/nvidia_nat_app/tests/compiler/test_compilation_stage.py
  • packages/nvidia_nat_app/tests/compiler/test_compiler.py
  • packages/nvidia_nat_app/tests/compiler/test_default_graph_compiler.py
  • packages/nvidia_nat_app/tests/compiler/test_pipelined_compiler.py
  • packages/nvidia_nat_app/tests/conftest.py
  • packages/nvidia_nat_app/tests/constraints/__init__.py
  • packages/nvidia_nat_app/tests/constraints/test_decorators.py
  • packages/nvidia_nat_app/tests/constraints/test_models.py
  • packages/nvidia_nat_app/tests/constraints/test_resolution.py
  • packages/nvidia_nat_app/tests/executors/__init__.py
  • packages/nvidia_nat_app/tests/executors/test_execution_state.py
  • packages/nvidia_nat_app/tests/executors/test_metrics.py
  • packages/nvidia_nat_app/tests/executors/test_result_handler.py
  • packages/nvidia_nat_app/tests/executors/test_runner.py
  • packages/nvidia_nat_app/tests/graph/__init__.py
  • packages/nvidia_nat_app/tests/graph/conftest.py
  • packages/nvidia_nat_app/tests/graph/test_access.py
  • packages/nvidia_nat_app/tests/graph/test_adapter.py
  • packages/nvidia_nat_app/tests/graph/test_analysis.py
  • packages/nvidia_nat_app/tests/graph/test_factory.py
  • packages/nvidia_nat_app/tests/graph/test_llm_detection.py
  • packages/nvidia_nat_app/tests/graph/test_models.py
  • packages/nvidia_nat_app/tests/graph/test_optimizer.py
  • packages/nvidia_nat_app/tests/graph/test_scheduling.py
  • packages/nvidia_nat_app/tests/graph/test_static_analysis.py
  • packages/nvidia_nat_app/tests/graph/test_topology.py
  • packages/nvidia_nat_app/tests/graph/test_types.py
  • packages/nvidia_nat_app/tests/graph/test_uncertainty_invariants.py
  • packages/nvidia_nat_app/tests/speculation/__init__.py
  • packages/nvidia_nat_app/tests/speculation/test_plan.py
  • packages/nvidia_nat_app/tests/speculation/test_safety.py
  • packages/nvidia_nat_app/tests/stages/__init__.py
  • packages/nvidia_nat_app/tests/stages/test_edge_classification.py
  • packages/nvidia_nat_app/tests/stages/test_extract.py
  • packages/nvidia_nat_app/tests/stages/test_llm_analysis.py
  • packages/nvidia_nat_app/tests/stages/test_node_analysis.py
  • packages/nvidia_nat_app/tests/stages/test_priority_assignment.py
  • packages/nvidia_nat_app/tests/stages/test_scheduling.py
  • packages/nvidia_nat_app/tests/stages/test_topology.py
  • packages/nvidia_nat_app/tests/stages/test_validate.py
  • packages/nvidia_nat_app/tests/test_api.py

- Updated `analyze_function` to correctly reference `result.all_writes` for write operations.
- Improved error handling in `run_speculation` to ensure target tasks are cancelled upon exceptions.
- Enhanced group merging logic in `find_parallel_groups` to respect node dependencies, preventing transitive merges.
- Added support for conditional edges in the `Graph` class, including branch handling and removal of old edges.
- Updated tests to cover new functionality, including decision task error handling and conditional edge behavior.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 6

♻️ Duplicate comments (1)
packages/nvidia_nat_app/src/nat_app/api.py (1)

237-237: Unused adapter variable triggers Ruff RUF059.

Both find_parallel_stages (Line 237) and speculative_opportunities (Line 440) unpack adapter but never use it.

Proposed fix
-    graph, adapter = build_graph_and_adapter(nodes, edges, self_state_attrs=self_state_attrs)
+    graph, _adapter = build_graph_and_adapter(nodes, edges, self_state_attrs=self_state_attrs)
-    graph, adapter = build_graph_and_adapter(
+    graph, _adapter = build_graph_and_adapter(
         nodes, edges, conditional_edges=conditional_edges,
         self_state_attrs=self_state_attrs,
     )

Also applies to: 440-443

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/api.py` at line 237, The unpacked but
unused adapter from build_graph_and_adapter inside find_parallel_stages and
speculative_opportunities is triggering RUF059; change the unpack to ignore the
second value (e.g., graph, _ = build_graph_and_adapter(...) or graph, _adapter =
build_graph_and_adapter(...)) so only graph is used and the linter warning is
suppressed, leaving build_graph_and_adapter calls and semantics unchanged.
🧹 Nitpick comments (3)
packages/nvidia_nat_app/tests/test_api.py (1)

231-240: Unused loop variable name — rename to _name.

Ruff B007 flags name as unused inside the loop body.

Proposed fix
-        for name, analysis in info.items():
+        for _name, analysis in info.items():
             assert "reads" in analysis
             assert "writes" in analysis
             assert "confidence" in analysis
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/test_api.py` around lines 231 - 240, In
test_info_has_reads_writes, the loop uses an unused variable name which triggers
Ruff B007; rename the loop variable from name to _name in the for loop that
iterates over info.items() (the test_info_has_reads_writes function surrounding
the call to find_parallel_stages) so the loop becomes for _name, analysis in
info.items(): while leaving the assertions on analysis unchanged.
packages/nvidia_nat_app/src/nat_app/api.py (1)

53-55: Verify re-exported symbols are part of the intended public API.

SpeculationPlan, partition_targets, and plan_speculation are imported with # noqa: F401 but are not used within this module. If these are intentional re-exports, consider adding them to an __all__ list in this module to make the intent explicit and prevent accidental removal.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/api.py` around lines 53 - 55, The three
imported-but-unused symbols (SpeculationPlan, partition_targets,
plan_speculation) are likely intended as re-exports; make that explicit by
adding them to this module's public API (e.g., append "SpeculationPlan",
"partition_targets", and "plan_speculation" to an __all__ list) and keep the
imports so they remain available; optionally remove or adjust the `# noqa: F401`
comments once __all__ documents the intent.
packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py (1)

339-346: Conservative has_dynamic_targets flag on resolved-but-unknown receivers.

When _resolve_receiver returns a concrete name that isn't in _llm_names (Line 343–344), the counter flags has_dynamic_targets = True. This is overly cautious: a call like logger.invoke("msg") would trigger partial confidence just because .invoke matches an invocation method name. Consider only flagging when the receiver is None (truly unresolvable) and trusting the resolved-name-not-in-llm-names case as a normal non-LLM call.

Proposed adjustment
         name = self._resolve_receiver(node.value)
         if name is None:
             self.has_dynamic_targets = True
             return False
-        if name not in self._llm_names:
-            self.has_dynamic_targets = True
-            return False
-        return True
+        return name in self._llm_names
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py` around lines 339
- 346, The code currently treats any resolved receiver name not found in
_llm_names as dynamic by setting has_dynamic_targets = True; change this so
has_dynamic_targets is only set when _resolve_receiver(...) returns None (truly
unresolvable). Specifically, in the block using name =
self._resolve_receiver(node.value) keep the early branch that sets
has_dynamic_targets = True and returns False only when name is None, but when
name is concrete and not in self._llm_names simply return False without flipping
has_dynamic_targets; update the logic in the method containing
_resolve_receiver, has_dynamic_targets, and _llm_names accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@packages/nvidia_nat_app/src/nat_app/api.py`:
- Around line 326-342: _run_parallel silently drops non-dict results from
execute_node while _run_sequential assigns them directly; make the return
contract consistent by validating execute_node's return value in both runners
and raising a clear error if it's not a dict. Specifically, in async functions
_run_sequential and _run_parallel, after each await execute_node(...) check
isinstance(result, dict) and if not raise a TypeError (or ValueError) that names
execute_node and the node name, instructing callers that execute_node must
return a dict; alternatively add a clear docstring above execute_node/_run_*
stating the dict contract and keep the validation in at least one runner to
fail-fast.

In `@packages/nvidia_nat_app/src/nat_app/executors/runner.py`:
- Around line 166-201: The loops over actually_cancelled/target_tasks (symbols:
actually_cancelled, target_tasks, execution_state.record_timeline_event,
chosen_results) currently swallow all exceptions; change them to handle
asyncio.CancelledError separately and use logger.exception(...) when catching
other exceptions so full stack traces are preserved; specifically, in the
cancellation branch await the task inside a try/except that excepts
asyncio.CancelledError (ignore) and catches Exception to call
logger.exception(...) and update execution_state (e.g., increment
tools_cancelled or mark failure) instead of pass, and in the collection branch
replace the bare except that does logger.debug(...) with an except
asyncio.CancelledError (ignore or handle as cancelled) and an except Exception
that calls logger.exception(...) and records the failure so retries/metrics are
accurate.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py`:
- Around line 851-904: The code calls next(iter(param_to_obj.values())) without
guarding for an empty param_to_obj in both the function-def branch (assigning
callee_param_to_obj at state_param) and the lambda branch (assigned to
resolved_param), which raises StopIteration when param_to_obj is empty; update
both sites to safely obtain the value (e.g., use next(iter(...), None) or check
if param_to_obj) and handle the None case by skipping analysis/returning None or
otherwise avoiding constructing callee_param_to_obj and invoking
_NodeASTVisitor/_extract_writes_from_expr when no object is available so the
analyzer doesn't crash when self_state_attrs passes an empty param_to_obj.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py`:
- Around line 59-62: Rename the unused parameter in the
_RaisingResolution.resolve method from label to _label to silence the ARG002
ruff warning; update the method signature in class _RaisingResolution
(resolve(self, _label: str)) so the parameter is explicitly marked unused and no
other behavior changes are required.
- Around line 81-377: Replace unused lambda parameter names "r" with "_" for all
get_decision call sites to silence ARG005; locate each run_speculation
invocation that passes get_decision=lambda r: ... (e.g., in tests like
test_basic_left_chosen, test_basic_right_chosen, test_metrics_tracked,
test_router_decision_recorded, test_timeline_events_recorded,
test_fast_target_already_done_before_router, test_non_dict_results,
test_cancel_map_references_node_not_in_targets,
test_cancelled_task_already_done_before_cancel,
test_chosen_target_raises_continues_without_crash,
test_get_decision_raises_cancels_targets,
test_resolution_resolve_raises_cancels_targets, etc.) and change the lambda
signature to lambda _: ... so the unused argument is clearly intentional.

In `@packages/nvidia_nat_app/tests/test_api.py`:
- Around line 118-120: The test_writes_detected contains a tautological
assertion: in test_writes_detected replace the redundant check `"ticker" in
info["writes"] or "ticker" in info["writes"]` with the intended assertion —
either remove the unnecessary `or` to simply assert `"ticker" in
info["writes"]`, or if a second key was intended (e.g., `"thesis"`), change the
second operand to the correct key so the test asserts `"ticker" in
info["writes"] or "thesis" in info["writes"]`; locate this inside the
test_writes_detected function which calls analyze_function(step_a) and asserts
against info["writes"].

---

Duplicate comments:
In `@packages/nvidia_nat_app/src/nat_app/api.py`:
- Line 237: The unpacked but unused adapter from build_graph_and_adapter inside
find_parallel_stages and speculative_opportunities is triggering RUF059; change
the unpack to ignore the second value (e.g., graph, _ =
build_graph_and_adapter(...) or graph, _adapter = build_graph_and_adapter(...))
so only graph is used and the linter warning is suppressed, leaving
build_graph_and_adapter calls and semantics unchanged.

---

Nitpick comments:
In `@packages/nvidia_nat_app/src/nat_app/api.py`:
- Around line 53-55: The three imported-but-unused symbols (SpeculationPlan,
partition_targets, plan_speculation) are likely intended as re-exports; make
that explicit by adding them to this module's public API (e.g., append
"SpeculationPlan", "partition_targets", and "plan_speculation" to an __all__
list) and keep the imports so they remain available; optionally remove or adjust
the `# noqa: F401` comments once __all__ documents the intent.

In `@packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py`:
- Around line 339-346: The code currently treats any resolved receiver name not
found in _llm_names as dynamic by setting has_dynamic_targets = True; change
this so has_dynamic_targets is only set when _resolve_receiver(...) returns None
(truly unresolvable). Specifically, in the block using name =
self._resolve_receiver(node.value) keep the early branch that sets
has_dynamic_targets = True and returns False only when name is None, but when
name is concrete and not in self._llm_names simply return False without flipping
has_dynamic_targets; update the logic in the method containing
_resolve_receiver, has_dynamic_targets, and _llm_names accordingly.

In `@packages/nvidia_nat_app/tests/test_api.py`:
- Around line 231-240: In test_info_has_reads_writes, the loop uses an unused
variable name which triggers Ruff B007; rename the loop variable from name to
_name in the for loop that iterates over info.items() (the
test_info_has_reads_writes function surrounding the call to
find_parallel_stages) so the loop becomes for _name, analysis in info.items():
while leaving the assertions on analysis unchanged.

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between dbfa73a and 8b9efa6.

📒 Files selected for processing (13)
  • packages/nvidia_nat_app/src/nat_app/api.py
  • packages/nvidia_nat_app/src/nat_app/executors/runner.py
  • packages/nvidia_nat_app/src/nat_app/graph/analysis.py
  • packages/nvidia_nat_app/src/nat_app/graph/llm_detection.py
  • packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py
  • packages/nvidia_nat_app/src/nat_app/graph/topology.py
  • packages/nvidia_nat_app/src/nat_app/graph/types.py
  • packages/nvidia_nat_app/tests/executors/test_runner.py
  • packages/nvidia_nat_app/tests/graph/test_analysis.py
  • packages/nvidia_nat_app/tests/graph/test_llm_detection.py
  • packages/nvidia_nat_app/tests/graph/test_topology.py
  • packages/nvidia_nat_app/tests/graph/test_types.py
  • packages/nvidia_nat_app/tests/test_api.py

- Updated `benchmark` function to ensure `execute_node` returns a dict, raising a TypeError for invalid return types.
- Improved error handling in `run_speculation` to log exceptions during task cancellation and after task completion.
- Refactored `_analyze_callee` to use a default object when retrieving parameters, enhancing robustness.
- Added tests to validate the new error handling for non-dict returns in sequential and parallel executions.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 3

♻️ Duplicate comments (1)
packages/nvidia_nat_app/src/nat_app/api.py (1)

98-98: 🛠️ Refactor suggestion | 🟠 Major

Prefix unused adapter with underscore to silence RUF059.

All three build_graph_and_adapter call sites unpack adapter but never use it. The past review flagged lines 237 and 445; line 98 has the same issue.

Proposed fix (all three sites)
-    graph, adapter = build_graph_and_adapter(nodes, edges, entry, conditional_edges, self_state_attrs)
+    graph, _adapter = build_graph_and_adapter(nodes, edges, entry, conditional_edges, self_state_attrs)
-    graph, adapter = build_graph_and_adapter(nodes, edges, self_state_attrs=self_state_attrs)
+    graph, _adapter = build_graph_and_adapter(nodes, edges, self_state_attrs=self_state_attrs)
-    graph, adapter = build_graph_and_adapter(
+    graph, _adapter = build_graph_and_adapter(
         nodes, edges, conditional_edges=conditional_edges,
         self_state_attrs=self_state_attrs,
     )

As per coding guidelines: "Use ruff (via ruff check --fix) for linting… fix warnings unless explicitly ignored in pyproject.toml."

Also applies to: 237-237, 445-448

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/api.py` at line 98, The call to
build_graph_and_adapter unpacks an unused variable named adapter which triggers
RUF059; change the unpacked name to a prefixed underscore (e.g., _adapter or _)
at this call site (and the other two sites referenced) so the unused value is
explicit and linter-suppressed, leaving graph as-is (i.e., replace "graph,
adapter = build_graph_and_adapter(...)" with "graph, _adapter =
build_graph_and_adapter(...)" or similar).
🧹 Nitpick comments (5)
packages/nvidia_nat_app/tests/test_api.py (2)

231-240: Rename unused loop variable name to _name (B007).

Proposed fix
-        for name, analysis in info.items():
+        for _name, analysis in info.items():
             assert "reads" in analysis
             assert "writes" in analysis
             assert "confidence" in analysis
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/test_api.py` around lines 231 - 240, The loop
in test_info_has_reads_writes uses an unused variable name "name" which should
be renamed to "_name" to satisfy the B007 lint rule; update the for statement in
the test_info_has_reads_writes function (for name, analysis in info.items()) to
use for _name, analysis in info.items() so only the used variable analysis
remains named.

384-414: Use raw strings for regex patterns in match= to silence RUF043.

The .* metacharacters in the match= parameter work correctly but ruff flags them (RUF043) because they aren't in raw strings, which could confuse readers about intent.

Proposed fix
-        with pytest.raises(TypeError, match="execute_node must return a dict.*got NoneType.*node"):
+        with pytest.raises(TypeError, match=r"execute_node must return a dict.*got NoneType.*node"):
-        with pytest.raises(TypeError, match="execute_node must return a dict.*got str.*node"):
+        with pytest.raises(TypeError, match=r"execute_node must return a dict.*got str.*node"):
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/test_api.py` around lines 384 - 414, Update the
pytest.raises calls in the tests to use raw-string regex literals to satisfy
RUF043: change the match arguments in
test_execute_node_non_dict_sequential_raises and
test_execute_node_non_dict_parallel_raises from plain strings like "execute_node
must return a dict.*got NoneType.*node" and "execute_node must return a
dict.*got str.*node" to raw strings r"execute_node must return a dict.*got
NoneType.*node" and r"execute_node must return a dict.*got str.*node"
respectively; no other behavior changes needed (tests still use the same
execute_node helpers and benchmark call).
packages/nvidia_nat_app/tests/executors/test_runner.py (1)

70-77: Extract repeated ExecutionState setup into a pytest fixture.

The pattern state = ExecutionState(); state.execution_start_time = 0.0 is repeated in every test method (~14 times). Per coding guidelines, frequently repeated code should be extracted into a fixture.

Proposed fixture
`@pytest.fixture`(name="execution_state")
def fixture_execution_state() -> ExecutionState:
    state = ExecutionState()
    state.execution_start_time = 0.0
    return state

Then in each test:

-    async def test_basic_left_chosen(self):
-        plan = _make_plan()
-        state = ExecutionState()
-        state.execution_start_time = 0.0
+    async def test_basic_left_chosen(self, execution_state):
+        plan = _make_plan()

As per coding guidelines: "Any frequently repeated code should be extracted into pytest fixtures. Pytest fixtures should define the name argument when applying the pytest.fixture decorator. The fixture function being decorated should be named using the fixture_ prefix."

Also applies to: 95-99, 112-115, 131-134, 146-149, 166-170

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py` around lines 70 - 77,
Extract the repeated ExecutionState setup into a pytest fixture by adding a
fixture function named fixture_execution_state decorated with
`@pytest.fixture`(name="execution_state") that creates ExecutionState(), sets
execution_start_time = 0.0, and returns it; then update tests (e.g.,
test_basic_left_chosen and other tests using state variable) to accept the
execution_state fixture instead of instantiating ExecutionState inline and
remove the duplicated two-line setup (state = ExecutionState();
state.execution_start_time = 0.0) so each test uses the injected execution_state
object.
packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py (1)

911-925: Consider constructing StaticAnalysisResult directly instead of field-by-field copy.

Both _analyze_callee (lines 911–925) and analyze_function_ast (lines 1068–1081) manually copy ~14 fields from the visitor to the result. A helper or direct construction would reduce the risk of missing a field if new flags are added.

Possible approach
result = StaticAnalysisResult(
    reads=visitor.reads,
    writes=visitor.writes,
    mutations=visitor.mutations,
    detected_special_calls=visitor.detected_special_calls,
    has_dynamic_keys=visitor.has_dynamic_keys,
    has_unresolved_calls=visitor.has_unresolved_calls,
    recursion_depth_hit=visitor.recursion_depth_hit,
    has_dynamic_exec=visitor.has_dynamic_exec,
    has_closure_write=visitor.has_closure_write,
    has_global_write=visitor.has_global_write,
    has_unknown_attr_access=visitor.has_unknown_attr_access,
    has_return_lambda_mutates_state=visitor.has_return_lambda_mutates_state,
    has_dynamic_attr=visitor.has_dynamic_attr,
    warnings=visitor.warnings,
)

Or extract a _visitor_to_result(visitor) -> StaticAnalysisResult helper to DRY both call sites.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
911 - 925, Replace the manual field-by-field copy into StaticAnalysisResult in
_analyze_callee (and similarly in analyze_function_ast) with a direct
construction or a small helper to map a visitor to a result to avoid missing
fields later; for example implement a helper function
_visitor_to_result(visitor) -> StaticAnalysisResult that builds and returns
StaticAnalysisResult(...) using visitor.reads, visitor.writes,
visitor.mutations, visitor.detected_special_calls, visitor.has_dynamic_keys,
visitor.has_unresolved_calls, visitor.recursion_depth_hit,
visitor.has_dynamic_exec, visitor.has_closure_write, visitor.has_global_write,
visitor.has_unknown_attr_access, visitor.has_return_lambda_mutates_state,
visitor.has_dynamic_attr, visitor.warnings and call that helper from
_analyze_callee and analyze_function_ast.
packages/nvidia_nat_app/src/nat_app/api.py (1)

53-55: Re-exports look intentional — consider adding __all__ to make the public API explicit.

Lines 53–55 re-export SpeculationPlan, partition_targets, and plan_speculation with noqa: F401 suppression. This is fine, but an __all__ list at the top of the module would make the public surface self-documenting and prevent accidental exposure of internal imports.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/api.py` around lines 53 - 55, The module
re-exports SpeculationPlan, partition_targets, and plan_speculation but lacks an
explicit __all__; add an __all__ list (e.g. __all__ = ["SpeculationPlan",
"partition_targets", "plan_speculation"]) near the top of the module to declare
the public API and avoid accidental exports, and update it if you later expose
additional public symbols from nat_app.api.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@packages/nvidia_nat_app/src/nat_app/executors/runner.py`:
- Around line 196-215: When a chosen target task raises (caught in the
try/except around awaiting the task), ensure we record the failure so metrics
reconcile: capture the exception (e.g., as exc) and before continuing call
execution_state.mark_node_completed(name, {}), record the duration with
execution_state.record_node_duration(name, time.time() - start), emit a timeline
event with execution_state.record_timeline_event(name, start, time.time()), and
increment the existing tools_cancelled (or appropriate failure counter) so
launched == completed + cancelled + failed; keep the logger.exception call and
then continue.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py`:
- Around line 28-36: Docstring incorrectly claims "Recursion depth limited to 3"
while the code's default parameter is max_recursion_depth=5; update the
documentation to match the code. Edit the top-level module docstring in
static_analysis.py to say "Recursion depth limited to 5" (or otherwise reflect
the actual default), and search for the identifier max_recursion_depth in
functions/classes that reference it to ensure the docstring wording is
consistent with the defaults defined (occurrences around the functions that
accept max_recursion_depth at the noted sites).
- Around line 484-489: The use of zip(..., strict=True) when iterating
target.elts and value.elts can raise ValueError on starred unpacking (e.g. a,
*rest = x, y, z); change the iteration in the tuple-assignment handling to not
use strict=True (i.e., use the default zip behavior or otherwise guard for
mismatched lengths) so starred remainders are skipped instead of crashing,
leaving the existing logic that calls _resolve_state_source(v) and assigns to
self._aliases[t.id] intact.

---

Duplicate comments:
In `@packages/nvidia_nat_app/src/nat_app/api.py`:
- Line 98: The call to build_graph_and_adapter unpacks an unused variable named
adapter which triggers RUF059; change the unpacked name to a prefixed underscore
(e.g., _adapter or _) at this call site (and the other two sites referenced) so
the unused value is explicit and linter-suppressed, leaving graph as-is (i.e.,
replace "graph, adapter = build_graph_and_adapter(...)" with "graph, _adapter =
build_graph_and_adapter(...)" or similar).

---

Nitpick comments:
In `@packages/nvidia_nat_app/src/nat_app/api.py`:
- Around line 53-55: The module re-exports SpeculationPlan, partition_targets,
and plan_speculation but lacks an explicit __all__; add an __all__ list (e.g.
__all__ = ["SpeculationPlan", "partition_targets", "plan_speculation"]) near the
top of the module to declare the public API and avoid accidental exports, and
update it if you later expose additional public symbols from nat_app.api.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py`:
- Around line 911-925: Replace the manual field-by-field copy into
StaticAnalysisResult in _analyze_callee (and similarly in analyze_function_ast)
with a direct construction or a small helper to map a visitor to a result to
avoid missing fields later; for example implement a helper function
_visitor_to_result(visitor) -> StaticAnalysisResult that builds and returns
StaticAnalysisResult(...) using visitor.reads, visitor.writes,
visitor.mutations, visitor.detected_special_calls, visitor.has_dynamic_keys,
visitor.has_unresolved_calls, visitor.recursion_depth_hit,
visitor.has_dynamic_exec, visitor.has_closure_write, visitor.has_global_write,
visitor.has_unknown_attr_access, visitor.has_return_lambda_mutates_state,
visitor.has_dynamic_attr, visitor.warnings and call that helper from
_analyze_callee and analyze_function_ast.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py`:
- Around line 70-77: Extract the repeated ExecutionState setup into a pytest
fixture by adding a fixture function named fixture_execution_state decorated
with `@pytest.fixture`(name="execution_state") that creates ExecutionState(), sets
execution_start_time = 0.0, and returns it; then update tests (e.g.,
test_basic_left_chosen and other tests using state variable) to accept the
execution_state fixture instead of instantiating ExecutionState inline and
remove the duplicated two-line setup (state = ExecutionState();
state.execution_start_time = 0.0) so each test uses the injected execution_state
object.

In `@packages/nvidia_nat_app/tests/test_api.py`:
- Around line 231-240: The loop in test_info_has_reads_writes uses an unused
variable name "name" which should be renamed to "_name" to satisfy the B007 lint
rule; update the for statement in the test_info_has_reads_writes function (for
name, analysis in info.items()) to use for _name, analysis in info.items() so
only the used variable analysis remains named.
- Around line 384-414: Update the pytest.raises calls in the tests to use
raw-string regex literals to satisfy RUF043: change the match arguments in
test_execute_node_non_dict_sequential_raises and
test_execute_node_non_dict_parallel_raises from plain strings like "execute_node
must return a dict.*got NoneType.*node" and "execute_node must return a
dict.*got str.*node" to raw strings r"execute_node must return a dict.*got
NoneType.*node" and r"execute_node must return a dict.*got str.*node"
respectively; no other behavior changes needed (tests still use the same
execute_node helpers and benchmark call).

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8b9efa6 and d869391.

📒 Files selected for processing (5)
  • packages/nvidia_nat_app/src/nat_app/api.py
  • packages/nvidia_nat_app/src/nat_app/executors/runner.py
  • packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py
  • packages/nvidia_nat_app/tests/executors/test_runner.py
  • packages/nvidia_nat_app/tests/test_api.py

- Introduced the `nvidia-nat-app` package to the project dependencies and configuration files.
- Updated static analysis to handle starred unpacking without crashing, increasing robustness.
- Enhanced error handling in `run_speculation` to record execution metrics for failed tasks.
- Added tests to validate the new functionality and ensure proper behavior of the static analysis.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 5

♻️ Duplicate comments (1)
packages/nvidia_nat_app/tests/executors/test_runner.py (1)

62-62: ⚠️ Potential issue | 🟡 Minor

Outstanding TRY003 ruff warnings: long inline exception messages.

Lines 62, 297, 327, and 351 each carry a ruff TRY003 ("Avoid specifying long messages outside the exception class") warning that was present in the previous review cycle's static analysis output but was not fixed when the ARG issues were addressed. Per the guidelines, ruff warnings should be resolved unless explicitly suppressed in pyproject.toml.

The quickest fix in test code is an inline # noqa: TRY003 suppression comment:

🔧 Proposed suppressions
-        raise RuntimeError("resolve failed")
+        raise RuntimeError("resolve failed")  # noqa: TRY003
-                raise ValueError("simulated failure")
+                raise ValueError("simulated failure")  # noqa: TRY003
-                raise ValueError("decision failed")
+                raise ValueError("decision failed")  # noqa: TRY003
-            raise RuntimeError("bad decision")
+            raise RuntimeError("bad decision")  # noqa: TRY003

Also applies to: 297-297, 327-327, 351-351

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py` at line 62, The test
file contains long inline exception messages triggering ruff TRY003 at the raise
sites (e.g., the RuntimeError("resolve failed") at the top of
tests/executors/test_runner.py and the similar raise statements later around the
other failures); fix by adding an inline suppression comment "# noqa: TRY003" to
each offending raise statement (the RuntimeError raises at the locations
referenced) so the test code silences TRY003 without changing behavior.
🧹 Nitpick comments (3)
pyproject.toml (1)

91-91: Consider whether an experimental package belongs in most.

The PR title calls nvidia-nat-app experimental, yet it is included in the most aggregate extra. Anyone running uv pip install "nvidia-nat[most]" will automatically pull it in. If the package's API surface is not yet stable, a safer landing would be to omit it from most until it graduates out of experimental status — consistent with how openpipe-art and ragaai are intentionally excluded.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@pyproject.toml` at line 91, The package "nvidia-nat-app" is experimental but
currently included in the "most" extras list; remove "nvidia-nat-app ==
{version}" from the "most" extra in pyproject.toml (the extras table named
"most") so it is not pulled in by uv pip install "nvidia-nat[most]" until it
graduates, mirroring how "openpipe-art" and "ragaai" are excluded.
packages/nvidia_nat_app/tests/executors/test_runner.py (2)

103-103: Redundant lambda wrappers for _slow_node — use the function directly.

lambda name: _slow_node(name) at lines 103, 120, 139, 154, 175, 194, and 256 is a no-op wrapper; _slow_node already matches the expected Callable[[str], Awaitable[...]] signature and can be passed directly.

♻️ Proposed refactor (apply to all seven occurrences)
-            run_node=lambda name: _slow_node(name),
+            run_node=_slow_node,
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py` at line 103, The test
uses redundant lambda wrappers like run_node=lambda name: _slow_node(name);
replace those lambda expressions by passing the coroutine function directly
(e.g., set run_node=_slow_node) wherever they appear so the Callable[[str],
Awaitable[...]] is provided directly; update all occurrences referencing
run_node and _slow_node in the test file to remove the no-op lambdas.

70-385: Extract repeated ExecutionState setup into a pytest fixture.

ExecutionState() followed by state.execution_start_time = 0.0 appears ~13 times across the test class. The coding guidelines explicitly state: "Any frequently repeated code should be extracted into pytest fixtures", and the fixture naming convention (fixture_ prefix, name argument on the decorator) must be followed.

♻️ Proposed refactor
+import pytest
+
+
+@pytest.fixture(name="state")
+def fixture_state() -> ExecutionState:
+    state = ExecutionState()
+    state.execution_start_time = 0.0
+    return state
+
+
 class TestRunSpeculativeRouter:
 
-    async def test_basic_left_chosen(self):
-        plan = _make_plan()
-        state = ExecutionState()
-        state.execution_start_time = 0.0
+    async def test_basic_left_chosen(self, state: ExecutionState):
+        plan = _make_plan()
         ...

Apply the same substitution to all methods in TestRunSpeculativeRouter.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/executors/test_runner.py` around lines 70 -
385, Extract the repeated ExecutionState creation and start-time assignment into
a pytest fixture named fixture_state (use `@pytest.fixture`(name="state")), which
returns an ExecutionState instance with state.execution_start_time = 0.0;
replace the manual state setup in tests like test_basic_left_chosen,
test_basic_right_chosen, test_metrics_tracked, test_router_decision_recorded,
test_timeline_events_recorded, test_no_cancel_map_entry, test_single_target,
test_fast_target_already_done_before_router, test_non_dict_results,
test_cancel_map_references_node_not_in_targets,
test_cancelled_task_already_done_before_cancel,
test_chosen_target_raises_continues_without_crash,
test_decision_task_raises_cancels_targets,
test_get_decision_raises_cancels_targets, and
test_resolution_resolve_raises_cancels_targets to accept the fixture parameter
state (remove local ExecutionState() and start-time lines); ensure all
references to ExecutionState in these tests use the injected state and that
run_speculation calls continue to pass state as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py`:
- Around line 15-109: The module and class docstrings reference code identifiers
(e.g., Send, Command, has_dynamic_exec) without inline code formatting which
triggers Vale warnings; update the docstrings in this file (top module docstring
and StaticAnalysisResult docstring) to wrap code entities in backticks (e.g.,
`Send`, `Command`, `has_dynamic_exec`) and any other plain identifiers mentioned
(like `state`, `param_to_obj`, `AccessSet`) so all code symbols are formatted as
inline code.
- Around line 1084-1089: The logging call using func.__name__ can raise
AttributeError for callable objects like functools.partial or instances; update
the logger.debug invocation in static_analysis.py to safely obtain a name (e.g.,
use a fallback via getattr or repr) so it never tries to access a missing
__name__; locate the logger.debug call referencing func.__name__ and replace the
name expression with a safe lookup (e.g., name = getattr(func, "__name__",
repr(func)) or similar) and use that variable in the log along with the existing
result.reads/result.writes/result.mutations/result.specials references.
- Around line 720-786: The three private helper methods _resolve_name,
_resolve_subscript_container, and _resolve_subscript_key lack return type
annotations; add explicit return types (e.g., from typing import Any, Optional)
so signatures become _resolve_name(self, name: str) -> Optional[Any],
_resolve_subscript_container(self, node: ast.expr) -> Optional[Any], and
_resolve_subscript_key(self, node: ast.expr) -> Optional[Any]; also ensure the
typing imports (Any, Optional) are present at the top of the module.
- Around line 1019-1023: The code uses state_param =
next(iter(effective_param_to_obj)) which raises StopIteration if param_to_obj is
empty; update the logic around effective_param_to_obj/state_param to validate
param_to_obj first and fail fast with a clear error (e.g., raise
ValueError("param_to_obj must not be empty") ) or choose a safe fallback (e.g.,
set state_param to a known default like _DEFAULT_OBJ) before calling
next(iter(...)); modify the block that constructs effective_param_to_obj and
sets state_param to perform this check and emit the clear validation error or
fallback.

In `@packages/nvidia_nat_app/tests/graph/test_static_analysis.py`:
- Around line 45-47: The helper functions (e.g., the inner def fn(state) in the
test file) declare an unused parameter `state`, causing Ruff ARG001; fix by
renaming the parameter to `_state` (or `_*` variant) or by referencing it
explicitly (e.g., assert _state is not None) so the linter sees it used; apply
the same change to all flagged helpers at the other ranges (lines around 75-77,
172-174, 351-354, 360-363, 424-427, 433-436, 442-445) to eliminate ARG001
warnings.

---

Duplicate comments:
In `@packages/nvidia_nat_app/tests/executors/test_runner.py`:
- Line 62: The test file contains long inline exception messages triggering ruff
TRY003 at the raise sites (e.g., the RuntimeError("resolve failed") at the top
of tests/executors/test_runner.py and the similar raise statements later around
the other failures); fix by adding an inline suppression comment "# noqa:
TRY003" to each offending raise statement (the RuntimeError raises at the
locations referenced) so the test code silences TRY003 without changing
behavior.

---

Nitpick comments:
In `@packages/nvidia_nat_app/tests/executors/test_runner.py`:
- Line 103: The test uses redundant lambda wrappers like run_node=lambda name:
_slow_node(name); replace those lambda expressions by passing the coroutine
function directly (e.g., set run_node=_slow_node) wherever they appear so the
Callable[[str], Awaitable[...]] is provided directly; update all occurrences
referencing run_node and _slow_node in the test file to remove the no-op
lambdas.
- Around line 70-385: Extract the repeated ExecutionState creation and
start-time assignment into a pytest fixture named fixture_state (use
`@pytest.fixture`(name="state")), which returns an ExecutionState instance with
state.execution_start_time = 0.0; replace the manual state setup in tests like
test_basic_left_chosen, test_basic_right_chosen, test_metrics_tracked,
test_router_decision_recorded, test_timeline_events_recorded,
test_no_cancel_map_entry, test_single_target,
test_fast_target_already_done_before_router, test_non_dict_results,
test_cancel_map_references_node_not_in_targets,
test_cancelled_task_already_done_before_cancel,
test_chosen_target_raises_continues_without_crash,
test_decision_task_raises_cancels_targets,
test_get_decision_raises_cancels_targets, and
test_resolution_resolve_raises_cancels_targets to accept the fixture parameter
state (remove local ExecutionState() and start-time lines); ensure all
references to ExecutionState in these tests use the injected state and that
run_speculation calls continue to pass state as before.

In `@pyproject.toml`:
- Line 91: The package "nvidia-nat-app" is experimental but currently included
in the "most" extras list; remove "nvidia-nat-app == {version}" from the "most"
extra in pyproject.toml (the extras table named "most") so it is not pulled in
by uv pip install "nvidia-nat[most]" until it graduates, mirroring how
"openpipe-art" and "ragaai" are excluded.

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d869391 and 6a5e5ba.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (5)
  • packages/nvidia_nat_app/src/nat_app/executors/runner.py
  • packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py
  • packages/nvidia_nat_app/tests/executors/test_runner.py
  • packages/nvidia_nat_app/tests/graph/test_static_analysis.py
  • pyproject.toml
🚧 Files skipped from review as they are similar to previous changes (1)
  • packages/nvidia_nat_app/src/nat_app/executors/runner.py

- Improved docstrings for clarity, using backticks for code references.
- Added validation to raise a ValueError if `param_to_obj` is empty in `analyze_function_ast`.
- Updated test cases to ensure proper handling of empty `param_to_obj` and adjusted function signatures for consistency.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
- Updated the `build_wheel.sh` script to skip the import and CLI test for `nvidia_nat_app`, which has a different structure.
- Added import verification for `nat_app` when processing other wheels, ensuring proper error logging for import failures.
- Cleaned up the `pyproject.toml` by removing commented-out sections for clarity.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
Copy link

@coderabbitai coderabbitai bot left a comment

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 (9)
ci/scripts/github/build_wheel.sh (1)

108-114: Add a success log after a passing nat_app import.

The nat branch implicitly confirms success via the nat --version output. The nat_app branch is silent on success, making CI logs harder to reason about when triaging failures.

📝 Proposed fix
             PYTHON_IMPORT_OUT=$(python -c "import nat_app" 2>&1)
             IMPORT_TEST_RESULT=$?
             if [[ ${IMPORT_TEST_RESULT} -ne 0 ]]; then
                 rapids-logger "Error, failed to import nat_app from wheel ${whl} with Python ${pyver}"
                 echo "${PYTHON_IMPORT_OUT}"
                 exit ${IMPORT_TEST_RESULT}
             fi
+            rapids-logger "Successfully imported nat_app from wheel ${whl} with Python ${pyver}"
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@ci/scripts/github/build_wheel.sh` around lines 108 - 114, The script
currently only logs on import failure of nat_app; add a success message when
import succeeds by inserting a rapids-logger call after the existing failure
check (use the same context variables IMPORT_TEST_RESULT, PYTHON_IMPORT_OUT,
whl, pyver). Specifically, after the if block that checks IMPORT_TEST_RESULT,
invoke rapids-logger with a clear success message like "Successfully imported
nat_app from wheel ${whl} with Python ${pyver}" (optionally echo
PYTHON_IMPORT_OUT for verbose output) so CI logs show a passing nat_app import.
packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py (7)

140-186: Bare set type hints lose semantic information.

visited_funcs stores id() return values (integers), but the annotation is set / set | None. Using set[int] would communicate intent and help static analysis tools like pyright.

Same applies at Line 810 in _analyze_callee.

Proposed fix
-        visited_funcs: set | None = None,
+        visited_funcs: set[int] | None = None,
-        self._visited_funcs: set = visited_funcs or set()
+        self._visited_funcs: set[int] = visited_funcs or set()

And similarly in _analyze_callee:

-    visited_funcs: set,
+    visited_funcs: set[int],
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
140 - 186, The type hints use bare set for visited function id values; change
the annotations from set / set | None to set[int] and set[int] | None for the
constructor parameter visited_funcs and the _visited_funcs attribute
(self._visited_funcs) in __init__ (function __init__ in static_analysis.py), and
likewise update the corresponding annotation in the _analyze_callee function
signature (where visited_funcs is accepted) to use set[int] | None so the intent
(ids are integers) is explicit for static checkers.

1137-1151: Missing docstring for _get_state_param_name.

Other private helpers (_extract_lambda_source, _get_lambda_param_name) have Google-style docstrings, but _get_state_param_name does not. As per coding guidelines, "Provide Google-style docstrings for every public module, class, function and CLI command." While this is a private function, consistency with the adjacent helpers would be beneficial.

Suggested docstring
 def _get_state_param_name(func_def: ast.FunctionDef | ast.AsyncFunctionDef) -> str | None:
+    """Extract the state parameter name from a function definition.
+
+    Skips ``self``/``cls`` if present; falls back to ``*args`` when no
+    positional parameters exist.
+
+    Args:
+        func_def: A FunctionDef or AsyncFunctionDef AST node.
+
+    Returns:
+        The state parameter name, or None if no suitable parameter is found.
+    """
     args = func_def.args
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
1137 - 1151, Add a Google-style docstring to the private helper
_get_state_param_name describing its purpose (determine the name of the
“state”/first relevant positional parameter for a function or return None), the
arguments (func_def: ast.FunctionDef | ast.AsyncFunctionDef), and the return
value (str | None), and briefly note handling of self/cls and vararg cases;
place the docstring immediately below the def _get_state_param_name(...) line to
match the style used by _extract_lambda_source and _get_lambda_param_name.

935-941: Public API docstring uses triple-quote style instead of single-line summary.

The first line of the docstring (""") is separate from the summary sentence, which starts on the next line. Per coding guidelines, "The first line of docstrings must be a concise description ending with a period." — typically this means the summary should start on the same line as the opening """.

Proposed adjustment
-    """
-    Analyze a function via AST to detect state reads and writes.
+    """Analyze a function via AST to detect state reads and writes.

As per coding guidelines, "The first line of docstrings must be a concise description ending with a period."

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
935 - 941, The public API docstring for analyze_function_ast currently places
the summary on the second line; update the docstring to have a single-line
summary that starts immediately after the opening triple quotes and ends with a
period (e.g., """Analyze a function's AST and produce a
StaticAnalysisResult."""), and keep any longer explanation on subsequent lines;
modify the docstring in the analyze_function_ast definition to follow this
single-line summary convention.

912-927: Extract visitor-to-result copying into a shared helper.

Lines 912-927 and 1073-1086 contain the same 15-line block that copies every attribute from the visitor to a StaticAnalysisResult. If a new flag is added, both sites must be updated in sync. A helper method (e.g., _result_from_visitor) would eliminate this duplication.

Suggested extraction
+def _result_from_visitor(visitor: _NodeASTVisitor) -> StaticAnalysisResult:
+    """Build a StaticAnalysisResult from a completed visitor."""
+    return StaticAnalysisResult(
+        reads=visitor.reads,
+        writes=visitor.writes,
+        mutations=visitor.mutations,
+        detected_special_calls=visitor.detected_special_calls,
+        has_dynamic_keys=visitor.has_dynamic_keys,
+        has_unresolved_calls=visitor.has_unresolved_calls,
+        recursion_depth_hit=visitor.recursion_depth_hit,
+        has_dynamic_exec=visitor.has_dynamic_exec,
+        has_closure_write=visitor.has_closure_write,
+        has_global_write=visitor.has_global_write,
+        has_unknown_attr_access=visitor.has_unknown_attr_access,
+        has_return_lambda_mutates_state=visitor.has_return_lambda_mutates_state,
+        has_dynamic_attr=visitor.has_dynamic_attr,
+        warnings=visitor.warnings,
+    )

Then replace both copy blocks with result = _result_from_visitor(visitor).

Also applies to: 1073-1086

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
912 - 927, The repeated 15-line block that copies attributes from the `visitor`
into a new `StaticAnalysisResult` instance should be extracted into a single
helper (e.g., `_result_from_visitor(visitor)`) to avoid duplication; implement
`_result_from_visitor` to instantiate `StaticAnalysisResult`, copy the listed
attributes (reads, writes, mutations, detected_special_calls, has_dynamic_keys,
has_unresolved_calls, recursion_depth_hit, has_dynamic_exec, has_closure_write,
has_global_write, has_unknown_attr_access, has_return_lambda_mutates_state,
has_dynamic_attr, warnings) from the provided `visitor`, return it, and then
replace both existing copy blocks (the one returning `result` at the end of the
current function and the similar block at lines ~1073-1086) with `result =
_result_from_visitor(visitor)` followed by the existing return.

356-382: Duplicated closure/global write detection pattern.

The block at Lines 364-371 (checking base_name in self._freevars vs. base_name not in param_names) is repeated nearly identically in visit_AugAssign (397-404), visit_Delete (421-428), and _handle_assign_target (446-453). Consider extracting a small helper to reduce repetition.

Suggested helper
+    def _flag_nonlocal_write(self, base_name: str):
+        """Flag closure or global write if base_name is not a tracked param."""
+        param_names = set(self._param_to_obj) | set(self._state_aliases)
+        if base_name in self._freevars:
+            self.has_closure_write = True
+        elif base_name not in param_names:
+            self.has_global_write = True

Then replace each occurrence with:

self._flag_nonlocal_write(base_name)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
356 - 382, Extract the duplicated closure/global write detection into a helper
(e.g., _flag_nonlocal_write(self, base_name)) and call it from visit_Subscript,
visit_AugAssign, visit_Delete and _handle_assign_target instead of repeating the
logic; the helper should assume base_name is not None, compute param_names =
set(self._param_to_obj) | set(self._state_aliases), set self.has_closure_write =
True if base_name in self._freevars, else set self.has_global_write = True if
base_name not in param_names. Replace the repeated blocks in the methods
mentioned with a single call like self._flag_nonlocal_write(base_name).

1004-1017: "_self_state_placeholder" is a magic string that could confuse future maintainers.

When self_state_attrs is provided but no regular params exist (class method with only self), the code creates a synthetic state_param = "_self_state_placeholder" with an empty effective_param_to_obj. A named constant and an inline comment explaining why the param name is irrelevant (since self.state access is handled by _self_state_attrs) would improve clarity.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
1004 - 1017, Replace the magic string "_self_state_placeholder" with a named
constant (e.g., SELF_STATE_PLACEHOLDER) and add a brief inline comment where
state_param is set explaining that the placeholder is only used to satisfy
downstream logic and has no semantic meaning because self_state_attrs handles
self.state attribute access; update references to state_param and
effective_param_to_obj in the same block (function: _get_state_param_name usage
and the branch that assigns state_param and effective_param_to_obj) to use the
new constant to improve clarity and maintainability.

1101-1118: _extract_lambda_source is fragile for multi-line or nested lambdas.

The heuristic source[idx:] followed by trailing-comma stripping works for simple dict-entry lambdas but will break for lambdas followed by closing brackets, other arguments, or multi-line expressions. Since the outer try/except SyntaxError catches failures gracefully, this is not a correctness bug — just a known limitation worth documenting.

Suggested documentation addition
 def _extract_lambda_source(source: str) -> str | None:
     """Extract a lambda expression from dict-entry or assignment source.
 
+    This is a best-effort heuristic that handles simple cases
+    (e.g. ``"key": lambda s: {...},``).  Complex multi-line or nested
+    lambdas may not parse; callers should handle ``None`` or subsequent
+    ``SyntaxError`` gracefully.
+
     Args:
         source: Raw source code that may contain a lambda.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py` around lines
1101 - 1118, Update the _extract_lambda_source function docstring to explicitly
state its heuristic limitation: it searches for the first "lambda" and slices to
the end, so it may fail for multi-line or nested lambdas, lambdas followed by
closing brackets, other arguments, or complex trailing punctuation; mention that
SyntaxError is caught by callers and recommend adding tests or replacing this
heuristic if more robust parsing is later required. Include the function name
_extract_lambda_source in the docstring so maintainers can quickly find the
implementation and rationale.
packages/nvidia_nat_app/tests/graph/test_static_analysis.py (1)

97-107: Consider strengthening the starred-unpacking assertion.

The assertion on Line 106 is deliberately lenient ("x" in ... or "y" in ...), but since zip(strict=False) pairs elements by index, all three reads (x, y, z) should be detected for the non-starred targets that align. If the analyzer correctly walks the RHS tuple, asserting all three are present would make the regression test more robust against silent regressions.

That said, the primary goal ("no crash") is achieved, so this is optional.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_app/tests/graph/test_static_analysis.py` around lines 97
- 107, The test test_starred_unpacking_no_crash currently asserts a lenient read
(either "x" or "y") but should require that all RHS names are detected; update
the assertion in test_starred_unpacking_no_crash (the call to
analyze_function_ast and the use of r.reads.all_fields_flat) to assert that "x",
"y", and "z" are present (e.g., require the set of expected fields is a subset
of r.reads.all_fields_flat) so the analyzer's tuple-walking for the RHS of the
starred unpack is validated more strictly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@ci/scripts/github/build_wheel.sh`:
- Line 107: The log message saying "Skipping nat import/CLI test for
nvidia_nat_app" is misleading because the script still executes the Python
import (import nat_app) while only the nat CLI check is skipped; update the
rapids-logger message to accurately reflect behavior (e.g., "Skipping nat CLI
test; performing import check for nvidia_nat_app") or alternatively remove the
import nat_app call if you intend to skip both checks; reference the
rapids-logger log invocation and the import nat_app statement when making the
change.

In `@packages/nvidia_nat_app/pyproject.toml`:
- Around line 21-23: Update the package namespace pattern in the setuptools
config and rename the source package directory: change the include value under
[tool.setuptools.packages.find] from ["nat_app*"] to ["nat.*"] and then
move/rename the code directory from src/nat_app/ to src/nat/ (and update any
import paths or references that point to nat_app to use nat) so the package
follows the nat.* API namespace convention used across the repo.

---

Nitpick comments:
In `@ci/scripts/github/build_wheel.sh`:
- Around line 108-114: The script currently only logs on import failure of
nat_app; add a success message when import succeeds by inserting a rapids-logger
call after the existing failure check (use the same context variables
IMPORT_TEST_RESULT, PYTHON_IMPORT_OUT, whl, pyver). Specifically, after the if
block that checks IMPORT_TEST_RESULT, invoke rapids-logger with a clear success
message like "Successfully imported nat_app from wheel ${whl} with Python
${pyver}" (optionally echo PYTHON_IMPORT_OUT for verbose output) so CI logs show
a passing nat_app import.

In `@packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py`:
- Around line 140-186: The type hints use bare set for visited function id
values; change the annotations from set / set | None to set[int] and set[int] |
None for the constructor parameter visited_funcs and the _visited_funcs
attribute (self._visited_funcs) in __init__ (function __init__ in
static_analysis.py), and likewise update the corresponding annotation in the
_analyze_callee function signature (where visited_funcs is accepted) to use
set[int] | None so the intent (ids are integers) is explicit for static
checkers.
- Around line 1137-1151: Add a Google-style docstring to the private helper
_get_state_param_name describing its purpose (determine the name of the
“state”/first relevant positional parameter for a function or return None), the
arguments (func_def: ast.FunctionDef | ast.AsyncFunctionDef), and the return
value (str | None), and briefly note handling of self/cls and vararg cases;
place the docstring immediately below the def _get_state_param_name(...) line to
match the style used by _extract_lambda_source and _get_lambda_param_name.
- Around line 935-941: The public API docstring for analyze_function_ast
currently places the summary on the second line; update the docstring to have a
single-line summary that starts immediately after the opening triple quotes and
ends with a period (e.g., """Analyze a function's AST and produce a
StaticAnalysisResult."""), and keep any longer explanation on subsequent lines;
modify the docstring in the analyze_function_ast definition to follow this
single-line summary convention.
- Around line 912-927: The repeated 15-line block that copies attributes from
the `visitor` into a new `StaticAnalysisResult` instance should be extracted
into a single helper (e.g., `_result_from_visitor(visitor)`) to avoid
duplication; implement `_result_from_visitor` to instantiate
`StaticAnalysisResult`, copy the listed attributes (reads, writes, mutations,
detected_special_calls, has_dynamic_keys, has_unresolved_calls,
recursion_depth_hit, has_dynamic_exec, has_closure_write, has_global_write,
has_unknown_attr_access, has_return_lambda_mutates_state, has_dynamic_attr,
warnings) from the provided `visitor`, return it, and then replace both existing
copy blocks (the one returning `result` at the end of the current function and
the similar block at lines ~1073-1086) with `result =
_result_from_visitor(visitor)` followed by the existing return.
- Around line 356-382: Extract the duplicated closure/global write detection
into a helper (e.g., _flag_nonlocal_write(self, base_name)) and call it from
visit_Subscript, visit_AugAssign, visit_Delete and _handle_assign_target instead
of repeating the logic; the helper should assume base_name is not None, compute
param_names = set(self._param_to_obj) | set(self._state_aliases), set
self.has_closure_write = True if base_name in self._freevars, else set
self.has_global_write = True if base_name not in param_names. Replace the
repeated blocks in the methods mentioned with a single call like
self._flag_nonlocal_write(base_name).
- Around line 1004-1017: Replace the magic string "_self_state_placeholder" with
a named constant (e.g., SELF_STATE_PLACEHOLDER) and add a brief inline comment
where state_param is set explaining that the placeholder is only used to satisfy
downstream logic and has no semantic meaning because self_state_attrs handles
self.state attribute access; update references to state_param and
effective_param_to_obj in the same block (function: _get_state_param_name usage
and the branch that assigns state_param and effective_param_to_obj) to use the
new constant to improve clarity and maintainability.
- Around line 1101-1118: Update the _extract_lambda_source function docstring to
explicitly state its heuristic limitation: it searches for the first "lambda"
and slices to the end, so it may fail for multi-line or nested lambdas, lambdas
followed by closing brackets, other arguments, or complex trailing punctuation;
mention that SyntaxError is caught by callers and recommend adding tests or
replacing this heuristic if more robust parsing is later required. Include the
function name _extract_lambda_source in the docstring so maintainers can quickly
find the implementation and rationale.

In `@packages/nvidia_nat_app/tests/graph/test_static_analysis.py`:
- Around line 97-107: The test test_starred_unpacking_no_crash currently asserts
a lenient read (either "x" or "y") but should require that all RHS names are
detected; update the assertion in test_starred_unpacking_no_crash (the call to
analyze_function_ast and the use of r.reads.all_fields_flat) to assert that "x",
"y", and "z" are present (e.g., require the set of expected fields is a subset
of r.reads.all_fields_flat) so the analyzer's tuple-walking for the RHS of the
starred unpack is validated more strictly.

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6a5e5ba and 5619ea2.

📒 Files selected for processing (4)
  • ci/scripts/github/build_wheel.sh
  • packages/nvidia_nat_app/pyproject.toml
  • packages/nvidia_nat_app/src/nat_app/graph/static_analysis.py
  • packages/nvidia_nat_app/tests/graph/test_static_analysis.py

…rification

- Modified the log message to specify that the nat CLI test is being skipped for `nvidia_nat_app`, while still verifying the import of `nat_app`.
- This change enhances clarity in the build process and ensures proper logging of import verification results.

Signed-off-by: Matthew Penn <mpenn@nvidia.com>
@mpenn
Copy link
Contributor Author

mpenn commented Feb 24, 2026

/merge

@rapids-bot rapids-bot bot merged commit 61715ce into NVIDIA:develop Feb 24, 2026
15 checks passed
Charlie-Yi-2002 pushed a commit to Charlie-Yi-2002/NeMo-Agent-Toolkit that referenced this pull request Mar 5, 2026
…ge (NVIDIA#1636)

Closes

Adds **nvidia-nat-app**, a framework-agnostic library of Agent Performance Primitives (APP) for accelerating agentic applications. This PR focuses provides graph analysis, parallelization, and speculative execution primitives that can be embedded in any agent framework.

- **Graph optimization** — `quick_optimize()` takes `nodes` and `edges` and returns parallel execution stages via a 6-stage pipeline: extract → validate → topology → node analysis → edge classification → scheduling.
- **AST-based static analysis** — `analyze_function()` and `classify_edge()` infer state reads/writes and edge necessity without executing code.
- **Speculative execution** — `plan_speculation()`, `speculative_opportunities()`, and `run_speculation()` support router-branch speculation with `@speculation_unsafe` and `SpeculationSafetyConfig`.
- **Benchmarking** — `benchmark()` compares sequential vs. parallel execution and custom strategies.
- **Framework integration** — `AbstractFrameworkAdapter` and `build_graph_and_adapter()` for framework-specific integration; public API uses plain Python types.
- **Constraints** — `@sequential`, `@depends_on`, `@has_side_effects`, and `OptimizationConfig` for controlling optimization.

`quick_optimize` | `analyze_function` | `classify_edge` | `find_parallel_stages` | `benchmark` | `speculative_opportunities` | `plan_speculation` | `partition_targets` | `run_speculation`

- Framework-agnostic, no external dependencies. Python 3.11–3.13. Experimental (API may change).

```python
from nat_app import quick_optimize

stages = quick_optimize(
    nodes={"parse": parse_fn, "research_a": fn_a, "research_b": fn_b, "synthesize": fn_c},
    edges=[("parse", "research_a"), ("parse", "research_b"),
           ("research_a", "synthesize"), ("research_b", "synthesize")],
)
```

- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing/index.md).
- 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.

* **New Features**
  * Add a complete framework‑agnostic optimization toolkit: graph model, static AST analysis, topology & scheduling, pipelined compiler, optimization stages, and high‑level APIs (quick_optimize, analyze_function, classify_edge, find_parallel_stages, benchmark, speculative_opportunities).
  * Speculation planning and execution: router‑branch strategy, resolution policies, planner, and a speculative runner; runtime utilities for execution state, metrics, result handling, constraints, and priority assignment.

* **Documentation**
  * Package metadata and PyPI documentation for the NVIDIA NAT‑APP subpackage.

* **Tests**
  * Extensive unit and integration tests covering compiler, graph, analysis, scheduling, speculation, stages, executors, and utilities.

Authors:
  - Matthew Penn (https://github.com/mpenn)

Approvers:
  - Dhruv Nandakumar (https://github.com/dnandakumar-nv)
  - David Gardner (https://github.com/dagardner-nv)
  - Will Killian (https://github.com/willkill07)

URL: NVIDIA#1636
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants