Add experimental nvidia-nat-app Agent Performance Primitives subpackage#1636
Add experimental nvidia-nat-app Agent Performance Primitives subpackage#1636rapids-bot[bot] merged 15 commits intoNVIDIA:developfrom
Conversation
…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>
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds 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
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)
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
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. Comment |
…re proper functionality. Signed-off-by: Matthew Penn <mpenn@nvidia.com>
- 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>
- 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>
There was a problem hiding this comment.
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
nameargument and use thefixture_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 toMinimalAdaptermethods.
extractandbuildare 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 resultAdd the import at the top:
+from typing import AnyAs 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 usingSequence[str]forpossible_targets.The coding guidelines prefer
collections.abcabstractions (Sequenceoverlist) 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.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.”♻️ 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🤖 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 repeatedExecutionStatesetup into a pytest fixture.The pattern
state = ExecutionState(); state.execution_start_time = 0.0is 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 parameterThen update each test method to accept
execution_stateas 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.valuestring comparison.Replace
ea.edge_type.value == "necessary"withea.edge_type == EdgeType.NECESSARYfor 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 similarnodes,edges, andconditional_edgesarguments. Consolidate these into reusable fixtures named with thefixture_prefix and set the fixturename=argument per coding guidelines.Examples from the test suite:
- Basic two-node graph:
nodes={"router": route_fn, "a": fn_a, "b": fn_b}withedges=[("router", "a"), ("router", "b")]and matchingconditional_edges- Three-node graph with merge: same router and branches plus
"merge": fn_mergewith additional edges to merge nodeThis 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 addingstrict=Truetozip()for extra safety.
group.node_namesandgroup.prioritiesshould be the same length; usingstrict=Truewill raiseValueErrorif 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_graphhas a stray comma in the signature.This is minor but noisy (and likely to be reformatted by yapf anyway).
As per coding guidelines: “Run yapf… for code formatting.”Proposed cleanup
-def _nested_conditional_parallel_graph(llm_analysis: dict[str, LLMCallInfo], ) -> CompilationContext: +def _nested_conditional_parallel_graph(llm_analysis: dict[str, LLMCallInfo]) -> CompilationContext:🤖 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 intest_fast_branch_gets_high_slow_gets_low().
_router_graph()already adds the target nodes and conditional edges for the passedbranches, so the extraadd_node("medium_path")andadd_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(...)insuccessors(),predecessors(), andto_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_sourcesexpose 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_setis aset, 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 prefernvidia-natandnatas namespace—please confirm the intentional exception for this experimental package.quick_optimize()doc:conditional_edgesdescription reads like{target_name: target_name}; should be{branch_name: target_node(s)}.speculative_opportunities()doc: thebranchesbullet 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 trustsopp.metadatashape; consider tightening.
plan()recomputes topology even thoughidentify()already received it, and it assumesopp.metadata["branches"]/["merge_nodes"]exist with the right types. A small guard (or a typedmetadataobject) 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:RouterBranchResolutionlogic is clear; consider an explicit “unknown decision” policy.Right now an unexpected
decision_resultdefaults 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; backticknat_app.This module docstring starts with a blank first line and mentions
nat_appwithout 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: UseTYPE_CHECKINGinstead of import-timeassert isinstance()for protocol conformance checks.While
ResolutionPolicyis marked@runtime_checkableso 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
⛔ Files ignored due to path filters (1)
packages/nvidia_nat_app/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (94)
packages/nvidia_nat_app/pyproject.tomlpackages/nvidia_nat_app/src/meta/pypi.mdpackages/nvidia_nat_app/src/nat_app/__init__.pypackages/nvidia_nat_app/src/nat_app/api.pypackages/nvidia_nat_app/src/nat_app/compiler/__init__.pypackages/nvidia_nat_app/src/nat_app/compiler/compilation_context.pypackages/nvidia_nat_app/src/nat_app/compiler/compilation_stage.pypackages/nvidia_nat_app/src/nat_app/compiler/compiler.pypackages/nvidia_nat_app/src/nat_app/compiler/default_graph_compiler.pypackages/nvidia_nat_app/src/nat_app/compiler/errors.pypackages/nvidia_nat_app/src/nat_app/compiler/optimizer.pypackages/nvidia_nat_app/src/nat_app/compiler/pipelined_compiler.pypackages/nvidia_nat_app/src/nat_app/constraints/__init__.pypackages/nvidia_nat_app/src/nat_app/constraints/decorators.pypackages/nvidia_nat_app/src/nat_app/constraints/models.pypackages/nvidia_nat_app/src/nat_app/constraints/resolution.pypackages/nvidia_nat_app/src/nat_app/executors/__init__.pypackages/nvidia_nat_app/src/nat_app/executors/execution_state.pypackages/nvidia_nat_app/src/nat_app/executors/metrics.pypackages/nvidia_nat_app/src/nat_app/executors/result_handler.pypackages/nvidia_nat_app/src/nat_app/executors/runner.pypackages/nvidia_nat_app/src/nat_app/graph/__init__.pypackages/nvidia_nat_app/src/nat_app/graph/access.pypackages/nvidia_nat_app/src/nat_app/graph/adapter.pypackages/nvidia_nat_app/src/nat_app/graph/analysis.pypackages/nvidia_nat_app/src/nat_app/graph/factory.pypackages/nvidia_nat_app/src/nat_app/graph/llm_detection.pypackages/nvidia_nat_app/src/nat_app/graph/models.pypackages/nvidia_nat_app/src/nat_app/graph/protocols.pypackages/nvidia_nat_app/src/nat_app/graph/scheduling.pypackages/nvidia_nat_app/src/nat_app/graph/static_analysis.pypackages/nvidia_nat_app/src/nat_app/graph/topology.pypackages/nvidia_nat_app/src/nat_app/graph/types.pypackages/nvidia_nat_app/src/nat_app/speculation/__init__.pypackages/nvidia_nat_app/src/nat_app/speculation/plan.pypackages/nvidia_nat_app/src/nat_app/speculation/planner.pypackages/nvidia_nat_app/src/nat_app/speculation/resolution.pypackages/nvidia_nat_app/src/nat_app/speculation/safety.pypackages/nvidia_nat_app/src/nat_app/speculation/strategies/__init__.pypackages/nvidia_nat_app/src/nat_app/speculation/strategies/base.pypackages/nvidia_nat_app/src/nat_app/speculation/strategies/router_branch.pypackages/nvidia_nat_app/src/nat_app/stages/__init__.pypackages/nvidia_nat_app/src/nat_app/stages/edge_classification.pypackages/nvidia_nat_app/src/nat_app/stages/extract.pypackages/nvidia_nat_app/src/nat_app/stages/llm_analysis.pypackages/nvidia_nat_app/src/nat_app/stages/node_analysis.pypackages/nvidia_nat_app/src/nat_app/stages/priority_assignment.pypackages/nvidia_nat_app/src/nat_app/stages/scheduling.pypackages/nvidia_nat_app/src/nat_app/stages/topology.pypackages/nvidia_nat_app/src/nat_app/stages/validate.pypackages/nvidia_nat_app/tests/__init__.pypackages/nvidia_nat_app/tests/compiler/__init__.pypackages/nvidia_nat_app/tests/compiler/test_compilation_context.pypackages/nvidia_nat_app/tests/compiler/test_compilation_stage.pypackages/nvidia_nat_app/tests/compiler/test_compiler.pypackages/nvidia_nat_app/tests/compiler/test_default_graph_compiler.pypackages/nvidia_nat_app/tests/compiler/test_pipelined_compiler.pypackages/nvidia_nat_app/tests/conftest.pypackages/nvidia_nat_app/tests/constraints/__init__.pypackages/nvidia_nat_app/tests/constraints/test_decorators.pypackages/nvidia_nat_app/tests/constraints/test_models.pypackages/nvidia_nat_app/tests/constraints/test_resolution.pypackages/nvidia_nat_app/tests/executors/__init__.pypackages/nvidia_nat_app/tests/executors/test_execution_state.pypackages/nvidia_nat_app/tests/executors/test_metrics.pypackages/nvidia_nat_app/tests/executors/test_result_handler.pypackages/nvidia_nat_app/tests/executors/test_runner.pypackages/nvidia_nat_app/tests/graph/__init__.pypackages/nvidia_nat_app/tests/graph/conftest.pypackages/nvidia_nat_app/tests/graph/test_access.pypackages/nvidia_nat_app/tests/graph/test_adapter.pypackages/nvidia_nat_app/tests/graph/test_analysis.pypackages/nvidia_nat_app/tests/graph/test_factory.pypackages/nvidia_nat_app/tests/graph/test_llm_detection.pypackages/nvidia_nat_app/tests/graph/test_models.pypackages/nvidia_nat_app/tests/graph/test_optimizer.pypackages/nvidia_nat_app/tests/graph/test_scheduling.pypackages/nvidia_nat_app/tests/graph/test_static_analysis.pypackages/nvidia_nat_app/tests/graph/test_topology.pypackages/nvidia_nat_app/tests/graph/test_types.pypackages/nvidia_nat_app/tests/graph/test_uncertainty_invariants.pypackages/nvidia_nat_app/tests/speculation/__init__.pypackages/nvidia_nat_app/tests/speculation/test_plan.pypackages/nvidia_nat_app/tests/speculation/test_safety.pypackages/nvidia_nat_app/tests/stages/__init__.pypackages/nvidia_nat_app/tests/stages/test_edge_classification.pypackages/nvidia_nat_app/tests/stages/test_extract.pypackages/nvidia_nat_app/tests/stages/test_llm_analysis.pypackages/nvidia_nat_app/tests/stages/test_node_analysis.pypackages/nvidia_nat_app/tests/stages/test_priority_assignment.pypackages/nvidia_nat_app/tests/stages/test_scheduling.pypackages/nvidia_nat_app/tests/stages/test_topology.pypackages/nvidia_nat_app/tests/stages/test_validate.pypackages/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>
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (1)
packages/nvidia_nat_app/src/nat_app/api.py (1)
237-237: Unusedadaptervariable triggers Ruff RUF059.Both
find_parallel_stages(Line 237) andspeculative_opportunities(Line 440) unpackadapterbut 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 variablename— rename to_name.Ruff B007 flags
nameas 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, andplan_speculationare imported with# noqa: F401but 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: Conservativehas_dynamic_targetsflag on resolved-but-unknown receivers.When
_resolve_receiverreturns a concrete name that isn't in_llm_names(Line 343–344), the counter flagshas_dynamic_targets = True. This is overly cautious: a call likelogger.invoke("msg")would trigger partial confidence just because.invokematches an invocation method name. Consider only flagging when the receiver isNone(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
📒 Files selected for processing (13)
packages/nvidia_nat_app/src/nat_app/api.pypackages/nvidia_nat_app/src/nat_app/executors/runner.pypackages/nvidia_nat_app/src/nat_app/graph/analysis.pypackages/nvidia_nat_app/src/nat_app/graph/llm_detection.pypackages/nvidia_nat_app/src/nat_app/graph/static_analysis.pypackages/nvidia_nat_app/src/nat_app/graph/topology.pypackages/nvidia_nat_app/src/nat_app/graph/types.pypackages/nvidia_nat_app/tests/executors/test_runner.pypackages/nvidia_nat_app/tests/graph/test_analysis.pypackages/nvidia_nat_app/tests/graph/test_llm_detection.pypackages/nvidia_nat_app/tests/graph/test_topology.pypackages/nvidia_nat_app/tests/graph/test_types.pypackages/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>
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
packages/nvidia_nat_app/src/nat_app/api.py (1)
98-98: 🛠️ Refactor suggestion | 🟠 MajorPrefix unused
adapterwith underscore to silence RUF059.All three
build_graph_and_adaptercall sites unpackadapterbut 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 variablenameto_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 inmatch=to silence RUF043.The
.*metacharacters in thematch=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 repeatedExecutionStatesetup into a pytest fixture.The pattern
state = ExecutionState(); state.execution_start_time = 0.0is 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 stateThen 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 constructingStaticAnalysisResultdirectly instead of field-by-field copy.Both
_analyze_callee(lines 911–925) andanalyze_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) -> StaticAnalysisResulthelper 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, andplan_speculationwithnoqa: F401suppression. 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
📒 Files selected for processing (5)
packages/nvidia_nat_app/src/nat_app/api.pypackages/nvidia_nat_app/src/nat_app/executors/runner.pypackages/nvidia_nat_app/src/nat_app/graph/static_analysis.pypackages/nvidia_nat_app/tests/executors/test_runner.pypackages/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>
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (1)
packages/nvidia_nat_app/tests/executors/test_runner.py (1)
62-62:⚠️ Potential issue | 🟡 MinorOutstanding 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 inpyproject.toml.The quickest fix in test code is an inline
# noqa: TRY003suppression 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: TRY003Also 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 inmost.The PR title calls
nvidia-nat-appexperimental, yet it is included in themostaggregate extra. Anyone runninguv 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 frommostuntil it graduates out of experimental status — consistent with howopenpipe-artandragaaiare 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_nodealready matches the expectedCallable[[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 repeatedExecutionStatesetup into a pytest fixture.
ExecutionState()followed bystate.execution_start_time = 0.0appears ~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,nameargument 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
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
packages/nvidia_nat_app/src/nat_app/executors/runner.pypackages/nvidia_nat_app/src/nat_app/graph/static_analysis.pypackages/nvidia_nat_app/tests/executors/test_runner.pypackages/nvidia_nat_app/tests/graph/test_static_analysis.pypyproject.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>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (9)
ci/scripts/github/build_wheel.sh (1)
108-114: Add a success log after a passingnat_appimport.The
natbranch implicitly confirms success via thenat --versionoutput. Thenat_appbranch 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: Baresettype hints lose semantic information.
visited_funcsstoresid()return values (integers), but the annotation isset/set | None. Usingset[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_namedoes 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._freevarsvs.base_name not in param_names) is repeated nearly identically invisit_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 = TrueThen 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_attrsis provided but no regular params exist (class method with onlyself), the code creates a syntheticstate_param = "_self_state_placeholder"with an emptyeffective_param_to_obj. A named constant and an inline comment explaining why the param name is irrelevant (sinceself.stateaccess 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_sourceis 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 outertry/except SyntaxErrorcatches 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 sincezip(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
📒 Files selected for processing (4)
ci/scripts/github/build_wheel.shpackages/nvidia_nat_app/pyproject.tomlpackages/nvidia_nat_app/src/nat_app/graph/static_analysis.pypackages/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>
|
/merge |
…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
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
quick_optimize()takesnodesandedgesand returns parallel execution stages via a 6-stage pipeline: extract → validate → topology → node analysis → edge classification → scheduling.analyze_function()andclassify_edge()infer state reads/writes and edge necessity without executing code.plan_speculation(),speculative_opportunities(), andrun_speculation()support router-branch speculation with@speculation_unsafeandSpeculationSafetyConfig.benchmark()compares sequential vs. parallel execution and custom strategies.AbstractFrameworkAdapterandbuild_graph_and_adapter()for framework-specific integration; public API uses plain Python types.@sequential,@depends_on,@has_side_effects, andOptimizationConfigfor controlling optimization.API
quick_optimize|analyze_function|classify_edge|find_parallel_stages|benchmark|speculative_opportunities|plan_speculation|partition_targets|run_speculationNotes
Example
By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
Documentation
Tests