feat(dynamo): opt-in lineage-scoped session affinity

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

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(

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,

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,
                     ),
                 )