feat(workers): Dynamo conversation-aware routing via nvext.session_control#1101
Draft
ajcasagrande wants to merge 5 commits into
Draft
feat(workers): Dynamo conversation-aware routing via nvext.session_control#1101ajcasagrande wants to merge 5 commits into
ajcasagrande wants to merge 5 commits into
Conversation
…ntrol Opt-in --use-dynamo-conv-aware-routing emits nvext.session_control in OpenAI-compatible request bodies so a Dynamo frontend can pin every turn of a replayed conversation to the same backend worker (reusing its prefix KV cache). Sticky sessions are keyed by X-Correlation-ID. Two wire contracts: - Modern (default, Dynamo >= v1.3.0-dev): action 'bind' on every non-final turn (idempotent, refreshes the router TTL), 'close' on the final turn so affinity is released instead of leaking to the TTL reaper (material at millions of sessions). - Legacy (--use-legacy-dynamo-session-control, Dynamo v1.2.x): 'open' exactly once per session, then 'close' on the final turn. The client tracks opened sessions and exposes discard_dynamo_session() for abandonment paths. Policy lives in the pure-function dynamo_session_control module; the overlay happens at the serialization chokepoint after the endpoint formats the request dict, so it is endpoint-agnostic and never mutates a cached Turn. The API-error console insight explains the HTTP 400 'unknown variant bind' rejection from pre-v1.3.0 Dynamo builds and points at the legacy flag. 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@167fb30a03f2317e12c1911cc7a200541477aad7Recommended 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@167fb30a03f2317e12c1911cc7a200541477aad7Last updated for commit: |
This comment has been minimized.
This comment has been minimized.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
WalkthroughAdds Dynamo session-control configuration, schema, and CLI documentation; request payload helpers and client wiring for modern and legacy routing; a console detector for unsupported bind errors; and tests covering config, payload, client, and detector behavior. ChangesDynamo Session-Control Feature
Estimated code review effort: 4 (Complex) | ~45 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
src/aiperf/common/models/model_endpoint_info.py (1)
230-238: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGetattr fallbacks duplicate
EndpointDefaultsliterals.
False,False,300here should ideally referenceEndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING,EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL, andEndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDSto avoid silent drift if those defaults ever change. This mirrors the existing pattern for othergetattr(ep, ..., <literal>)calls in this method, so it's a pre-existing style rather than a new regression.♻️ Proposed fix
- use_dynamo_conv_aware_routing=getattr( - ep, "use_dynamo_conv_aware_routing", False - ), - use_legacy_dynamo_session_control=getattr( - ep, "use_legacy_dynamo_session_control", False - ), - dynamo_session_timeout_seconds=getattr( - ep, "dynamo_session_timeout_seconds", 300 - ), + use_dynamo_conv_aware_routing=getattr( + ep, + "use_dynamo_conv_aware_routing", + EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING, + ), + use_legacy_dynamo_session_control=getattr( + ep, + "use_legacy_dynamo_session_control", + EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL, + ), + dynamo_session_timeout_seconds=getattr( + ep, + "dynamo_session_timeout_seconds", + EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS, + ),🤖 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/common/models/model_endpoint_info.py` around lines 230 - 238, The getattr fallbacks in the model endpoint mapping are hardcoded literals and should use the corresponding EndpointDefaults values instead. Update the `model_endpoint_info` conversion logic so `use_dynamo_conv_aware_routing`, `use_legacy_dynamo_session_control`, and `dynamo_session_timeout_seconds` fall back to `EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING`, `EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL`, and `EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS`, matching the existing style used elsewhere in this method.src/aiperf/config/flags/cli_config.py (1)
392-410: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAliased flag name may read as a standalone "session control" toggle.
--use-dynamo-session-controlis an alias foruse_dynamo_conv_aware_routing, while a separate--use-legacy-dynamo-session-controlflag exists for the wire-contract variant. The similar naming could lead users to believe these are two independent options rather than one enabling the feature and the other selecting its wire format. Not a functional issue sincevalidate_dynamo_session_control_coherentinendpoint.pyalready rejects the invalid combination.🤖 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/config/flags/cli_config.py` around lines 392 - 410, The alias name in use_dynamo_conv_aware_routing can be mistaken for a separate standalone session-control toggle, so update the CLI flag naming/description to make it clear that --use-dynamo-session-control is only an alias for the same feature and that --use-legacy-dynamo-session-control is the wire-format variant. Adjust the CLIParameter definition and Field description in cli_config.py to explicitly distinguish the feature toggle from the legacy contract option, so users can tell they are not independent settings.src/aiperf/workers/inference_client.py (1)
126-126: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFeature silently no-ops for non-dict payloads.
When
use_dynamo_conv_aware_routingis enabled butpayloadisn't a dict (e.g. multipart/form-data endpoints), the session_control overlay is silently skipped with no log, which could confuse users who enabled the flag expecting it to apply.🤖 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` at line 126, The Dynamo conversation-aware routing overlay is skipped silently when payload is not a dict, so add explicit handling in inference_client.py around the use_dynamo_conv_aware_routing check in the inference_client path. In the logic that applies the session_control overlay, keep the existing dict-based branch but add a clear log when the flag is enabled and payload is a non-dict type so users know the feature was intentionally not applied. Make sure the relevant symbols like use_dynamo_conv_aware_routing and session_control are referenced so the behavior is easy to find and adjust.
🤖 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/workers/inference_client.py`:
- Around line 129-149: The legacy session tracking in _send_request_internal
updates _dynamo_opened_sessions before transport.send_request completes, so a
failed send can leave the session marked open even though the server never
received it. Track whether this call just added session_id in the legacy/open
path around build_session_control and merge_session_control, and if
self.transport.send_request raises, roll back that optimistic add by removing
the session_id from _dynamo_opened_sessions before rethrowing. Keep the
final-turn discard behavior unchanged so the set only reflects opens that were
actually sent successfully.
- Around line 151-160: The `discard_dynamo_session` helper is currently unused,
so abandoned sessions can remain in `_dynamo_opened_sessions` and the legacy set
can grow. Wire `discard_dynamo_session` into the
session-abandonment/cancellation path in `InferenceClient`, using the existing
`_send_request_to_transport` flow as the reference point, so sessions that stop
before a final turn are explicitly removed from the legacy tracking set.
---
Nitpick comments:
In `@src/aiperf/common/models/model_endpoint_info.py`:
- Around line 230-238: The getattr fallbacks in the model endpoint mapping are
hardcoded literals and should use the corresponding EndpointDefaults values
instead. Update the `model_endpoint_info` conversion logic so
`use_dynamo_conv_aware_routing`, `use_legacy_dynamo_session_control`, and
`dynamo_session_timeout_seconds` fall back to
`EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING`,
`EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL`, and
`EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS`, matching the existing style
used elsewhere in this method.
In `@src/aiperf/config/flags/cli_config.py`:
- Around line 392-410: The alias name in use_dynamo_conv_aware_routing can be
mistaken for a separate standalone session-control toggle, so update the CLI
flag naming/description to make it clear that --use-dynamo-session-control is
only an alias for the same feature and that --use-legacy-dynamo-session-control
is the wire-format variant. Adjust the CLIParameter definition and Field
description in cli_config.py to explicitly distinguish the feature toggle from
the legacy contract option, so users can tell they are not independent settings.
In `@src/aiperf/workers/inference_client.py`:
- Line 126: The Dynamo conversation-aware routing overlay is skipped silently
when payload is not a dict, so add explicit handling in inference_client.py
around the use_dynamo_conv_aware_routing check in the inference_client path. In
the logic that applies the session_control overlay, keep the existing dict-based
branch but add a clear log when the flag is enabled and payload is a non-dict
type so users know the feature was intentionally not applied. Make sure the
relevant symbols like use_dynamo_conv_aware_routing and session_control are
referenced so the behavior is easy to find and adjust.
🪄 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: aaaa1904-ff1b-4c64-95c1-65bf8416d45e
📒 Files selected for processing (10)
docs/cli-options.mdsrc/aiperf/common/models/model_endpoint_info.pysrc/aiperf/config/endpoint.pysrc/aiperf/config/flags/_converter_endpoint.pysrc/aiperf/config/flags/_section_fields.pysrc/aiperf/config/flags/cli_config.pysrc/aiperf/exporters/console_api_error_exporter.pysrc/aiperf/workers/dynamo_session_control.pysrc/aiperf/workers/inference_client.pytests/unit/workers/test_dynamo_session_control.py
…control fields The generate-config-schema pre-commit hook failed in CI because the new use_dynamo_conv_aware_routing, use_legacy_dynamo_session_control, and dynamo_session_timeout_seconds EndpointConfig fields were not reflected in the generated schema. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
…insight Codecov flagged the patch-coverage gaps in PR #1101. Add targeted unit tests for: - InferenceClient._send_request_to_transport overlay: modern bind on non-final turns, close on final turns, the legacy open -> bare session_id -> close lifecycle with per-session open tracking, preservation of unrelated nvext keys, the routing-disabled and non-dict-payload guards, and discard_dynamo_session idempotency. - EndpointConfig.validate_dynamo_session_control_coherent error path (legacy flag without conv-aware routing) plus accepted combinations and defaults. - DynamoSessionControlDetector: JSON-wrapped, raw, and JSON-non-dict 'unknown variant bind' messages, unrelated/empty/None summaries, and the exporter panel rendering for the insight. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unit/exporters/test_console_api_error_exporter.py (1)
131-137: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
orjson.dumpsinstead of stdlibjson.dumps.Test fixtures build the mock error
messagewithjson.dumps(...). Per coding guidelines, JSON serialization in this codebase should always useorjson.dumps/orjson.loads. Noteorjson.dumpsreturnsbytes, so callers here would need.decode()to keepmessageastr.As per coding guidelines, "Always use
orjson.loads(s),orjson.dumps(d)for JSON serialization/deserialization".♻️ Example fix
- json.dumps( + orjson.dumps( { "message": "Failed to deserialize the JSON body into the target type: " "nvext.session_control.action: unknown variant `bind`, " "expected `open` or `close` at line 1 column 100" } - ), + ).decode(),Also applies to: 145-145, 200-202
🤖 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/exporters/test_console_api_error_exporter.py` around lines 131 - 137, Replace the stdlib json.dumps usage in the test fixture with orjson.dumps and decode the resulting bytes back to a string so the mock message remains a str. Update the affected test cases in test_console_api_error_exporter.py where the error message is constructed, keeping the same message content but switching serialization to orjson per the project’s JSON handling guidelines.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tests/unit/exporters/test_console_api_error_exporter.py`:
- Around line 131-137: Replace the stdlib json.dumps usage in the test fixture
with orjson.dumps and decode the resulting bytes back to a string so the mock
message remains a str. Update the affected test cases in
test_console_api_error_exporter.py where the error message is constructed,
keeping the same message content but switching serialization to orjson per the
project’s JSON handling guidelines.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f268dc17-b022-45dc-9f29-c23d04db0129
📒 Files selected for processing (4)
src/aiperf/config/schema/aiperf-config.schema.jsontests/unit/config/test_endpoint.pytests/unit/exporters/test_console_api_error_exporter.pytests/unit/workers/test_inference_client.py
…al eviction) Cancelled/abandoned sessions never dispatch their final turn, so the inline legacy-'open' discard in InferenceClient cannot fire for them; Worker._release_and_evict_for_terminal now drops the session from the legacy Dynamo open-tracking set alongside the SessionManager eviction, keeping the set bounded on cancellation-heavy runs. Idempotent and a no-op for the modern (bind) path; the non-final non-cancelled warmup turn never hits this path, so cross-phase 'open' retention is preserved. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
ajcasagrande
added a commit
that referenced
this pull request
Jul 7, 2026
…insight Port of 154ac54 (PR #1101, ajc/dynamo-conv-aware-routing) with agentx adaptations: - _sent_payload decodes the orjson bytes agentx canonicalises dict payloads into before the transport call (the carve asserted on the raw dict). - Restore the isinstance(formatted_payload, dict) guard on the Dynamo session-control overlay in InferenceClient._send_request_to_transport. The carve had it (73e8046) but agentx's payload-bytes refactor of the block dropped it; the ported non-dict-payload test caught the regression (merge_session_control raises ValueError on str payloads). (cherry picked from commit 154ac54) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
ajcasagrande
added a commit
that referenced
this pull request
Jul 7, 2026
…acking Test-only port of b604173 (PR #1101, ajc/dynamo-conv-aware-routing). The src fix (Worker._release_and_evict_for_terminal calling inference_client.discard_dynamo_session before SessionManager eviction) was already present on agentx, which additionally routes the terminal-context-overflow disposition through the same discard. (cherry picked from commit b604173, test hunk only) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
The carve re-authored this slice's commits and dropped the co-author trailer; carry it forward so the squash merge credits the original work. Co-Authored-By: weireweire <20922698+weireweire@users.noreply.github.com> Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Opt-in
--use-dynamo-conv-aware-routingemitsnvext.session_controlin OpenAI-compatible request bodies so a Dynamo frontend can pin every turn of a replayed conversation to the same backend worker (reusing its prefix KV cache). Sticky sessions are keyed by the stableX-Correlation-ID.Two wire contracts:
action: bindon every non-final turn — idempotent and refreshes the router TTL so long inter-turn delays can't silently expire affinity — andcloseon the final turn so affinity is released immediately instead of leaking to the TTL reaper (material on high-session-cardinality runs).--use-legacy-dynamo-session-control, Dynamo v1.2.x):openexactly once per session,closeon the final turn; the client tracks opened sessions (openis not idempotent) and exposesdiscard_dynamo_session()for abandonment paths.Design: policy is a pure-function module (
build_session_control/merge_session_control); the overlay happens at the serialization chokepoint after the endpoint formats the request dict — endpoint-agnostic, never mutates a cached Turn. A new API-error console insight explains the HTTP 400unknown variant \bind`` rejection from pre-v1.3.0 Dynamo and points at the legacy flag.All off by default; zero behavior change unless opted in. CLI docs regenerated.
Testing
pytest tests/unit/workers/ tests/unit/config/ tests/unit/exporters/ tests/unit/transports/— 2498 passed.Summary by CodeRabbit
nvext.session_controlinto requests when enabled, handling modern bind/close and legacy open/close lifecycle behavior.binddue to unsupported session-control behavior.