From ea59b0137aba0a8e894d3bb268954bf236914930 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Tue, 7 Jul 2026 10:31:33 -0700 Subject: [PATCH] feat(dynamo): opt-in lineage-scoped session affinity Add --dynamo-session-affinity-scope {conversation,lineage} (default conversation, no behavior change). With lineage, fork/spawn child conversations bind Dynamo session affinity with the agent-tree root correlation id so the whole lineage co-locates with the shared parent prefix; children never close the shared key (only the root conversation's final turn closes, TTL covers leftovers). --- src/aiperf/common/config/config_defaults.py | 1 + src/aiperf/common/config/endpoint_config.py | 18 ++++++++++++++++++ .../common/models/model_endpoint_info.py | 7 +++++++ src/aiperf/workers/inference_client.py | 15 +++++++++++++-- 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/src/aiperf/common/config/config_defaults.py b/src/aiperf/common/config/config_defaults.py index fd9e627341..4f476c5fae 100644 --- a/src/aiperf/common/config/config_defaults.py +++ b/src/aiperf/common/config/config_defaults.py @@ -49,6 +49,7 @@ class EndpointDefaults: DOWNLOAD_VIDEO_CONTENT = False REQUEST_CONTENT_TYPE = None USE_DYNAMO_CONV_AWARE_ROUTING = False + DYNAMO_SESSION_AFFINITY_SCOPE = "conversation" DYNAMO_SESSION_TIMEOUT_SECONDS = 300 # Readiness probe defaults. Timeout 0 disables the probe (the default); # any positive value enables it. Interval is only consulted when the diff --git a/src/aiperf/common/config/endpoint_config.py b/src/aiperf/common/config/endpoint_config.py index 7861dc5c84..04d5dcc0da 100644 --- a/src/aiperf/common/config/endpoint_config.py +++ b/src/aiperf/common/config/endpoint_config.py @@ -350,6 +350,24 @@ def url(self) -> str: ), ] = EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS + dynamo_session_affinity_scope: Annotated[ + Literal["conversation", "lineage"], + Field( + description=( + "Scope of the Dynamo session affinity key when " + "--use-dynamo-conv-aware-routing is enabled. 'conversation' " + "(default) binds each conversation with its own id. 'lineage' " + "binds fork/spawn child conversations with the agent-tree root " + "id so the whole lineage co-locates with the shared parent " + "prefix (for deployments without KV-event prefix indexing)." + ), + ), + CLIParameter( + name=("--dynamo-session-affinity-scope",), + group=Groups.ENDPOINT, + ), + ] = EndpointDefaults.DYNAMO_SESSION_AFFINITY_SCOPE + connection_reuse_strategy: Annotated[ ConnectionReuseStrategy, Field( diff --git a/src/aiperf/common/models/model_endpoint_info.py b/src/aiperf/common/models/model_endpoint_info.py index dbfe02e39c..6da86fccf3 100644 --- a/src/aiperf/common/models/model_endpoint_info.py +++ b/src/aiperf/common/models/model_endpoint_info.py @@ -123,6 +123,10 @@ class EndpointInfo(AIPerfBaseModel): ge=1, description="Timeout in seconds for Dynamo nvext.session_control sessions.", ) + dynamo_session_affinity_scope: str = Field( + default=EndpointDefaults.DYNAMO_SESSION_AFFINITY_SCOPE, + description="Scope of the Dynamo session affinity key ('conversation' or 'lineage').", + ) connection_reuse_strategy: ConnectionReuseStrategy = Field( default=EndpointDefaults.CONNECTION_REUSE_STRATEGY, description="Transport connection reuse strategy.", @@ -179,6 +183,9 @@ def from_user_config(cls, user_config: UserConfig) -> "EndpointInfo": dynamo_session_timeout_seconds=( user_config.endpoint.dynamo_session_timeout_seconds ), + dynamo_session_affinity_scope=( + user_config.endpoint.dynamo_session_affinity_scope + ), connection_reuse_strategy=user_config.endpoint.connection_reuse_strategy, download_video_content=user_config.endpoint.download_video_content, request_content_type=user_config.endpoint.request_content_type, diff --git a/src/aiperf/workers/inference_client.py b/src/aiperf/workers/inference_client.py index a0b4ee2c32..6d233028ac 100644 --- a/src/aiperf/workers/inference_client.py +++ b/src/aiperf/workers/inference_client.py @@ -145,11 +145,22 @@ async def _send_request_to_transport( # excluded by the dataset-load guard, so it is not handled here. endpoint = self.model_endpoint.endpoint if endpoint.use_dynamo_conv_aware_routing: + session_id = request_info.x_correlation_id + is_final_turn = request_info.is_final_turn + if endpoint.dynamo_session_affinity_scope == "lineage": + root = request_info.root_correlation_id + if root and root != request_info.x_correlation_id: + # Child conversations share the lineage affinity key; + # they must never close it while siblings may still + # run, so only the root conversation's final turn + # closes (the TTL reaper covers any leftovers). + session_id = root + is_final_turn = False formatted_payload = merge_session_control( formatted_payload, build_session_control( - session_id=request_info.x_correlation_id, - is_final_turn=request_info.is_final_turn, + session_id=session_id, + is_final_turn=is_final_turn, timeout_seconds=endpoint.dynamo_session_timeout_seconds, ), )