feat(routing): session-routing plugin layer for per-session router affinity#1127
feat(routing): session-routing plugin layer for per-session router affinity#1127ajcasagrande wants to merge 4 commits into
Conversation
…finity
Unify every mechanism for telling an external router which session a request
belongs to behind a single --session-routing flag backed by a new
session_routing plugin category. One mode per run, parameterized by repeatable
--session-routing-opt key=value pairs that populate a per-mode typed generic
Options model (extra="forbid"), coerced and canonicalized at config resolution
so downstream code always sees typed values.
Four modes:
- dynamo_headers : X-Dynamo-Session-ID on every request (plus
X-Dynamo-Parent-Session-ID on subagent children).
- dynamo_nvext : nvext.session_control body metadata -- bind on every
non-final turn, close on the final turn
(opt: timeout_seconds, default 300). Mutates the body,
so it is incompatible with any verbatim-bytes request path.
- smg_routing_key : X-SMG-Routing-Key for SGLang Model Gateway stickiness.
- session_id_header: an arbitrary configured header (opt: header_name).
The plugin is instantiated once per worker by InferenceClient and invoked at
the request-serialization chokepoint. Each transform receives a RoutingContext
of per-request lineage/finality facts stamped issuer-side from
SessionTreeRegistry state. Finality is deliberately conservative: is_tree_final
is False whenever indeterminate, is_parent_final is None for roots or when
unknown, and a has_branches SPAWN gate keeps any turn that will spawn
descendants on its return from ever stamping tree-final (SPAWN children register
only at return-intercept, AFTER issue-time stamping). on_session_end fires
strictly after a session's last worker-side activity on every terminal path
(final turn, cancellation, cancel-before-start via the done-callback path) and
is an idempotent no-op by default.
Behavior change: the exgentic loader no longer auto-stamps x-dynamo-session-id
on every turn. Run with --session-routing dynamo_headers to restore that
stamping (now available for any dataset and endpoint, not just exgentic).
A four-mode wire-level integration test drives a real aiperf profile subprocess
per mode against the in-repo mock server with --export-level raw and asserts the
exact per-record header / nvext.session_control lifecycle equals each session's
X-Correlation-ID, with 2-sessions x 3-turns integrity guards.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
ajcasagrande
requested review from
FrankD412,
debermudez,
ganeshku1,
ilana-n,
lkomali and
matthewkotila
as code owners
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@a865111e3f55f5d2aa87e6d8e9ee284d067f9e64Recommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@a865111e3f55f5d2aa87e6d8e9ee284d067f9e64Last updated for commit: |
WalkthroughThis PR adds session-aware routing plugins and per-tree finality tracking, then wires both through configuration, timing, workers, runtime, tests, and docs. ChangesSession Routing
Session Tree Finality
Estimated code review effort: 4 (Complex) | ~75 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tests/unit/plugin/test_session_routing_registry.py (1)
21-30: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep these new tests aligned with the repo's Python test conventions.
These functions are still missing explicit parameter/return annotations, and
SessionRoutingTypeis imported inside the test body instead of at module scope. As per coding guidelines, "Annotate all functions with type hints for every parameter and return value." and "Keep imports at the top of the file."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/plugin/test_session_routing_registry.py` around lines 21 - 30, The new tests in test_session_routing_registry should follow the repo’s Python conventions by adding explicit type hints to every test function parameter and return value, including the plugin lookup test’s name argument and both test functions’ return types. Move the SessionRoutingType import out of test_session_routing_enum_generated and place it at module scope with the other imports so the test file matches the project’s import style.Source: Coding guidelines
tests/unit/workers/test_worker.py (1)
303-440: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBring these new tests in line with the repo's Python test conventions.
The new test methods are missing explicit parameter/return annotations, and
Turnis imported inside the test bodies instead of at module scope. As per coding guidelines, "Annotate all functions with type hints for every parameter and return value." and "Keep imports at the top of the file."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/workers/test_worker.py` around lines 303 - 440, Update the new tests in TestSessionRoutingTerminalHooks and TestCreateRequestInfoFinality to follow the repo’s Python conventions by adding explicit type annotations to every test method parameter and return type, including helper methods like _done_callback. Move the Turn import out of the test bodies and place it at module scope with the other imports. Keep the existing test behavior intact while aligning the signatures and imports of test_release_and_evict_for_terminal_notifies_session_end, test_done_callback_not_returned_branch_notifies_session_end, test_done_callback_already_returned_does_not_double_notify, and the _create_request_info tests.Source: Coding guidelines
tests/unit/timing/test_session_tree_finality.py (1)
28-105: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winApply the repo's test typing/import conventions throughout this file.
The new test cases are missing explicit
-> Noneannotations, andRequestInfo/Creditare imported inside the test body instead of at the top of the file. As per coding guidelines, "Annotate all functions with type hints for every parameter and return value." and "Keep imports at the top of the file."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/timing/test_session_tree_finality.py` around lines 28 - 105, The new tests in this module should follow the repo’s typing and import conventions: add explicit return type annotations of None to every test function in this diff, including the helper-style tests like test_root_terminal_unknown_tree_is_none and test_finality_flows_credit_to_request_info. Move the RequestInfo and Credit imports out of test_finality_flows_credit_to_request_info and place them with the other top-level imports in the file. Keep the existing test names and assertions, and make sure any helpers referenced here still use the same symbols (_make_registry, _registry_with_tree, CreditPhase, RequestInfo, Credit) after the import cleanup.Source: Coding guidelines
tests/unit/timing/test_session_tree_wiring.py (1)
463-463: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
next()over list-comprehension + index for single-element lookup.Flagged by Ruff (RUF015).
♻️ Proposed fix
- root_credit0 = [c for c in router.sent if c.agent_depth == 0][0] + root_credit0 = next(c for c in router.sent if c.agent_depth == 0)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/timing/test_session_tree_wiring.py` at line 463, Replace the single-element lookup in the test with next() instead of building a list comprehension and indexing it. Update the root_credit0 assignment in test_session_tree_wiring to use next(...) over router.sent filtered by agent_depth == 0, keeping the same behavior while satisfying Ruff RUF015.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/aiperf/credit/issuer.py`:
- Around line 98-106: In the issuer initialization logic for
`session_tree_registry_enabled` and `self._session_tree_registry`, the force-on
path still falls back to `None` when no `session_tree_registry` is provided.
Update the constructor handling so `session_tree_registry_enabled=True` either
requires a registry or creates/attaches one explicitly, rather than silently
no-oping; keep the existing force-off behavior unchanged. Use the `Issuer`
initializer and `session_tree_registry_enabled` / `_session_tree_registry`
decision block to locate the fix.
---
Nitpick comments:
In `@tests/unit/plugin/test_session_routing_registry.py`:
- Around line 21-30: The new tests in test_session_routing_registry should
follow the repo’s Python conventions by adding explicit type hints to every test
function parameter and return value, including the plugin lookup test’s name
argument and both test functions’ return types. Move the SessionRoutingType
import out of test_session_routing_enum_generated and place it at module scope
with the other imports so the test file matches the project’s import style.
In `@tests/unit/timing/test_session_tree_finality.py`:
- Around line 28-105: The new tests in this module should follow the repo’s
typing and import conventions: add explicit return type annotations of None to
every test function in this diff, including the helper-style tests like
test_root_terminal_unknown_tree_is_none and
test_finality_flows_credit_to_request_info. Move the RequestInfo and Credit
imports out of test_finality_flows_credit_to_request_info and place them with
the other top-level imports in the file. Keep the existing test names and
assertions, and make sure any helpers referenced here still use the same symbols
(_make_registry, _registry_with_tree, CreditPhase, RequestInfo, Credit) after
the import cleanup.
In `@tests/unit/timing/test_session_tree_wiring.py`:
- Line 463: Replace the single-element lookup in the test with next() instead of
building a list comprehension and indexing it. Update the root_credit0
assignment in test_session_tree_wiring to use next(...) over router.sent
filtered by agent_depth == 0, keeping the same behavior while satisfying Ruff
RUF015.
In `@tests/unit/workers/test_worker.py`:
- Around line 303-440: Update the new tests in TestSessionRoutingTerminalHooks
and TestCreateRequestInfoFinality to follow the repo’s Python conventions by
adding explicit type annotations to every test method parameter and return type,
including helper methods like _done_callback. Move the Turn import out of the
test bodies and place it at module scope with the other imports. Keep the
existing test behavior intact while aligning the signatures and imports of
test_release_and_evict_for_terminal_notifies_session_end,
test_done_callback_not_returned_branch_notifies_session_end,
test_done_callback_already_returned_does_not_double_notify, and the
_create_request_info tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cf4de605-64c6-4027-bceb-3f277944e23a
📒 Files selected for processing (40)
docs/benchmark-datasets.mddocs/cli-options.mddocs/plugins/plugin-system.mdsrc/aiperf/common/models/model_endpoint_info.pysrc/aiperf/common/models/record_models.pysrc/aiperf/config/endpoint.pysrc/aiperf/config/flags/_converter_endpoint.pysrc/aiperf/config/flags/_section_fields.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/config/schema/aiperf-config.schema.jsonsrc/aiperf/credit/issuer.pysrc/aiperf/credit/structs.pysrc/aiperf/dataset/loader/exgentic.pysrc/aiperf/plugin/categories.yamlsrc/aiperf/plugin/enums.pysrc/aiperf/plugin/plugins.pysrc/aiperf/plugin/plugins.yamlsrc/aiperf/plugin/schema/plugins.schema.jsonsrc/aiperf/timing/_branch_orchestrator_spawn.pysrc/aiperf/timing/_branch_orchestrator_state.pysrc/aiperf/timing/branch_orchestrator.pysrc/aiperf/timing/conversation_source.pysrc/aiperf/timing/phase/runner.pysrc/aiperf/timing/session_tree.pysrc/aiperf/timing/strategies/fixed_schedule.pysrc/aiperf/timing/strategies/user_centric_rate.pysrc/aiperf/workers/inference_client.pysrc/aiperf/workers/session_routing.pysrc/aiperf/workers/worker.pytests/integration/test_session_routing_raw_export.pytests/unit/config/test_session_routing_config.pytests/unit/credit/test_issuer_finality.pytests/unit/dataset/loader/test_exgentic.pytests/unit/plugin/test_session_routing_registry.pytests/unit/timing/strategies/test_user_centric_rate.pytests/unit/timing/test_session_tree_finality.pytests/unit/timing/test_session_tree_wiring.pytests/unit/workers/test_inference_client.pytests/unit/workers/test_session_routing.pytests/unit/workers/test_worker.py
💤 Files with no reviewable changes (2)
- src/aiperf/dataset/loader/exgentic.py
- tests/unit/dataset/loader/test_exgentic.py
| self._session_tree_registry = ( | ||
| session_tree_registry | ||
| if ( | ||
| session_tree_registry_enabled | ||
| if session_tree_registry_enabled is not None | ||
| else session_tree_registry is not None | ||
| ) | ||
| else None | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
session_tree_registry_enabled=True silently no-ops when no registry is supplied.
Tracing the four (enabled, registry) combinations: forcing session_tree_registry_enabled=True without also passing session_tree_registry still resolves self._session_tree_registry to None (the true-branch just evaluates to session_tree_registry, i.e. None). This contradicts the docstring's "force it on/off" framing for the True case — force-off works, force-on silently doesn't, with no error surfaced. A test/caller trying to force-enable tree accounting without wiring a registry would fail silently.
🛡️ Proposed fix
+ if session_tree_registry_enabled and session_tree_registry is None:
+ raise ValueError(
+ "session_tree_registry_enabled=True requires a session_tree_registry instance"
+ )
self._session_tree_registry = (
session_tree_registry
if (
session_tree_registry_enabled
if session_tree_registry_enabled is not None
else session_tree_registry is not None
)
else None
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self._session_tree_registry = ( | |
| session_tree_registry | |
| if ( | |
| session_tree_registry_enabled | |
| if session_tree_registry_enabled is not None | |
| else session_tree_registry is not None | |
| ) | |
| else None | |
| ) | |
| if session_tree_registry_enabled and session_tree_registry is None: | |
| raise ValueError( | |
| "session_tree_registry_enabled=True requires a session_tree_registry instance" | |
| ) | |
| self._session_tree_registry = ( | |
| session_tree_registry | |
| if ( | |
| session_tree_registry_enabled | |
| if session_tree_registry_enabled is not None | |
| else session_tree_registry is not None | |
| ) | |
| else None | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/aiperf/credit/issuer.py` around lines 98 - 106, In the issuer
initialization logic for `session_tree_registry_enabled` and
`self._session_tree_registry`, the force-on path still falls back to `None` when
no `session_tree_registry` is provided. Update the constructor handling so
`session_tree_registry_enabled=True` either requires a registry or
creates/attaches one explicitly, rather than silently no-oping; keep the
existing force-off behavior unchanged. Use the `Issuer` initializer and
`session_tree_registry_enabled` / `_session_tree_registry` decision block to
locate the fix.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![dynamo-review-agent[bot]](https://avatars.githubusercontent.com/u/201626793?s=60&v=4){kind=small}
| # onto the structured body, endpoint-agnostic, after the payload dict is | ||
| # in hand. transform_body returns a copy, so this never mutates a cached | ||
| # Turn.raw_payload dict (the copy-on-write contract is load-bearing here). | ||
| if routing_ctx is not None and isinstance(payload, dict): |
There was a problem hiding this comment.
Body-mutating routing plugins are applied to Turn.raw_payload dicts, violating the raw payload/PAYLOAD_BYTES verbatim replay contract by silently changing authored request bodies. Fix: reject mutates_body routing when the current turn uses raw_payload or another verbatim payload path.
🤖 AI Fix
In src/aiperf/workers/inference_client.py, in InferenceClient._send_request_to_transport, before self._routing.transform_body(...), check raw_payload is not None and self._routing.mutates_body and raise a clear NotImplementedError or ValueError; add a test for dynamo_nvext plus raw_payload.
| # tree here is finality bookkeeping only -- no slot is released. | ||
| child_root = self._child_root.pop(child_corr, None) | ||
| if self._session_tree_registry is not None and child_root is not None: | ||
| self._session_tree_registry.on_descendant_done(child_root) |
There was a problem hiding this comment.
A terminal child decrements and can retire the session tree before intercept() registers branches declared on that same final child turn, so grandchildren spawned from a final child lose live tree state and get conservative/wrong finality stamps. Fix: register same-credit child branches before descendant completion retires the tree, or defer on_descendant_done until after interception.
🤖 AI Fix
In src/aiperf/timing/branch_orchestrator.py and src/aiperf/credit/callback_handler.py, change the child-final return flow so branch registration from BranchOrchestrator.intercept(credit) happens before _handle_child_done() calls SessionTreeRegistry.on_descendant_done(...) for that child.
… from main port) Backport of the finality-gate fix from the main-targeted session-routing port (ai-dynamo#1127, ref commit dd9bfdbc8). Defect: SessionTreeRegistry.is_last_tree_request gated on has_forks, which is FORK-only. SPAWN children register at return-intercept -- AFTER the issuer's _finality_for_issue stamped the credit -- so a final turn declaring only SPAWN branches could stamp a wrong is_tree_final=True while children are pending, violating the conservative contract. On this branch the SessionTreeRegistry is engaged only for AGENTIC_REPLAY (phase_orchestrator.py builds it iff is_agentic_replay), so the reactive dag_jsonl exposure is unreachable here (registry=None -> _finality_for_issue returns (None, False)) and agentic replay pre-registers recorded subagents. The gate is applied as defense-in-depth and to keep ajc/agentx-on-main convergent with the main port. - TurnToSend.has_branches (any-mode branch flag; has_forks stays FORK-only for the sticky router) derived from turn-metadata branch_ids and stamped at every construction seam: build_first_turn, build_turn_at_index, from_previous_credit, fixed_schedule turn-0, and both branch_orchestrator join seams via PendingBranchJoin.parent_has_branches_on_gated_turn. - is_last_tree_request kwarg renamed has_forks -> has_branches; _finality_for_issue and dispatch_join_turn pass it through. - Pre-session fold intentionally not ported: the unregistered dispatch_pre_session_branches path only runs where the registry is None. - Regression tests: SPAWN-declaring final turn never tree-final (unit + real issuer), metadata-driven build_first_turn stamping guard, positive control. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
ajcasagrande
left a comment
ajcasagrande
left a comment
There was a problem hiding this comment.
Self-review with independent verification — three fresh-context review passes (correctness/concurrency, config/conventions, tests/DX) plus runtime reproduction against the in-repo mock server (receipts kept on the branch under artifacts/repro-runtime-20260707/; wire evidence confirms all four modes behave as documented, including non-default opt plumbing timeout_seconds=123 -> bind {timeout: 123} / close).
Fix order:
- The exgentic fixed-schedule affinity-grouping behavior change needs an explicit decision — restore grouping via a dataset-provided affinity key, or document the regression (see inline comment).
mutates_bodydocstring overclaims an enforcement gate that does not exist on main.- Wrap
on_session_end/ per-request hook calls so a faulty plugin cannot skip session eviction or masquerade as server errors. - Conservative-only tree-retire ordering nit.
- Two test-strength nits.
What holds up well: no wrong-True is_tree_final path found under adversarial tracing (the has_branches gate + register-before-dispatch ordering held); copy-on-write body transforms protect shared raw_payload; canonicalized opts round-trip typed to workers; error messages are actionable; generated artifacts fully in sync; the wire-level integration test is deterministic and value-asserting.
| @@ -360,7 +360,6 @@ def _parse_span( | |||
| + _normalize_messages(attributes.get("gen_ai.input.messages")), | |||
| raw_tools=_normalize_tools(attributes.get("gen_ai.tool.definitions")), | |||
There was a problem hiding this comment.
Important — silent fixed-schedule behavior change. The removed stamp carried the shared source-session id across the N exploded single-turn conversations (the deleted test asserted exactly this: every exploded conversation had x-dynamo-session-id: session-1), grouping a recorded agent session under one Dynamo affinity key. --session-routing dynamo_headers keys on x_correlation_id, which fixed-schedule mints fresh per conversation (fixed_schedule.py:100) — so the spans of one recorded session no longer share an affinity key, and the migration advice in docs/benchmark-datasets.md does not restore that grouping.
The old behavior had its own flaw (a recorded id baked into the dataset collides across dataset-recycled replays), so neither is strictly better — but the change should be explicit, not silent. Suggest: document the behavior change here and in docs/benchmark-datasets.md, and consider a follow-up where loaders can supply an optional dataset lineage key on RoutingContext for modes to opt into (restores grouping without the replay-collision).
| class SessionRoutingBase(ABC, Generic[OptionsT]): # noqa: B024 # ABC marks the plugin protocol; every method has a working default so passthrough subclasses instantiate directly. | ||
| """Base for session-routing plugins (``session_routing`` category).""" | ||
|
|
||
| mutates_body: ClassVar[bool] = False |
There was a problem hiding this comment.
Minor — docstring overclaims. This says mutates_body gates the plugin off "the verbatim PAYLOAD_BYTES mmap fast path at dataset build, cache hit, and runtime" — none of that machinery exists on main, and the flag currently has zero consumers (grep-verified). docs/plugins/plugin-system.md already words it honestly ("protocol metadata that marks a mode as incompatible with any verbatim-bytes request path"); suggest matching that wording here and in categories.yaml/plugins.yaml, so a future bytes-path optimization doesn't trust a gate that isn't there and silently drop session_control.
| return state.outstanding <= 0 | ||
| return (not state.root_pending) and state.outstanding == 1 | ||
|
|
||
| def register_descendants(self, root_corr: str, n: int = 1) -> None: |
There was a problem hiding this comment.
Minor (conservative direction) — retire-before-register ordering. On a credit return, child-completion decrement (callback_handler.py step 4b) runs before intercept-registration (step 5). For root(done) -> C(last outstanding) whose final turn declares SPAWN grandchildren, the tree retires before they register; they buffer into _pending_descendants under a root that never reopens, so is_tree_final can never fire for that tree and the buffer entry lives until phase end (release_all only iterates _trees).
Direction is safely conservative — the has_branches gate independently keeps C's own stamp False — so not merge-blocking. Follow-up suggestion: decrement after intercept for child credits (mirroring the root's spawn-before-clear ordering in branch_orchestrator.py), and drain _pending_descendants in release_all.
| self.transport = TransportClass(model_endpoint=self.model_endpoint) | ||
| self.attach_child_lifecycle(self.transport) | ||
|
|
||
| def notify_session_end(self, x_correlation_id: str) -> None: |
There was a problem hiding this comment.
Minor — hook robustness + error attribution. notify_session_end is unwrapped and is the first statement of _release_and_evict_for_terminal, so a stateful plugin whose on_session_end raises (docs demand idempotency, not non-raising) skips session eviction and fork-pin release. Similarly, a plugin raising in headers()/transform_body() surfaces as an error record labeled "Error calling inference server API", hiding the actual culprit.
Suggest try/except-log around this pass-through (the hook should never break core lifecycle) and a routing-specific prefix when a plugin raises mid-request.
| assert registry.open_count() == 0 | ||
|
|
||
|
|
||
| def test_finality_flows_credit_to_request_info(): |
There was a problem hiding this comment.
Minor — docstring vs body. This claims to "catch a missed plumb touch," but it only asserts field-name presence on Credit/RequestInfo — deleting the plumb kwargs in _create_request_info keeps it green. The real value-level guard is test_worker.py::test_create_request_info_plumbs_finality_from_credit. Suggest fixing this docstring (or folding this test into that one).
|
|
||
| _NUM_SESSIONS = 2 | ||
| _TURNS_PER_SESSION = 3 | ||
| _TIMEOUT_SECONDS = 300 |
There was a problem hiding this comment.
Minor — assertion not load-bearing. _TIMEOUT_SECONDS = 300 equals DynamoNvextOptions.timeout_seconds' default, so the lifecycle assertion cannot distinguish "opt plumbed through CLI -> config -> worker" from "opt silently ignored." A manual run with timeout_seconds=123 confirms the plumb works end-to-end. Suggest a non-default value here so the test is load-bearing.
- exgentic(docs): document fixed-schedule Dynamo affinity grouping change (dynamo_headers keys on the live per-conversation id; each exploded span gets its own key; source-session grouping needs a dataset affinity key) - mutates_body: reword docstring/categories.yaml/plugins.yaml to honest semantics (protocol metadata, not enforced by any gate here today); regenerate plugins.schema.json - session_tree: resurrect a retired tree when a descendant's final-turn SPAWN registers grandchildren after that descendant's own on_descendant_done retired it, so the grandchild's genuinely-last credit can stamp is_tree_final; drain _pending_descendants/_retired_roots in release_all - inference_client: swallow+warn a raising on_session_end so eviction proceeds; attribute raising headers()/transform_body() to the routing plugin (not the server) in the error record - tests(finality): fix misleading schema-guard docstring; add resurrect + release_all-drain regressions; add hook-robustness + error-attribution tests - test(integration): non-default session-routing timeout (123) so the numeric plumb assertion is load-bearing Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
There was a problem hiding this comment.
♻️ Duplicate comments (1)
src/aiperf/workers/inference_client.py (1)
169-185: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftVerbatim
raw_payloadbodies are still silently overwritten bytransform_body.This mirrors a previously-raised concern: when
request_info.turns[-1].raw_payloadis adict, it's still passed straight intoself._routing.transform_body(payload, routing_ctx)at line 181 whenever routing is enabled, with no check for whether the routing modemutates_bodyand the turn is meant to be sent byte-for-byte. The new comment at lines 175-178 only explains that the cachedTurn.raw_payloadisn't mutated in place (copy-on-write) — that protects the source object, but the actual bytes sent on the wire for a verbatim-payload turn are still altered, which is the actual contract violation the earlier review flagged.If the incompatibility between body-mutating routing modes and verbatim payload paths is meant to be enforced (per the PR description "marks body-mutating routing modes as incompatible with verbatim-bytes request paths"), that enforcement isn't visible in this runtime code path — please confirm it's guarded at config/startup validation, otherwise a raw_payload turn using e.g.
dynamo_nvextwould have its body silently changed at send time.🔍 Verification script
#!/bin/bash # Look for a mutates_body flag/property on session routing plugins and any # validation that rejects it together with raw_payload usage. rg -n 'mutates_body' -C3 rg -n 'raw_payload' -C3 src/aiperf/workers/session_routing.py src/aiperf/config/endpoint.py 2>/dev/null ast-grep run --pattern 'class $NAME(SessionRoutingBase) { $$$ }' --lang python src/aiperf/workers/session_routing.py 2>/dev/null🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiperf/workers/inference_client.py` around lines 169 - 185, The verbatim raw_payload path in inference_client.py can still be modified by _routing.transform_body when routing_ctx is set, which breaks byte-for-byte request sending. Update the InferenceClient payload assembly flow so that request_info.turns[-1].raw_payload is exempt from body-mutating routing, or gate the transform behind the routing mode’s mutates_body flag. If the incompatibility is meant to be rejected, add/verify validation around the session-routing configuration so modes like dynamo_nvext cannot be used with verbatim payload turns.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@src/aiperf/workers/inference_client.py`:
- Around line 169-185: The verbatim raw_payload path in inference_client.py can
still be modified by _routing.transform_body when routing_ctx is set, which
breaks byte-for-byte request sending. Update the InferenceClient payload
assembly flow so that request_info.turns[-1].raw_payload is exempt from
body-mutating routing, or gate the transform behind the routing mode’s
mutates_body flag. If the incompatibility is meant to be rejected, add/verify
validation around the session-routing configuration so modes like dynamo_nvext
cannot be used with verbatim payload turns.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 66b32b6b-b56e-47a8-8c40-31a151739a75
📒 Files selected for processing (12)
docs/benchmark-datasets.mdsrc/aiperf/plugin/categories.yamlsrc/aiperf/plugin/plugins.yamlsrc/aiperf/plugin/schema/plugins.schema.jsonsrc/aiperf/timing/session_tree.pysrc/aiperf/workers/inference_client.pysrc/aiperf/workers/session_routing.pytests/integration/test_session_routing_raw_export.pytests/unit/timing/test_session_tree_finality.pytests/unit/timing/test_session_tree_wiring.pytests/unit/workers/test_inference_client.pytests/unit/workers/test_worker.py
✅ Files skipped from review due to trivial changes (1)
- docs/benchmark-datasets.md
🚧 Files skipped from review as they are similar to previous changes (9)
- src/aiperf/plugin/schema/plugins.schema.json
- src/aiperf/plugin/categories.yaml
- tests/unit/workers/test_worker.py
- src/aiperf/plugin/plugins.yaml
- tests/integration/test_session_routing_raw_export.py
- tests/unit/workers/test_inference_client.py
- tests/unit/timing/test_session_tree_wiring.py
- src/aiperf/workers/session_routing.py
- src/aiperf/timing/session_tree.py
…stness, docs, tests) Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
There was a problem hiding this comment.
Review: session-routing plugin layer
Overall: strong PR -- clean plugin abstraction, fail-fast typed config, and a genuinely careful conservative-finality design. Verified locally against the in-repo mock server: focused unit suite 161 passed, all four modes emit their exact documented wire format, copy-on-write body transform holds, and every config-validation branch behaves as specified. No correctness defects found. The two items below are a docstring-accuracy fix and a scope question -- neither blocks merge.
Suggested fix order:
- F2 (discussion, no code): decide whether the finality subsystem lands now or splits to a follow-up once a consumer exists.
Behavior change to call out in release notes: the exgentic loader no longer auto-stamps x-dynamo-session-id; migration is --session-routing dynamo_headers (now works for any dataset).
What's working well:
- Copy-on-write
transform_bodyis correct and asserted -- the cachedTurn.raw_payloaddict is never mutated. - Config validation is fail-fast and typed: opts-without-mode error,
extra="forbid"on every Options model, string->int coercion + canonicalization at resolution,ge=1bound enforced. - Finality is conservative-by-construction; the
has_branchesSPAWN gate correctly prevents a spawning final turn (including a single-turn root with pre-session SPAWN children) from ever stampingis_tree_final=Trueprematurely. on_session_endis fired idempotently on every terminal path (final turn, cancel, cancel-before-start).- Test depth: 161 focused unit tests + a four-mode wire-level integration test.
On test coverage -- emission vs. acceptance (see the inline note on test_session_routing_raw_export.py): the integration test reads request_headers/payload from AIPerf's raw export, so it proves AIPerf emits the correct format -- it does not exercise real Dynamo/SGLang acceptance (the mock server has zero references to session_control/X-SMG-Routing-Key/etc. and ignores them). That's expected for a client tool, and the golden assertions already guard AIPerf-side regressions. The gap is upstream contract drift (an external server renaming a field), which nothing in the repo would catch. A low-cost mitigation is suggested inline; a real end-to-end guard is optional and belongs behind a marker.
| @@ -0,0 +1,254 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
Finality subsystem has no in-tree consumer yet -- scope question, not a defect (discussion).
No shipped plugin reads is_tree_final, is_parent_final, or root_correlation_id:
grep -rn 'ctx\.\(is_tree_final\|is_parent_final\|root_correlation_id\)' src/ -> 0 hits. The four built-ins consume only ctx.x_correlation_id, ctx.parent_correlation_id, and the pre-existing ctx.is_final_turn.
Yet the PR ships the full SessionTreeRegistry (254 lines), issuer finality stamping, BranchOrchestrator wiring (~105 lines), and Credit/TurnToSend/RequestInfo plumbing to populate these facts -- roughly 600 LoC of conservative-finality machinery plus its correctness burden (SPAWN gate, pre-session fold) that no built-in exercises for its stated purpose.
This is clearly deliberate, well-tested, forward-looking extension surface -- so this is a judgment call, not a request. Worth deciding explicitly: land now alongside the plugins, or split the finality ledger into a follow-up once a consumer (e.g. a tree-final-aware routing mode) exists? Happy either way if it's intentional.
There was a problem hiding this comment.
I did this on purpose to avoid accidental breakage of a tested topology
| turn_index, keyed by X-Correlation-ID.""" | ||
| result = await cli.run(_build_cmd(url, mode=mode, opts=opts), timeout=300.0) | ||
|
|
||
| records = list(result.raw_records or []) |
There was a problem hiding this comment.
These assertions prove emission, not upstream acceptance -- consider a drift guard (low).
Everything checked here comes from result.raw_records, i.e. AIPerf's own record of what it sent. The mock server never inspects session_control, X-Dynamo-Session-ID, or X-SMG-Routing-Key (grep confirms zero references), so this is a golden test of AIPerf's output format, not proof that Dynamo/SGLang would accept it. That's the right thing to test for a client tool -- just flagging that "upstream renamed action/nvext.session_control" would sail past every test in the repo.
Ranked mitigations, most-aligned-with-the-apparatus first:
- Do now (cheap, high value): treat these assertions as a documented golden contract -- add a one-line comment per mode in
session_routing.pyciting the upstream spec/version each targets (the PR already notes "current upstream Dynamo main does not implementsession_control"). Makes drift a conscious decision instead of a silent surprise. - Optional, later: a periodic (nightly/weekly, not per-PR) e2e job against a real Dynamo/SGLang deployment, gated behind an opt-in marker like the existing
slow/performancedeselection. This is the only thing that actually catches upstream drift; it needs real infra and is flaky/expensive, so it belongs off the fast path -- consistent with how-m integration/sloware already handled here. - Avoid: teaching the mock server to emulate Dynamo/SGLang acceptance -- that just duplicates the golden format in a second place to keep in sync without testing the real server.
No change required for this PR to merge; (1) is a ~2-line addition worth doing while the format is fresh.
…finity Port of the session-routing feature (ai-dynamo#1127 / ajc/agentx-on-main 60bce6b + ad658fa + de4cd3f) onto the agentx-v1.0 line, adapted to this branch's older config system and worker structure, plus fixes from a 7-dimension adversarial fan-out review of the port. One plugin category (session_routing) + one CLI flag (--session-routing <mode>, with repeatable --session-routing-opt key=value) subsume every mechanism for telling an external router which session a request belongs to. "Session" (live instance keyed by x_correlation_id), not "conversation" (dataset template), is the affinity identity. Built-in modes: - dynamo_headers (preset): X-Dynamo-Session-ID plus X-Dynamo-Parent-Session-ID on subagent children. - dynamo_nvext: nvext.session_control bind/close for session_control-era Dynamo builds; option timeout_seconds (default 300). - smg_routing_key (preset): X-SMG-Routing-Key for the SGLang Model Gateway manual policy. - session_id_header (preset): additive X-Session-ID, parameterless. - identity_headers: fully generic tiered identity headers emitting exactly what is configured, nothing by default. Options (each a comma-separable name list): session (this session's correlation ID), parent (immediate parent's ID, omitted for roots), root (session-tree root's ID -- whole-tree affinity, e.g. pinning an entire agent tree to one replica for prefix-cache locality). At least one name across tiers; case-insensitive uniqueness ACROSS tiers (a name in two tiers would silently overwrite one tier's value with the other's). Every header preset is documented with its identity_headers equivalent. Architecture: - SessionRoutingBase[OptionsT]: generic typed protocol (headers / transform_body / on_session_end); per-plugin Pydantic Options models (extra=forbid). The CLI surface lives on the classic EndpointConfig (CLIParameter annotations): --session-routing-opt pairs are parsed, validated fail-fast, and canonicalized to typed values into a runtime-stamped session_routing_opts dict at config resolution. Commas inside opt values are preserved (wrap-only BeforeValidator, not parse_str_or_list, whose list branch comma-splits inside elements) and interpreted by the plugin's Options model. - RoutingContext carries x/parent/root correlation IDs plus issue-time lineage finality (is_parent_final, is_tree_final) stamped on credits from SessionTreeRegistry state; is_tree_final gates on the any-mode has_branches flag (NOT FORK-only has_forks) so a SPAWN-declaring final turn never stamps a wrong tree-final while children are pending. - on_session_end fires on all four worker terminal paths: final turn and cancellation (the _process_credit finally block), terminal context overflow (agentic_replay lane recycles send no final/cancel credit; classified via is_context_overflow_response), and cancel-before-start (done callback). Plugin exceptions in headers() are attributed to the plugin as error records; on_session_end exceptions are logged and swallowed so cleanup never breaks eviction. - Body-mutating plugins are gated off the verbatim PAYLOAD_BYTES mmap path everywhere it can arise on this branch: _preformat_payloads bails to the structured-turns path whenever a body-mutator (cache-bust or a mutates_body routing mode) is active, _select_mmap_format refuses natively-raw datasets, preformat-promoted PAYLOAD_BYTES cache entries are downgraded to a MISS at both lookup sites (a rebuild stays structured), and source-loaded entries hard-fail at cache-hit adoption with an actionable message. Removed (superseded): --use-dynamo-conv-aware-routing / --use-dynamo-session-control, --use-legacy-dynamo-session-control, --dynamo-session-timeout-seconds, the dynamo_session_control module, AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID, AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID, and the base_transports env-toggle header blocks. Also: Dynamo error-insight detector arms for both the unknown-variant bind and unknown-field session_control rejections; five-mode wire-level raw-export integration test (incl. identity_headers with session+root tiers); binding unit tests for every terminal-path notify, the overflow classifier, credit finality stamping, preformat bails, the cache-hit downgrade matrix, the identity_headers tier matrix, and preset equivalence; docs updated (plugin guide session-routing section, generated CLI/env docs). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…finity Port of the session-routing feature (ai-dynamo#1127 / ajc/agentx-on-main 60bce6b + ad658fa + de4cd3f) onto the agentx-v1.0 line, adapted to this branch's classic config system and worker structure, then hardened by two adversarial fan-out reviews (7 reviewers x 3-skeptic verification each; 24 confirmed findings applied) and performance-tuned with microbenchmarks. One plugin category (session_routing) + one CLI flag (--session-routing <mode>, with repeatable --session-routing-opt key=value) subsume every mechanism for telling an external router which session a request belongs to. "Session" (live instance keyed by x_correlation_id), not "conversation" (dataset template), is the affinity identity. Built-in modes: - dynamo_headers (preset): X-Dynamo-Session-ID plus X-Dynamo-Parent-Session-ID on subagent children. - dynamo_nvext: nvext.session_control bind/close for session_control-era Dynamo builds; option timeout_seconds (default 300). Plugin values win over dataset-shipped session_control keys. - smg_routing_key (preset): X-SMG-Routing-Key for the SGLang Model Gateway manual policy. - session_id_header (preset): additive X-Session-ID, parameterless. - identity_headers: fully generic tiered identity headers emitting exactly what is configured, nothing by default. Options (each a comma-separable name list): session (this session's correlation ID), parent (immediate parent's ID, omitted for roots), root (session-tree root's ID -- whole-tree affinity, e.g. pinning an entire agent tree to one replica for prefix-cache locality). At least one name across tiers; names validated as RFC 9110 tokens (config-time rejection of spaces/colons/CRLF) and case-insensitively unique across all tiers combined; list-form (YAML) inputs stripped identically to CLI comma-splits. Every header preset is documented with its identity_headers equivalent. Architecture: - SessionRoutingBase[OptionsT]: generic typed protocol (headers / transform_body / on_session_end); per-plugin Pydantic Options models (extra=forbid). CONTRACT enforced at worker init: a plugin overrides transform_body iff it declares mutates_body=True (fail-fast on split-brain plugins that would otherwise silently ship unrouted bodies on fast-path datasets). Capability detection recognizes both class-level overrides and instance-bound methods. - CLI surface on the classic EndpointConfig (CLIParameter annotations): --session-routing-opt pairs are parsed, validated fail-fast, and canonicalized to typed values into a runtime-stamped session_routing_opts dict at config resolution. Commas inside opt values are preserved (wrap-only BeforeValidator; parse_str_or_list's list branch comma-splits inside elements) and interpreted by the plugin's Options model. - RoutingContext (msgspec.Struct, frozen, gc=False) carries x/parent/ root correlation IDs plus issue-time lineage finality (is_parent_final, is_tree_final) stamped on credits from SessionTreeRegistry state; is_tree_final gates on the any-mode has_branches flag (NOT FORK-only has_forks) so a SPAWN-declaring final turn never stamps a wrong tree-final while children are pending; has_branches is stamped and test-bound at every TurnToSend construction seam. - on_session_end fires on every terminal disposition and only terminal dispositions: final turn and cancellation (the _process_credit finally block), terminal context overflow (gated to agentic_replay, the only strategy that recycles the lane; classified via is_context_overflow_response), cancel-before-start (done callback, gated on cancelled so a failed CreditReturn send on a non-final turn does not fire mid-session), and a worker-stop sweep for sessions still resident at phase end. Plugin exceptions in headers()/transform_body() are attributed to the plugin as error records; on_session_end exceptions are logged and swallowed so cleanup never breaks eviction. - Body-mutating plugins are gated off the verbatim PAYLOAD_BYTES mmap path everywhere it can arise: _preformat_payloads bails to the structured-turns path whenever a body-mutator (cache-bust or a mutates_body routing mode) is active, _select_mmap_format refuses natively-raw datasets, preformat-promoted PAYLOAD_BYTES cache entries are downgraded to a MISS at both lookup sites (a rebuild stays structured), source-loaded entries hard-fail at cache-hit adoption with an actionable message, and body-mutator runs skip post-run cache population so the shared cache key (which excludes routing settings) never pins later feature-free runs off the fast path. Performance (microbenchmarked, best-of-7 x 200k, real classes): RoutingContext build 511 -> 63 ns (frozen slots dataclass -> msgspec.Struct; frozen-dataclass __init__ routes object.__setattr__ per field); full per-request routing block 3.0-3.6x faster across modes (dynamo_headers 669 -> 186 ns, dynamo_nvext 903 -> 300 ns) via init-time capability flags that skip no-op base-method calls, an nvext-absent fast path in transform_body, and init-time snapshots of pydantic options. Measured non-wins deliberately rejected: cached bound-method pointers (defeat CPython 3.12's adaptive specialization), dict.fromkeys at small N, prebuilt-template merges. Removed (superseded): --use-dynamo-conv-aware-routing / --use-dynamo-session-control, --use-legacy-dynamo-session-control, --dynamo-session-timeout-seconds, the dynamo_session_control module, AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID, AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID, and the base_transports env-toggle header blocks. Verification: six-mode wire-level raw-export integration matrix (incl. identity_headers with session+root tiers) against the in-repo mock server, plus live end-to-end runs of every mode confirming per-turn headers / bind-close lifecycle / body-untouched invariants on the wire; binding unit tests for every terminal-path notify (positive and negative), the overflow classifier, capability flags and skip semantics, both plugin-contract init rejections, credit finality stamping at all construction seams, preformat bails, the cache-hit downgrade matrix and its lookup-site wiring, cache-population skip, identity_headers tier/validation matrix, and preset equivalence; docs updated across the plugin guide, generated CLI/env docs, and README. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
2 participants
Summary
Session affinity in front of a multi-replica inference deployment has, until
now, meant hand-editing dataset headers or patching a loader. This PR
consolidates every "tell the router which session this request belongs to"
mechanism behind a single
--session-routingflag backed by a newsession_routingplugin category. One mode per run; parameterize it withrepeatable
--session-routing-opt key=value.Four modes:
dynamo_headers— stampsX-Dynamo-Session-IDon every request (plusX-Dynamo-Parent-Session-IDon subagent children) for a Dynamo frontendrunning
--router-session-affinity-ttl-secs.dynamo_nvext— writesnvext.session_controlbody metadata:bindonevery non-final turn,
closeon the final turn (--session-routing-opt timeout_seconds=…, default 300). Only for Dynamo builds that implementsession_control.smg_routing_key— emitsX-SMG-Routing-Keyfor SGLang Model Gatewaystickiness.
session_id_header— emits an arbitrary configured header(
--session-routing-opt header_name=X-Affinity).Example commands:
The selected plugin is instantiated once per worker by
InferenceClientandinvoked at the request-serialization chokepoint.
Design notes
Options. Each mode declares its own PydanticOptionsmodel (
extra="forbid", so unknown opt keys are rejected at startup).Repeated
--session-routing-optpairs populate it; values are coerced to themodel's field types and canonicalized at config resolution, so downstream
code always sees typed values.
--session-routing-optwithout--session-routingis an error, and parameterless modes reject every opt key.mutates_body. A plugin setsmutates_body = Truewhen it rewrites thepayload (only
dynamo_nvextdoes). It marks the mode as incompatible with anyverbatim-bytes request path; header-only modes leave the body untouched.
transform_bodynever mutates its input (the request path shares cachedTurn.raw_payloaddicts with the dataset) — it returns a new dict.RoutingContextfinality facts. Each transform receives per-requestlineage/finality facts stamped issuer-side from
SessionTreeRegistrystate.The contract is deliberately conservative:
is_tree_finalisFalsewhenever indeterminate,
is_parent_finalisNonefor roots or when unknown,and a
has_branchesSPAWN gate keeps any turn that will spawn descendants onits return from ever stamping tree-final — SPAWN children register only at
return-intercept, after issue-time stamping, so a turn that has not yet
spawned its children can never be treated as the tree's provably-last request.
on_session_endcontract. Fires strictly after a session's lastworker-side activity, on every terminal path (final turn, cancellation, and
cancel-before-start via the done-callback path). It is post-session cleanup
only and must be idempotent (default no-op).
Behavior changes
x-dynamo-session-idon everyturn.
--session-routing dynamo_headersto restore thestamping. It now applies to any dataset and endpoint, not just exgentic, and
emits the canonical
X-Dynamo-Session-IDheader keyed on the session'scorrelation id.
Testing
registry, session-tree finality + wiring, credit-issuer finality stamping,
worker / inference-client plumbing, and the routing-plugin transforms.
(
tests/integration/test_session_routing_raw_export.py): runs a realaiperf profilesubprocess per mode against the in-repo mock server with--export-level raw, then asserts the exact per-record header /nvext.session_controllifecycle equals each session'sX-Correlation-ID,with 2-sessions × 3-turns integrity guards (session count, turn count, stable
session id, header-mode-does-not-mutate-body).
TurnToSendbuilt in the user-centric-rate return path carrieshas_branches=Trueiff the next turn's metadata declaresbranch_ids, so abranching turn is never mis-stamped as a leaf.
uv run pytest -n auto tests/unit/→ 13894 passed,make validate-plugin-schemasandpre-commit run --all-filesboth green.🤖 Generated with Claude Code
Summary by CodeRabbit
--session-routingand repeatable--session-routing-optfor profiles and services, including new session-routing plugin modes and configuration validation.