diff --git a/docs/benchmark-datasets.md b/docs/benchmark-datasets.md index f0d2f3f4e4..6ee1a3ec00 100644 --- a/docs/benchmark-datasets.md +++ b/docs/benchmark-datasets.md @@ -67,7 +67,9 @@ This document describes datasets that AIPerf can use to generate stimulus. Addit ## Exgentic Agent Trace Replay -The Exgentic loaders stream recorded agent sessions directly from Hugging Face. `exgentic` is pinned to v1 revision `70036b93a04e61b0ea2706a68b962f4f26774587`; `exgentic_v2` is independently pinned to v2 revision `4b8ad4ab198438e5a170f9171c19c6a2cf7c1814`. Each replays successful, positive-token chat call snapshots. Recorded messages, system instructions, tool definitions, output-token limits, request controls, and call start times are preserved. Tools are not executed, and live responses are not added to later requests. Every request carries the source session as `x-dynamo-session-id` for Dynamo agentic tracing while AIPerf retains its own request correlation ID. +The Exgentic loaders stream recorded agent sessions directly from Hugging Face. `exgentic` is pinned to v1 revision `70036b93a04e61b0ea2706a68b962f4f26774587`; `exgentic_v2` is independently pinned to v2 revision `4b8ad4ab198438e5a170f9171c19c6a2cf7c1814`. Each replays successful, positive-token chat call snapshots. Recorded messages, system instructions, tool definitions, output-token limits, request controls, and call start times are preserved. Tools are not executed, and live responses are not added to later requests. The loaders no longer stamp a session header themselves; AIPerf keeps a stable per-conversation correlation ID, and pairing the run with `--session-routing dynamo_headers` sends `X-Dynamo-Session-ID` for Dynamo session affinity and agentic tracing. + +This changes fixed-schedule affinity grouping. `--session-routing dynamo_headers` keys affinity on the live per-conversation correlation ID, and in `--fixed-schedule` replay the loader explodes one recorded agent session into many single-turn span-conversations — so each span now gets its own key and routes independently, and Dynamo prefix-affinity no longer groups a recorded session's spans onto one worker. The removed stamp instead put the shared recorded session ID on every span, grouping them (but that recorded ID also collided across dataset-recycled replays, binding unrelated live conversations to a single router entry). The two behaviors serve different fidelity goals; restoring source-session grouping without the recycling collision would need a dataset-provided affinity key (possible future work). Provide a finite materialization bound through `--num-conversations`, `--num-dataset-entries`, or `--request-count`. `--benchmark-duration` limits request issuance, not dataset setup. diff --git a/docs/cli-options.md b/docs/cli-options.md index 8905f0bf94..7f5dd8fe7c 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -379,6 +379,15 @@ Content type for request body serialization. By default, requests are sent as 'a HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`). +#### `--session-routing` `` + +Session-aware routing mode: stamps per-session identity on every request for router affinity. Built-ins: dynamo_headers (X-Dynamo-Session-ID + parent header), dynamo_nvext (nvext.session_control bind/close request-body metadata), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (custom additive header). Parameterize with --session-routing-opt. +
_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`]_ + +#### `--session-routing-opt` `` + +Repeatable key=value option for the selected --session-routing mode (e.g. --session-routing-opt timeout_seconds=600), validated against the plugin's Options model. + ### Tokenizer #### `--tokenizer` `` @@ -1614,7 +1623,7 @@ Explore AIPerf plugins: aiperf plugins [category] [type] #### `--category` `` Category to explore. -
_Choices: [`accuracy_benchmark`, `accuracy_grader`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `convergence_criterion`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `gpu_telemetry_processor`, `plot`, `public_dataset_loader`, `ramp`, `record_processor`, `results_processor`, `search_planner`, `search_recipe`, `search_recipe_post_process`, `server_metrics_processor`, `service`, `service_manager`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_ +
_Choices: [`accuracy_benchmark`, `accuracy_grader`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `convergence_criterion`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `gpu_telemetry_processor`, `plot`, `public_dataset_loader`, `ramp`, `record_processor`, `results_processor`, `search_planner`, `search_recipe`, `search_recipe_post_process`, `server_metrics_processor`, `service`, `service_manager`, `session_routing`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_ #### `--name` `` @@ -1776,6 +1785,15 @@ Content type for request body serialization. By default, requests are sent as 'a HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`). +#### `--session-routing` `` + +Session-aware routing mode: stamps per-session identity on every request for router affinity. Built-ins: dynamo_headers (X-Dynamo-Session-ID + parent header), dynamo_nvext (nvext.session_control bind/close request-body metadata), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (custom additive header). Parameterize with --session-routing-opt. +
_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`]_ + +#### `--session-routing-opt` `` + +Repeatable key=value option for the selected --session-routing mode (e.g. --session-routing-opt timeout_seconds=600), validated against the plugin's Options model. + ### Tokenizer #### `--tokenizer` `` diff --git a/docs/plugins/plugin-system.md b/docs/plugins/plugin-system.md index b26be93ed6..db9752fe2e 100644 --- a/docs/plugins/plugin-system.md +++ b/docs/plugins/plugin-system.md @@ -100,7 +100,7 @@ for entry, cls in plugins.iter_all(PluginType.ENDPOINT): ## Plugin Categories -AIPerf supports 33 plugin categories organized by function, including `api_router` and `public_dataset_loader`: +AIPerf supports 34 plugin categories organized by function, including `api_router` and `public_dataset_loader`: ### Timing Categories @@ -129,6 +129,50 @@ AIPerf supports 33 plugin categories organized by function, including `api_route | `endpoint` | `EndpointType` | API endpoint implementations (chat, completions, embeddings, etc.) | | `transport` | `TransportType` | Network transport (HTTP via aiohttp) | +### Session Routing Category + +| Category | Enum | Description | +|----------|------|-------------| +| `session_routing` | `SessionRoutingType` | Stamps per-session identity (headers or body metadata) onto outbound requests so an external router pins every turn of a session to one worker; selected via `--session-routing` | + +**Purpose.** A session-routing plugin gives an external router (SGLang Model Gateway, Dynamo, a +generic session-affinity load balancer) the identity of the session a request belongs to, so all of +a conversation's turns re-land on the replica holding its KV prefix. Exactly one mode runs per +invocation; the selected plugin is instantiated once per worker by `InferenceClient` and invoked at +the request-serialization chokepoint. The base class is `SessionRoutingBase` +(`src/aiperf/workers/session_routing.py`); passthrough is the default (no headers, body unchanged). + +**Built-ins:** + +| Name | Class | Description | +|------|-------|-------------| +| `dynamo_headers` | `DynamoHeadersRouting` | Dynamo session affinity via `X-Dynamo-Session-ID`, plus `X-Dynamo-Parent-Session-ID` on subagent children. No options. | +| `dynamo_nvext` | `DynamoNvextRouting` | Dynamo session affinity via `nvext.session_control` request-body metadata (bind on non-final turns, close on the final turn). Only for Dynamo builds that implement `session_control`. Option: `timeout_seconds` (default 300). | +| `smg_routing_key` | `SmgRoutingKeyRouting` | SGLang Model Gateway `manual`-policy stickiness via `X-SMG-Routing-Key`. No options. | +| `session_id_header` | `SessionIdHeaderRouting` | Generic additive session-affinity header carrying the session's correlation ID. Option: `header_name` (default `X-Session-ID`). | + +**Options (`--session-routing-opt key=value`).** Each plugin exposes an `Options` Pydantic model +(`extra="forbid"`, so unknown keys are rejected at startup). Repeated `--session-routing-opt` +pairs populate it; values are coerced to the model's field types and canonicalized at config +resolution, so downstream code always sees typed values. `--session-routing-opt` without +`--session-routing` is an error, and parameterless modes (`dynamo_headers`, `smg_routing_key`) +reject every opt key. + +**`mutates_body`.** A plugin sets the `mutates_body` class var to `True` when `transform_body` +changes the payload (only `dynamo_nvext` does); header-only modes leave the body untouched. It is +protocol metadata that marks a mode as incompatible with any verbatim-bytes request path. +`transform_body` must never mutate its input — the request path shares cached `Turn.raw_payload` +dicts with the dataset — so it returns a new dict. + +**`on_session_end` contract.** Fires strictly after the session's last worker-side activity, on +every terminal path (final turn, cancellation, cancel-before-start). It is post-session cleanup +only and must be idempotent (default no-op). + +**Stateful-plugin rule.** A stateful plugin must key its instance state on `ctx.x_correlation_id` +only. A session tree deliberately spans workers, so tree-keyed worker state fragments across +processes. For tree-scoped behavior, use the stateless per-request context facts +`root_correlation_id` and `is_tree_final` instead of accumulating state. + ### Processing Categories | Category | Enum | Description | diff --git a/src/aiperf/common/models/model_endpoint_info.py b/src/aiperf/common/models/model_endpoint_info.py index e5e98bb62b..50ed2e87ff 100644 --- a/src/aiperf/common/models/model_endpoint_info.py +++ b/src/aiperf/common/models/model_endpoint_info.py @@ -122,6 +122,14 @@ def _redact_headers(self, value: list[tuple[str, str]]) -> list[tuple[str, str]] default=EndpointDefaults.USE_SERVER_TOKEN_COUNT, description="Use server-reported token counts from API usage fields instead of client-side tokenization.", ) + session_routing: str | None = Field( + default=None, + description="Selected session-routing plugin name (None = off).", + ) + session_routing_opts: dict[str, Any] = Field( + default_factory=dict, + description="Raw opts for the routing plugin's Options model.", + ) connection_reuse_strategy: ConnectionReuseStrategy = Field( default=EndpointDefaults.CONNECTION_REUSE_STRATEGY, description="Transport connection reuse strategy.", @@ -207,6 +215,10 @@ def from_run(cls, run: BenchmarkRun) -> ModelEndpointInfo: api_key=ep.api_key, use_legacy_max_tokens=ep.use_legacy_max_tokens, use_server_token_count=ep.use_server_token_count, + session_routing=( + str(ep.session_routing) if ep.session_routing is not None else None + ), + session_routing_opts=dict(ep.session_routing_opts), connection_reuse_strategy=ep.connection_reuse, download_video_content=ep.download_video_content, request_content_type=ep.request_content_type, diff --git a/src/aiperf/common/models/record_models.py b/src/aiperf/common/models/record_models.py index 00e62196cd..815898c18a 100644 --- a/src/aiperf/common/models/record_models.py +++ b/src/aiperf/common/models/record_models.py @@ -653,6 +653,25 @@ class RequestInfo(RecordContext): description="Whether this is the final turn in the conversation. " "Used by per-conversation connection strategy to release the connection lease.", ) + root_correlation_id: str | None = Field( + default=None, + description="The x_correlation_id of the depth-0 root of this request's session TREE. " + "Stable across the whole tree (root + every descendant subagent at any depth); equals " + "x_correlation_id for a root session. Sourced from the originating Credit. Per-request " + "transport context read by the session-routing transform to stamp tree-scoped identity " + "on the outbound request; it stays worker-side and is not carried onto the exported record.", + ) + is_parent_final: bool | None = Field( + default=None, + description="Parent conversation had already returned its final turn at " + "credit-issue time; None for roots or when not determinable. Sourced from " + "the originating Credit.", + ) + is_tree_final: bool = Field( + default=False, + description="Provably the last request this session tree will send " + "(conservative False when indeterminate). Sourced from the originating Credit.", + ) url_index: int | None = Field( default=None, ge=0, diff --git a/src/aiperf/config/endpoint.py b/src/aiperf/config/endpoint.py index a5a8fee1e0..0b2bac96d5 100644 --- a/src/aiperf/config/endpoint.py +++ b/src/aiperf/config/endpoint.py @@ -30,6 +30,7 @@ from aiperf.config.loader.parsing import normalize_http_urls from aiperf.plugin.enums import ( EndpointType, + SessionRoutingType, TransportType, URLSelectionStrategy, ) @@ -316,6 +317,34 @@ def _redact_headers(self, value: dict[str, str]) -> dict[str, str]: ), ] + session_routing: Annotated[ + SessionRoutingType | None, + Field( + default=None, + description=( + "Session-routing transform stamping per-session identity on " + "every request so an external router (SGLang Model Gateway, Dynamo, " + "generic session-affinity LBs) can pin a session's turns to " + "one replica. One mode per run; parameterize with " + "--session-routing-opt key=value." + ), + ), + ] + + session_routing_opts: Annotated[ + dict[str, Any], + Field( + default_factory=dict, + description=( + "Mode-specific options for --session-routing, validated " + "against the selected plugin's Options model (unknown keys and " + "invalid values are rejected at startup). Values are " + "canonicalized/coerced to the plugin's Options model types at " + "validation, so downstream consumers see typed values." + ), + ), + ] + wait_for_model_timeout: Annotated[ float, Field( @@ -525,3 +554,29 @@ def _validate_request_content_type(self) -> Self: f"video_generation); endpoint type {self.type} does not." ) return self + + @model_validator(mode="after") + def validate_session_routing(self) -> Self: + """Fail fast: opts require a mode; opts must satisfy the plugin's Options. + + Canonicalizes the opts to the plugin's Options model types so downstream + consumers (including the pickled BenchmarkRun that reaches workers) carry + coerced values (e.g. ``{"timeout_seconds": 600}``, not ``"600"``). + """ + if self.session_routing is None: + if self.session_routing_opts: + raise ValueError( + "--session-routing-opt requires --session-routing to select a mode." + ) + return self + from aiperf.plugin import plugins + from aiperf.plugin.enums import PluginType + + routing_cls = plugins.get_class( + PluginType.SESSION_ROUTING, str(self.session_routing) + ) + options = routing_cls.Options(**self.session_routing_opts) + canonical = options.model_dump(mode="json", exclude_unset=True) + if canonical != self.session_routing_opts: + self.session_routing_opts = canonical + return self diff --git a/src/aiperf/config/flags/_converter_endpoint.py b/src/aiperf/config/flags/_converter_endpoint.py index afe4773c61..10c1f69599 100644 --- a/src/aiperf/config/flags/_converter_endpoint.py +++ b/src/aiperf/config/flags/_converter_endpoint.py @@ -38,6 +38,21 @@ def _endpoint_template_from_extra( } +def _parse_routing_opts(values: list[str]) -> dict[str, str]: + opts: dict[str, str] = {} + for item in values: + key, separator, value = item.partition("=") + key, value = key.strip(), value.strip() + if not separator or not key or not value: + raise ValueError( + f"Invalid --session-routing-opt {item!r}; expected non-empty key=value" + ) + if key in opts: + raise ValueError(f"Duplicate --session-routing-opt key {key!r}") + opts[key] = value + return opts + + def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None: from aiperf.plugin.enums import EndpointType @@ -72,6 +87,7 @@ def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None: "download_video_content": "download_video_content", "request_content_type": "request_content_type", "session_header": "session_header", + "session_routing": "session_routing", } @@ -97,6 +113,8 @@ def build_endpoint(cli: CLIConfig) -> dict[str, Any]: extra = dict(cli.extra_inputs) _endpoint_template_from_extra(endpoint, extra) endpoint["extra"] = extra + if "session_routing_opt" in cli_set and cli.session_routing_opt: + endpoint["session_routing_opts"] = _parse_routing_opts(cli.session_routing_opt) _endpoint_template_fallback(endpoint) return endpoint diff --git a/src/aiperf/config/flags/_section_fields.py b/src/aiperf/config/flags/_section_fields.py index 15acafabd1..404bea9576 100644 --- a/src/aiperf/config/flags/_section_fields.py +++ b/src/aiperf/config/flags/_section_fields.py @@ -24,6 +24,8 @@ "model_selection_strategy", "request_content_type", "session_header", + "session_routing", + "session_routing_opt", "streaming", "timeout_seconds", "transport", diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index 45afb8b7c8..6f73ff169f 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -81,6 +81,7 @@ GPUTelemetryCollectorType, PublicDatasetType, SearchPlannerType, + SessionRoutingType, TransportType, UIType, URLSelectionStrategy, @@ -389,6 +390,42 @@ class CLIConfig(BaseConfig): ), ] = None + session_routing: Annotated[ + SessionRoutingType | None, + Field( + description=( + "Session-aware routing mode: stamps per-session identity on " + "every request for router affinity. Built-ins: dynamo_headers " + "(X-Dynamo-Session-ID + parent header), dynamo_nvext " + "(nvext.session_control bind/close request-body metadata), " + "smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway " + "manual policy), session_id_header (custom additive header). " + "Parameterize with --session-routing-opt." + ), + ), + CLIParameter( + name=("--session-routing",), + group=Groups.ENDPOINT, + ), + ] = None + + session_routing_opt: Annotated[ + list[str], + Field( + default_factory=list, + description=( + "Repeatable key=value option for the selected --session-routing " + "mode (e.g. --session-routing-opt timeout_seconds=600), " + "validated against the plugin's Options model." + ), + ), + CLIParameter( + name=("--session-routing-opt",), + consume_multiple=True, + group=Groups.ENDPOINT, + ), + ] + @property def url(self) -> str: """Return the first URL for backward compatibility.""" diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index cbf1475601..cc96641370 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -7214,6 +7214,24 @@ "description": "HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header. Useful when the inference server expects a custom session-affinity header (e.g. `--session-header X-Session-ID`).", "title": "Sessionheader" }, + "sessionRouting": { + "anyOf": [ + { + "$ref": "#/$defs/SessionRoutingType" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Session-routing transform stamping per-session identity on every request so an external router (SGLang Model Gateway, Dynamo, generic session-affinity LBs) can pin a session's turns to one replica. One mode per run; parameterize with --session-routing-opt key=value." + }, + "sessionRoutingOpts": { + "additionalProperties": true, + "description": "Mode-specific options for --session-routing, validated against the selected plugin's Options model (unknown keys and invalid values are rejected at startup). Values are canonicalized/coerced to the plugin's Options model types at validation, so downstream consumers see typed values.", + "title": "Sessionroutingopts", + "type": "object" + }, "waitForModelTimeout": { "default": 0.0, "description": "Enable a pre-flight readiness probe by setting this to a positive value (seconds). aiperf applies this timeout to each URL/model probe before starting the benchmark, aborting with a non-zero exit if any probe times out. For multiple URLs or models, worst-case wall-clock time can be roughly this timeout multiplied by the number of URL/model probes. The probe strategy is controlled by `--wait-for-model-mode`, which defaults to sending a 1-token inference request. 0 (default) disables the probe. Eliminates the need for external shell-based readiness loops in containers and Kubernetes recipes.", @@ -11880,6 +11898,16 @@ "title": "ServiceRunType", "type": "string" }, + "SessionRoutingType": { + "enum": [ + "dynamo_headers", + "dynamo_nvext", + "smg_routing_key", + "session_id_header" + ], + "title": "SessionRoutingType", + "type": "string" + }, "SobolSweep": { "additionalProperties": false, "description": "Sobol quasi-Monte-Carlo sweep — low-discrepancy joint coverage.\n\nNote on ``scramble=False``: scipy's raw Sobol sequence starts at the\norigin (the first row is exactly all-zeros), so the first variant\nwill land on ``(lo, lo, ...)`` for every dimension. With small\n``samples`` this produces degenerate corner-only coverage. Prefer\nthe default ``scramble=True`` (Owen-scrambled) unless you have a\nspecific reason to want raw Sobol points.", diff --git a/src/aiperf/credit/issuer.py b/src/aiperf/credit/issuer.py index 41a14da326..d09f344f4b 100644 --- a/src/aiperf/credit/issuer.py +++ b/src/aiperf/credit/issuer.py @@ -30,6 +30,7 @@ from aiperf.timing.phase.progress_tracker import PhaseProgressTracker from aiperf.timing.phase.stop_conditions import StopConditionChecker from aiperf.timing.request_cancellation import RequestCancellationSimulator + from aiperf.timing.session_tree import SessionTreeRegistry class CreditIssuer: @@ -62,6 +63,8 @@ def __init__( cancellation_policy: RequestCancellationSimulator, lifecycle: PhaseLifecycle, url_selection_strategy: URLSelectionStrategyProtocol | None = None, + session_tree_registry: SessionTreeRegistry | None = None, + session_tree_registry_enabled: bool | None = None, ) -> None: """Initialize credit issuer. @@ -75,6 +78,14 @@ def __init__( lifecycle: Phase lifecycle for timestamp data. url_selection_strategy: Optional URL selection strategy for multi-URL load balancing. If None, url_index will be None in credits. + session_tree_registry: Optional per-tree finality ledger (DAG + datasets only). When engaged, a root session start opens a tree + and the issuer stamps ``is_parent_final`` / ``is_tree_final`` on + every emitted credit from live tree state. None on non-DAG paths + (finality stays conservative ``(None, False)``). + session_tree_registry_enabled: Explicit engage override. When None, + the registry is engaged iff one was supplied (DAG runs engage it + in every phase). Set True/False to force it on/off (tests). """ self._phase = phase self._stop_checker = stop_checker @@ -84,6 +95,54 @@ def __init__( self._cancellation_policy = cancellation_policy self._lifecycle = lifecycle self._url_selection_strategy = url_selection_strategy + 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 + ) + + def _open_session_tree(self, turn: TurnToSend) -> None: + """Open a session tree for a root session-start credit just admitted. + + No-op when tree accounting is not engaged (non-DAG). A root session + start is always depth 0, so the tree root id is the root's own + ``x_correlation_id``. + """ + if self._session_tree_registry is not None: + self._session_tree_registry.open_tree( + turn.effective_root_correlation_id, self._phase, root_pending=True + ) + + def _finality_for_issue(self, turn: TurnToSend) -> tuple[bool | None, bool]: + """Issue-time lineage finality from ``SessionTreeRegistry`` state. + + Conservative by spec: returns ``None``/``False`` whenever indeterminate + (including the non-DAG path where no registry is engaged). + """ + registry = self._session_tree_registry + if registry is None: + return None, False + root_id = turn.effective_root_correlation_id + is_root = turn.parent_correlation_id is None + is_parent_final: bool | None = None + if not is_root and turn.parent_correlation_id == root_id: + # v1: parent finality is determinable only when the parent IS the + # root (the registry tracks per-tree, not per-intermediate-node). + is_parent_final = registry.root_terminal(root_id) + is_tree_final = registry.is_last_tree_request( + root_id, + is_final_turn=turn.is_final_turn, + is_root_credit=is_root, + # Any-mode branch flag, NOT the FORK-only has_forks: a final turn + # declaring SPAWN branches spawns descendants at return-intercept, + # after this stamp, so it must never read as tree-final. + has_branches=turn.has_branches, + ) + return is_parent_final, is_tree_final def can_acquire_and_start_new_session(self) -> bool: """Check if a session slot can be acquired and a new session can be started.""" @@ -134,6 +193,7 @@ async def issue_credit(self, turn: TurnToSend) -> bool: ) if not acquired: return False + self._open_session_tree(turn) # Prefill concurrency: one slot per request, released when TTFT arrives. # Limits concurrent prompt processing which is the GPU-intensive phase. @@ -182,6 +242,7 @@ async def try_issue_credit(self, turn: TurnToSend) -> bool | None: ) if not acquired: return None # No slot - credit not issued + self._open_session_tree(turn) acquired = self._concurrency_manager.try_acquire_prefill_slot( self._phase, can_proceed_fn @@ -219,6 +280,8 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: else None ) + is_parent_final, is_tree_final = self._finality_for_issue(turn) + credit = Credit( id=credit_index, phase=self._phase, @@ -231,7 +294,10 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: url_index=url_index, agent_depth=turn.agent_depth, parent_correlation_id=turn.parent_correlation_id, + root_correlation_id=turn.root_correlation_id, has_forks=turn.has_forks, + is_parent_final=is_parent_final, + is_tree_final=is_tree_final, branch_mode=turn.branch_mode, ) @@ -304,7 +370,9 @@ async def dispatch_join_turn(self, pending: PendingBranchJoin) -> bool: num_turns=pending.parent_num_turns, agent_depth=pending.parent_agent_depth, parent_correlation_id=pending.parent_parent_correlation_id, + root_correlation_id=pending.parent_root_correlation_id, has_forks=pending.parent_has_forks_on_gated_turn, + has_branches=pending.parent_has_branches_on_gated_turn, branch_mode=pending.parent_branch_mode, ) return await self._dispatch_dag_turn(turn) diff --git a/src/aiperf/credit/structs.py b/src/aiperf/credit/structs.py index 2ce51b958d..fbc031bd60 100644 --- a/src/aiperf/credit/structs.py +++ b/src/aiperf/credit/structs.py @@ -63,7 +63,23 @@ class Credit( url_index: int | None = None agent_depth: int = 0 parent_correlation_id: str | None = None + root_correlation_id: str | None = None + """x_correlation_id of the depth-0 root of this credit's session TREE. + + Stable across the whole tree: the root carries its own x_correlation_id + (left None on the wire when it equals x_correlation_id to keep the struct + small), and every descendant (child, subchild) inherits the root's id. + This is the key used for per-tree finality accounting + (``SessionTreeRegistry``) and is persisted in the export so analysis groups + a tree under one lane. Effective value is + ``root_correlation_id or x_correlation_id``.""" has_forks: bool = False + is_parent_final: bool | None = None + """Parent conversation had already returned its final turn at issue time. + None for roots / when not determinable. Issue-time stamp, never copied.""" + is_tree_final: bool = False + """Provably the last request the whole session tree will send (conservative + False when indeterminate). Issue-time stamp from SessionTreeRegistry.""" branch_mode: ConversationBranchMode = ConversationBranchMode.FORK """DAG branch mode for this credit. Ignored when parent_correlation_id is None (i.e. for root sessions). FORK = inherit parent turn_list; SPAWN = @@ -73,6 +89,11 @@ class Credit( def is_final_turn(self) -> bool: return self.turn_index == self.num_turns - 1 + @property + def effective_root_correlation_id(self) -> str: + """Tree root id, defaulting to this credit's own ``x_correlation_id``.""" + return self.root_correlation_id or self.x_correlation_id + class CreditContext( Struct, omit_defaults=True, kw_only=True, tag_field="t", tag="cctx" @@ -126,13 +147,33 @@ class TurnToSend(Struct, frozen=True): num_turns: int agent_depth: int = 0 parent_correlation_id: str | None = None + root_correlation_id: str | None = None + """x_correlation_id of the depth-0 root of this turn's session TREE. + + None for a root turn (the root IS its own tree root); set on every + descendant to the root's id. Propagated onto the issued ``Credit`` and used + for per-tree finality accounting. Effective value is + ``root_correlation_id or x_correlation_id``.""" has_forks: bool = False + has_branches: bool = False + """True iff the originating turn declares ANY branch (FORK or SPAWN) in its + metadata ``branch_ids``. Superset of ``has_forks``, which is FORK-only and + owned by the sticky router's deferred-eviction logic — do not conflate the + two. Consumed by finality stamping: a turn that will spawn descendants on + its return can never be the tree's provably-last request, even when the + registry shows nothing outstanding yet (SPAWN children register only at + return-intercept, AFTER issue-time stamping).""" branch_mode: ConversationBranchMode = ConversationBranchMode.FORK @property def is_final_turn(self) -> bool: return self.turn_index == self.num_turns - 1 + @property + def effective_root_correlation_id(self) -> str: + """Tree root id, defaulting to this turn's own ``x_correlation_id``.""" + return self.root_correlation_id or self.x_correlation_id + @classmethod def from_previous_credit( cls, credit: Credit, next_meta: "TurnMetadata | None" = None @@ -143,7 +184,10 @@ def from_previous_credit( credit: The previous turn's credit. next_meta: Metadata for the NEW turn being built. When provided, the ``has_forks`` flag is derived from it so the sticky - router can defer parent-entry eviction until DAG children drain. + router can defer parent-entry eviction until DAG children drain, + and ``has_branches`` (any-mode) is derived from its + ``branch_ids`` so finality stamping stays conservative on + spawning turns. """ return cls( conversation_id=credit.conversation_id, @@ -152,6 +196,8 @@ def from_previous_credit( num_turns=credit.num_turns, agent_depth=credit.agent_depth, parent_correlation_id=credit.parent_correlation_id, + root_correlation_id=credit.root_correlation_id, has_forks=next_meta.has_forks if next_meta is not None else False, + has_branches=bool(next_meta.branch_ids) if next_meta is not None else False, branch_mode=credit.branch_mode, ) diff --git a/src/aiperf/dataset/loader/exgentic.py b/src/aiperf/dataset/loader/exgentic.py index 4316f7abdd..761defd7b1 100644 --- a/src/aiperf/dataset/loader/exgentic.py +++ b/src/aiperf/dataset/loader/exgentic.py @@ -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")), extra_body=_request_extra_body(attributes), - extra_headers={"x-dynamo-session-id": session_id}, ) except (KeyError, TypeError, ValueError) as error: raise DatasetLoaderError( diff --git a/src/aiperf/plugin/categories.yaml b/src/aiperf/plugin/categories.yaml index 22e9dbfd82..3662a46c66 100644 --- a/src/aiperf/plugin/categories.yaml +++ b/src/aiperf/plugin/categories.yaml @@ -140,6 +140,23 @@ transport: Manages connection pooling, streaming, error handling, and TCP configuration. One-to-one mapping based on transport_type configuration. +# ============================================================================= +# Session Routing Category +# ============================================================================= + +session_routing: + protocol: aiperf.workers.session_routing:SessionRoutingBase + enum: SessionRoutingType + description: | + Session-routing transforms stamp per-session identity onto outbound + requests (headers or body metadata) so an external router in front of + multiple replicas can pin every turn of a session to one worker. + Selected via --session-routing; parameterized via repeatable + --session-routing-opt key=value validated against the plugin's + Options model. mutates_body is protocol metadata marking a mode as + body-mutating / incompatible with any verbatim-bytes request path; it is + not currently enforced by any gate on this codebase. + # ============================================================================= # Processing Categories # ============================================================================= diff --git a/src/aiperf/plugin/enums.py b/src/aiperf/plugin/enums.py index 506df476e0..1699969bdb 100644 --- a/src/aiperf/plugin/enums.py +++ b/src/aiperf/plugin/enums.py @@ -12,7 +12,7 @@ from aiperf.plugin import plugins from aiperf.plugin.extensible_enums import create_enum -__all__ = ["APIRouterType", "APIRouterTypeStr", "AccuracyBenchmarkType", "AccuracyBenchmarkTypeStr", "AccuracyGraderType", "AccuracyGraderTypeStr", "ArrivalPattern", "ArrivalPatternStr", "CommClientType", "CommClientTypeStr", "CommunicationBackend", "CommunicationBackendStr", "ComposerType", "ComposerTypeStr", "ConsoleExporterType", "ConsoleExporterTypeStr", "ConvergenceCriterionType", "ConvergenceCriterionTypeStr", "CustomDatasetType", "CustomDatasetTypeStr", "DataExporterType", "DataExporterTypeStr", "DatasetBackingStoreType", "DatasetBackingStoreTypeStr", "DatasetClientStoreType", "DatasetClientStoreTypeStr", "DatasetFormat", "DatasetFormatStr", "DatasetSamplingStrategy", "DatasetSamplingStrategyStr", "EndpointType", "EndpointTypeStr", "GPUTelemetryCollectorType", "GPUTelemetryCollectorTypeStr", "GPUTelemetryProcessorType", "GPUTelemetryProcessorTypeStr", "PhaseType", "PhaseTypeStr", "PlotType", "PlotTypeStr", "PluginType", "PluginTypeStr", "PublicDatasetType", "PublicDatasetTypeStr", "RampType", "RampTypeStr", "RecordProcessorType", "RecordProcessorTypeStr", "ResultsProcessorType", "ResultsProcessorTypeStr", "SearchPlannerType", "SearchPlannerTypeStr", "SearchRecipePostProcessType", "SearchRecipePostProcessTypeStr", "SearchRecipeType", "SearchRecipeTypeStr", "ServerMetricsProcessorType", "ServerMetricsProcessorTypeStr", "ServiceRunType", "ServiceRunTypeStr", "ServiceType", "ServiceTypeStr", "TimingMode", "TimingModeStr", "TransportType", "TransportTypeStr", "UIType", "UITypeStr", "URLSelectionStrategy", "URLSelectionStrategyStr", "ZMQProxyType", "ZMQProxyTypeStr"] +__all__ = ["APIRouterType", "APIRouterTypeStr", "AccuracyBenchmarkType", "AccuracyBenchmarkTypeStr", "AccuracyGraderType", "AccuracyGraderTypeStr", "ArrivalPattern", "ArrivalPatternStr", "CommClientType", "CommClientTypeStr", "CommunicationBackend", "CommunicationBackendStr", "ComposerType", "ComposerTypeStr", "ConsoleExporterType", "ConsoleExporterTypeStr", "ConvergenceCriterionType", "ConvergenceCriterionTypeStr", "CustomDatasetType", "CustomDatasetTypeStr", "DataExporterType", "DataExporterTypeStr", "DatasetBackingStoreType", "DatasetBackingStoreTypeStr", "DatasetClientStoreType", "DatasetClientStoreTypeStr", "DatasetFormat", "DatasetFormatStr", "DatasetSamplingStrategy", "DatasetSamplingStrategyStr", "EndpointType", "EndpointTypeStr", "GPUTelemetryCollectorType", "GPUTelemetryCollectorTypeStr", "GPUTelemetryProcessorType", "GPUTelemetryProcessorTypeStr", "PhaseType", "PhaseTypeStr", "PlotType", "PlotTypeStr", "PluginType", "PluginTypeStr", "PublicDatasetType", "PublicDatasetTypeStr", "RampType", "RampTypeStr", "RecordProcessorType", "RecordProcessorTypeStr", "ResultsProcessorType", "ResultsProcessorTypeStr", "SearchPlannerType", "SearchPlannerTypeStr", "SearchRecipePostProcessType", "SearchRecipePostProcessTypeStr", "SearchRecipeType", "SearchRecipeTypeStr", "ServerMetricsProcessorType", "ServerMetricsProcessorTypeStr", "ServiceRunType", "ServiceRunTypeStr", "ServiceType", "ServiceTypeStr", "SessionRoutingType", "SessionRoutingTypeStr", "TimingMode", "TimingModeStr", "TransportType", "TransportTypeStr", "UIType", "UITypeStr", "URLSelectionStrategy", "URLSelectionStrategyStr", "ZMQProxyType", "ZMQProxyTypeStr"] # Plugin Protocol Categories if TYPE_CHECKING: @@ -73,6 +73,10 @@ TransportType = plugins.create_enum(PluginType.TRANSPORT, "TransportType", module=__name__) """Dynamic enum for transport. Example: TransportType.HTTP""" +SessionRoutingTypeStr: TypeAlias = str +SessionRoutingType = plugins.create_enum(PluginType.SESSION_ROUTING, "SessionRoutingType", module=__name__) +"""Dynamic enum for session routing. Example: SessionRoutingType.DYNAMO_HEADERS, SessionRoutingType.SESSION_ID_HEADER, SessionRoutingType.SMG_ROUTING_KEY""" + RecordProcessorTypeStr: TypeAlias = str RecordProcessorType = plugins.create_enum(PluginType.RECORD_PROCESSOR, "RecordProcessorType", module=__name__) """Dynamic enum for record processor. Example: RecordProcessorType.ACCURACY_RECORD, RecordProcessorType.OUTPUTS_JSON, RecordProcessorType.RAW_RECORD_WRITER""" diff --git a/src/aiperf/plugin/plugins.py b/src/aiperf/plugin/plugins.py index c16d6abf1b..0ad87f5b8b 100644 --- a/src/aiperf/plugin/plugins.py +++ b/src/aiperf/plugin/plugins.py @@ -938,7 +938,7 @@ def _load_package_metadata( from aiperf.orchestrator.convergence.base import ConvergenceCriterion from aiperf.orchestrator.search_planner.base import SearchPlanner from aiperf.plot.core.plot_type_handlers import PlotTypeHandlerProtocol - from aiperf.plugin.enums import APIRouterType, AccuracyBenchmarkType, AccuracyGraderType, ArrivalPattern, CommClientType, CommunicationBackend, ComposerType, ConsoleExporterType, ConvergenceCriterionType, CustomDatasetType, DataExporterType, DatasetBackingStoreType, DatasetClientStoreType, DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, GPUTelemetryProcessorType, PlotType, PluginType, PluginTypeStr, PublicDatasetType, RampType, RecordProcessorType, ResultsProcessorType, SearchPlannerType, SearchRecipePostProcessType, SearchRecipeType, ServerMetricsProcessorType, ServiceRunType, ServiceType, TimingMode, TransportType, UIType, URLSelectionStrategy, ZMQProxyType + from aiperf.plugin.enums import APIRouterType, AccuracyBenchmarkType, AccuracyGraderType, ArrivalPattern, CommClientType, CommunicationBackend, ComposerType, ConsoleExporterType, ConvergenceCriterionType, CustomDatasetType, DataExporterType, DatasetBackingStoreType, DatasetClientStoreType, DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, GPUTelemetryProcessorType, PlotType, PluginType, PluginTypeStr, PublicDatasetType, RampType, RecordProcessorType, ResultsProcessorType, SearchPlannerType, SearchRecipePostProcessType, SearchRecipeType, ServerMetricsProcessorType, ServiceRunType, ServiceType, SessionRoutingType, TimingMode, TransportType, UIType, URLSelectionStrategy, ZMQProxyType from aiperf.post_processors.base_metrics_processor import BaseMetricsProcessor from aiperf.post_processors.protocols import RecordProcessorProtocol from aiperf.search_recipes._base import SearchRecipe @@ -949,6 +949,7 @@ def _load_package_metadata( from aiperf.timing.url_samplers import URLSelectionStrategyProtocol from aiperf.transports.base_transports import TransportProtocol from aiperf.ui.protocols import AIPerfUIProtocol + from aiperf.workers.session_routing import SessionRoutingBase from aiperf.zmq.zmq_proxy_base import BaseZMQProxy from typing import Literal, overload # @@ -1002,6 +1003,10 @@ def get_class(category: Literal[PluginType.TRANSPORT, "transport"], name_or_clas @overload def iter_all(category: Literal[PluginType.TRANSPORT, "transport"]) -> Iterator[tuple[PluginEntry, type[TransportProtocol]]]: ... @overload + def get_class(category: Literal[PluginType.SESSION_ROUTING, "session_routing"], name_or_class_path: SessionRoutingType | str) -> type[SessionRoutingBase]: ... + @overload + def iter_all(category: Literal[PluginType.SESSION_ROUTING, "session_routing"]) -> Iterator[tuple[PluginEntry, type[SessionRoutingBase]]]: ... + @overload def get_class(category: Literal[PluginType.RECORD_PROCESSOR, "record_processor"], name_or_class_path: RecordProcessorType | str) -> type[RecordProcessorProtocol]: ... @overload def iter_all(category: Literal[PluginType.RECORD_PROCESSOR, "record_processor"]) -> Iterator[tuple[PluginEntry, type[RecordProcessorProtocol]]]: ... diff --git a/src/aiperf/plugin/plugins.yaml b/src/aiperf/plugin/plugins.yaml index d72121a9d0..e3d99d2b66 100644 --- a/src/aiperf/plugin/plugins.yaml +++ b/src/aiperf/plugin/plugins.yaml @@ -437,6 +437,45 @@ transport: transport_type: http url_schemes: [http, https] +# ============================================================================= +# Session Routing +# ============================================================================= +# Session-routing transforms stamp per-session identity onto outbound +# requests so an external router can pin every turn to one worker. +# Selected via --session-routing; one instance per worker. +# ============================================================================= +session_routing: + dynamo_headers: + class: aiperf.workers.session_routing:DynamoHeadersRouting + description: | + Dynamo session affinity via X-Dynamo-Session-ID plus + X-Dynamo-Parent-Session-ID on subagent children. Pair with a Dynamo + frontend running --router-session-affinity-ttl-secs. No options. + + dynamo_nvext: + class: aiperf.workers.session_routing:DynamoNvextRouting + description: | + Dynamo session affinity via nvext.session_control request-body + metadata: bind on every non-final turn, close on the final turn. + Only for Dynamo builds that implement session_control (current + upstream main does not; use dynamo_headers there). Sets mutates_body + (protocol metadata: body-mutating / incompatible with any verbatim-bytes + request path; not enforced by any gate on this codebase today). + Option: timeout_seconds (default 300). + + smg_routing_key: + class: aiperf.workers.session_routing:SmgRoutingKeyRouting + description: | + SGLang Model Gateway stickiness: emits X-SMG-Routing-Key with the + session's correlation ID for the gateway's manual routing policy + (--policy manual). No options. + + session_id_header: + class: aiperf.workers.session_routing:SessionIdHeaderRouting + description: | + Generic additive session-affinity header carrying the session's + correlation ID. Option: header_name (default X-Session-ID). + # ============================================================================= # Dataset Composers # ============================================================================= @@ -1711,7 +1750,9 @@ public_dataset_loader: Exgentic agent LLM traces from HuggingFace. Replays complete recorded message and tool snapshots with recorded output limits. Fixed-schedule mode preserves call timestamps and overlaps; other modes replay sessions - closed-loop with recorded inter-turn delays. + closed-loop with recorded inter-turn delays. Pair with `--session-routing + dynamo_headers` for Dynamo session affinity (the loader no longer stamps a + session header). metadata: hf_dataset_name: Exgentic/agent-llm-traces hf_split: train diff --git a/src/aiperf/plugin/schema/plugins.schema.json b/src/aiperf/plugin/schema/plugins.schema.json index 3e9265d4f0..58a5765a6f 100644 --- a/src/aiperf/plugin/schema/plugins.schema.json +++ b/src/aiperf/plugin/schema/plugins.schema.json @@ -106,6 +106,14 @@ "$ref": "#/$defs/TransportPlugin" } }, + "session_routing": { + "title": "Session Routing Plugins", + "type": "object", + "description": "Session-routing transforms stamp per-session identity onto outbound\nrequests (headers or body metadata) so an external router in front of\nmultiple replicas can pin every turn of a session to one worker.\nSelected via --session-routing; parameterized via repeatable\n--session-routing-opt key=value validated against the plugin's\nOptions model. mutates_body is protocol metadata marking a mode as\nbody-mutating / incompatible with any verbatim-bytes request path; it is\nnot currently enforced by any gate on this codebase.", + "additionalProperties": { + "$ref": "#/$defs/SessionRoutingPlugin" + } + }, "record_processor": { "title": "Record Processor Plugins", "type": "object", @@ -1050,6 +1058,47 @@ "title": "Transport Plugin", "description": "Transports handle the network layer for sending requests to inference servers.\nManages connection pooling, streaming, error handling, and TCP configuration.\nOne-to-one mapping based on transport_type configuration." }, + "SessionRoutingPlugin": { + "type": "object", + "properties": { + "class": { + "description": "Python class that implements this plugin entry. Use 'module.path:ClassName' format, e.g., 'aiperf.endpoints.openai_chat:ChatEndpoint'.", + "title": "Class", + "type": "string" + }, + "description": { + "default": "", + "description": "Brief explanation of what this plugin type does and when to use it.", + "title": "Description", + "type": "string" + }, + "priority": { + "default": 0, + "description": "Conflict resolution priority. When multiple packages register the same type name, the one with higher priority wins. Use 0 for normal plugins, higher values to override built-in implementations.", + "title": "Priority", + "type": "integer" + }, + "metadata": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Category-specific configuration for this plugin type. The allowed fields depend on the category's metadata_class in categories.yaml.", + "title": "Metadata" + } + }, + "required": [ + "class" + ], + "title": "Session Routing Plugin", + "description": "Session-routing transforms stamp per-session identity onto outbound\nrequests (headers or body metadata) so an external router in front of\nmultiple replicas can pin every turn of a session to one worker.\nSelected via --session-routing; parameterized via repeatable\n--session-routing-opt key=value validated against the plugin's\nOptions model. mutates_body is protocol metadata marking a mode as\nbody-mutating / incompatible with any verbatim-bytes request path; it is\nnot currently enforced by any gate on this codebase." + }, "RecordProcessorPlugin": { "type": "object", "properties": { diff --git a/src/aiperf/timing/_branch_orchestrator_spawn.py b/src/aiperf/timing/_branch_orchestrator_spawn.py index 1c79b2a51a..8612b66f3b 100644 --- a/src/aiperf/timing/_branch_orchestrator_spawn.py +++ b/src/aiperf/timing/_branch_orchestrator_spawn.py @@ -42,8 +42,20 @@ class BranchOrchestratorSpawnMixin: which spawns worker subprocesses. """ - async def _fire_pre_session_children(self, branch: ConversationBranchInfo) -> None: - """Issue turn-0 for every child of a pre-session SPAWN branch.""" + async def _fire_pre_session_children( + self, branch: ConversationBranchInfo, parent_conversation_id: str + ) -> None: + """Issue turn-0 for every child of a pre-session SPAWN branch. + + Each issued child is tracked as a LIVE pre-session descendant of + ``parent_conversation_id``. No root session exists yet (pre-dispatch + runs before sampling mints any root correlation id), so the children + cannot be registered with the session-tree registry here; the + orchestrator folds the still-live set into the first root instance's + tree at that root's turn-0 return + (``_fold_pre_session_descendants``), keeping the root's final-turn + ``is_tree_final`` stamp conservative while these children run. + """ for child_cid in branch.child_conversation_ids: try: child_session = self._cs.start_pre_session_child(child_cid) @@ -54,6 +66,11 @@ async def _fire_pre_session_children(self, branch: ConversationBranchInfo) -> No issued = await self._issuer.dispatch_first_turn(child_session) if issued: self.stats.children_spawned += 1 + child_corr = child_session.x_correlation_id + self._pre_child_conv[child_corr] = parent_conversation_id + self._pre_live_by_conv.setdefault(parent_conversation_id, set()).add( + child_corr + ) else: # ``dispatch_first_turn`` only returns False under # stop-condition refusal (cap reached). Tally as truncated. @@ -84,6 +101,14 @@ async def _spawn_children_and_register_gates( if all_children: self._descendant_counts.setdefault(parent_corr, 0) self._descendant_counts[parent_corr] += len(all_children) + # Per-tree finality ledger keys on the ROOT (not the parent), so a + # depth>1 spawn still folds into the depth-0 root's tree. Registered + # BEFORE the dispatch gather; a child whose dispatch does not land is + # decremented again in _rollback_failed_child. + if self._session_tree_registry is not None: + self._session_tree_registry.register_descendants( + credit.effective_root_correlation_id, len(all_children) + ) # If any expected gate had zero children actually register, still # create a future-join entry so the drain logic sees it and fires. @@ -180,6 +205,7 @@ def _start_one_child( parent_correlation_id=parent_corr, child_conversation_id=child_conv_id, agent_depth=parent_depth + 1, + root_correlation_id=credit.effective_root_correlation_id, branch_mode=branch.mode, ) except Exception: @@ -187,6 +213,9 @@ def _start_one_child( self.stats.children_errored += 1 return None child_corr = child.x_correlation_id + # Capture the child's tree root so descendant-done can key the registry + # by root even for depth>1 descendants. + self._child_root[child_corr] = credit.effective_root_correlation_id self._child_modes[child_corr] = branch.mode # FORK-mode children sticky-route to the parent's worker; SPAWN-mode # children do not register a refcount. @@ -237,6 +266,10 @@ def _rollback_failed_child( child_corr = child.x_correlation_id child_mode = per_child_branch_mode.get(child_corr) self._child_modes.pop(child_corr, None) + # Undo the register_descendants that ran for this child at spawn time. + rolled_root = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and rolled_root is not None: + self._session_tree_registry.on_descendant_done(rolled_root) entries = self._child_to_join.pop(child_corr, []) for entry in entries: if entry.prereq_key is None: @@ -257,10 +290,17 @@ def _rollback_failed_child( self._sticky_router.release_child_routing(parent_corr) if parent_corr in self._descendant_counts: self._descendant_counts[parent_corr] -= 1 - # Three-way classification of non-True gather results: - # * BaseException -> genuine error. - # * False -> stop-condition refusal; tally as truncated. - # * None -> issuer suppressed silently; observable no-op. + self._tally_failed_child_result(result, child_corr) + self.stats.children_spawned -= 1 + + def _tally_failed_child_result(self, result: object, child_corr: str) -> None: + """Fold a non-True ``dispatch_first_turn`` gather result into stats. + + Three-way classification of non-True gather results: + * BaseException -> genuine error. + * False -> stop-condition refusal; tally as truncated. + * None -> issuer suppressed silently; observable no-op. + """ if isinstance(result, BaseException): logger.error( "dispatch_first_turn failed for child %s", @@ -279,7 +319,6 @@ def _rollback_failed_child( child_corr, ) self.stats.children_errored += 1 - self.stats.children_spawned -= 1 async def _drain_vestigial_gates(self, parent_corr: str) -> None: """Drain any zero-outstanding gates created in this spawn cycle. diff --git a/src/aiperf/timing/_branch_orchestrator_state.py b/src/aiperf/timing/_branch_orchestrator_state.py index 941d8a6657..62dee76e4d 100644 --- a/src/aiperf/timing/_branch_orchestrator_state.py +++ b/src/aiperf/timing/_branch_orchestrator_state.py @@ -72,10 +72,20 @@ class PendingBranchJoin: parent_num_turns: int parent_agent_depth: int = 0 parent_parent_correlation_id: str | None = None + parent_root_correlation_id: str | None = None + """Root of the parent's session tree (the parent credit's own + ``root_correlation_id``): None when the parent IS the tree root, set for a + nested parent. Re-stamped onto the gated join ``TurnToSend`` so the resumed + parent keeps its tree lineage for finality accounting.""" gated_turn_index: int | None = None outstanding: dict[str, PrereqState] = field(default_factory=dict) parent_branch_mode: ConversationBranchMode = ConversationBranchMode.FORK parent_has_forks_on_gated_turn: bool = False + parent_has_branches_on_gated_turn: bool = False + """True iff the gated turn itself declares ANY branch (FORK or SPAWN) in its + ``branch_ids``. Superset of the FORK-only flag above; re-stamped onto the + join ``TurnToSend`` so finality stamping stays conservative when the + resumed turn will spawn further descendants.""" is_blocked: bool = False created_at_ns: int = field(default_factory=time.monotonic_ns) diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index 1ec3f64da1..d8320386f2 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -34,6 +34,7 @@ from aiperf.credit.sticky_router import StickyCreditRouter from aiperf.credit.structs import Credit from aiperf.timing.conversation_source import ConversationSource, SampledSession + from aiperf.timing.session_tree import SessionTreeRegistry __all__ = [ "BranchOrchestrator", @@ -84,11 +85,32 @@ def __init__( sticky_router: StickyCreditRouter | None = None, *, benchmark_id: str = "unknown", + session_tree_registry: SessionTreeRegistry | None = None, ) -> None: self._cs = conversation_source self._issuer = credit_issuer self._sticky_router = sticky_router self._benchmark_id = benchmark_id + # Per-tree finality ledger (DAG datasets only). The orchestrator drives + # descendant registration + descendant/root terminal transitions; the + # issuer opens the tree and reads it for finality stamping. None on + # non-DAG paths. Finality bookkeeping only -- never a slot ledger. + self._session_tree_registry = session_tree_registry + # child_x_correlation_id -> its tree's root_correlation_id, captured at + # spawn so descendant-done can key the registry by root (parent != root + # for depth>1). Popped on completion/rollback. + self._child_root: dict[str, str] = {} + # Pre-session SPAWN children fire once per phase BEFORE any root + # session is sampled, so no root correlation id exists yet at dispatch. + # Track them per parent conversation until the first root instance's + # turn-0 return folds them into that root's session tree + # (register_descendants), so the root's final turn cannot stamp + # is_tree_final=True while a pre-session child is still live. + # _pre_child_conv: pre-child corr -> parent conversation_id (popped at + # the child's terminal). _pre_live_by_conv: conversation_id -> live + # pre-child corrs not yet folded into a tree (popped at fold). + self._pre_child_conv: dict[str, str] = {} + self._pre_live_by_conv: dict[str, set[str]] = {} self._child_modes: dict[str, ConversationBranchMode] = {} # Two-level pending-join state: a "future" join is registered at # spawn time and promoted to "active" once the parent reaches the @@ -157,7 +179,7 @@ async def dispatch_pre_session_branches(self) -> None: continue if branch.branch_id not in turn0_branch_ids: continue - await self._fire_pre_session_children(branch) + await self._fire_pre_session_children(branch, conv.conversation_id) self._pre_dispatched_branches.add( (conv.conversation_id, branch.branch_id) ) @@ -186,8 +208,65 @@ async def intercept(self, credit: Credit) -> bool: branch_ids = self.get_branch_ids(credit) if branch_ids: await self._spawn_children_and_register_gates(credit, branch_ids) + # Fold this conversation's live pre-session SPAWN children into the + # root instance's tree at its turn-0 return -- the earliest moment + # the root's correlation id meets the orchestrator (pre-dispatch + # fires before sampling, so no root id was knowable then). Must run + # BEFORE the root-terminal transition below so a single-turn root's + # tree does not retire while its pre-session children are live. + if credit.agent_depth == 0 and credit.turn_index == 0: + self._fold_pre_session_descendants(credit) + # Root's terminal turn: clear tree root-pending AFTER final-turn + # spawns have registered their descendants (register runs inside + # _spawn_children_and_register_gates above), so the tree does not + # transiently read as drained. + if ( + self._session_tree_registry is not None + and credit.agent_depth == 0 + and credit.is_final_turn + ): + self._session_tree_registry.on_root_terminal( + credit.effective_root_correlation_id + ) return self._maybe_suspend_parent(credit) + def _fold_pre_session_descendants(self, credit: Credit) -> None: + """Attach live pre-session SPAWN children to a root instance's tree. + + Runs at the root's turn-0 return. Folds ONCE, into the first root + instance of the conversation (the live set is popped); children that + already terminated were removed from the set on their return and are + not counted. Multi-instance runs of the same template attach the + one-shot pre children to the first instance's tree only. + """ + live = self._pre_live_by_conv.pop(credit.conversation_id, None) + if not live: + return + root_corr = credit.effective_root_correlation_id + for child_corr in live: + self._child_root[child_corr] = root_corr + if self._session_tree_registry is not None: + self._session_tree_registry.register_descendants(root_corr, len(live)) + + def _resolve_pre_session_descendant(self, child_corr: str) -> None: + """Terminal accounting for a pre-session SPAWN child. + + Before its conversation's fold: drop it from the live set so the fold + does not count an already-finished child. After the fold: decrement its + tree's outstanding count. No-op for non-pre-session children. + """ + conv_id = self._pre_child_conv.pop(child_corr, None) + if conv_id is None: + return + live = self._pre_live_by_conv.get(conv_id) + if live is not None: + live.discard(child_corr) + if not live: + self._pre_live_by_conv.pop(conv_id, None) + root_corr = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and root_corr is not None: + self._session_tree_registry.on_descendant_done(root_corr) + def _ensure_future_join( self, credit: Credit, @@ -200,21 +279,25 @@ def _ensure_future_join( pending = gates_for_parent.get(gated_idx) if pending is None: has_forks = False + has_branches = False if 0 <= gated_idx < len(parent_meta.turns): has_forks = bool( getattr(parent_meta.turns[gated_idx], "has_forks", False) ) + has_branches = bool(parent_meta.turns[gated_idx].branch_ids) pending = PendingBranchJoin( parent_x_correlation_id=parent_corr, parent_conversation_id=credit.conversation_id, parent_num_turns=len(parent_meta.turns), parent_agent_depth=credit.agent_depth, parent_parent_correlation_id=credit.parent_correlation_id, + parent_root_correlation_id=credit.root_correlation_id, gated_turn_index=gated_idx, parent_branch_mode=getattr( credit, "branch_mode", ConversationBranchMode.FORK ), parent_has_forks_on_gated_turn=has_forks, + parent_has_branches_on_gated_turn=has_branches, ) # Phase 3 fan-in seed: pre-populate every prereq_key declared # by the gated turn with an unregistered PrereqState so the @@ -355,6 +438,7 @@ async def on_child_leaf_reached(self, child_x_correlation_id: str) -> None: """Called when a child session reaches its final turn.""" if self._cleaning_up: return + self._resolve_pre_session_descendant(child_x_correlation_id) entries = self._child_to_join.get(child_x_correlation_id) if not entries: return @@ -370,6 +454,7 @@ async def on_child_stopped(self, child_x_correlation_id: str) -> None: """ if self._cleaning_up: return + self._resolve_pre_session_descendant(child_x_correlation_id) entries = self._child_to_join.get(child_x_correlation_id) if not entries: return @@ -383,6 +468,12 @@ async def _handle_child_done( self._child_to_join.pop(child_corr, None) # Every entry shares the same parent_correlation_id by construction. parent = entries[0].parent_correlation_id + # Registry descendant-done keys on the tree ROOT (not the parent), so + # depth>1 descendants still decrement the right tree. Retiring a drained + # 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) child_mode = self._child_modes.pop(child_corr, None) if ( child_mode == ConversationBranchMode.FORK @@ -421,6 +512,7 @@ async def on_child_errored(self, child_x_correlation_id: str) -> None: """ if self._cleaning_up: return + self._resolve_pre_session_descendant(child_x_correlation_id) entries = self._child_to_join.get(child_x_correlation_id) if not entries: return @@ -436,6 +528,9 @@ async def _handle_child_errored_fail_fast( parent = entries[0].parent_correlation_id errored_mode = self._child_modes.pop(child_corr, None) self._child_to_join.pop(child_corr, None) + errored_root = self._child_root.pop(child_corr, None) + if self._session_tree_registry is not None and errored_root is not None: + self._session_tree_registry.on_descendant_done(errored_root) # Collect all tracked children for this parent as potential orphans. orphans = [ @@ -459,6 +554,9 @@ async def _handle_child_errored_fail_fast( for orphan in orphans: self._child_to_join.pop(orphan, None) + orphan_root = self._child_root.pop(orphan, None) + if self._session_tree_registry is not None and orphan_root is not None: + self._session_tree_registry.on_descendant_done(orphan_root) orphan_mode = self._child_modes.pop(orphan, None) if ( orphan_mode == ConversationBranchMode.FORK @@ -508,9 +606,14 @@ def cleanup(self) -> None: self._abort_observer = None self._log_stats() self._log_leaks() + if self._session_tree_registry is not None: + self._session_tree_registry.release_all() self._active_joins.clear() self._future_joins.clear() self._child_to_join.clear() + self._child_root.clear() + self._pre_child_conv.clear() + self._pre_live_by_conv.clear() self._child_modes.clear() self._descendant_counts.clear() self._parent_locks.clear() diff --git a/src/aiperf/timing/conversation_source.py b/src/aiperf/timing/conversation_source.py index b301449a35..2972e56b2a 100644 --- a/src/aiperf/timing/conversation_source.py +++ b/src/aiperf/timing/conversation_source.py @@ -46,6 +46,10 @@ class SampledSession: parent's accumulated message history and pins to the same worker; SPAWN starts with a fresh context. Ignored when parent_correlation_id is None. + root_correlation_id: The x_correlation_id of the depth-0 root of this + session's tree. None on a depth-0 root (it is its own tree root); + inherited by children/subchildren so per-tree finality accounting + can key on ``effective_root_correlation_id``. """ conversation_id: str @@ -53,6 +57,7 @@ class SampledSession: x_correlation_id: str agent_depth: int = 0 parent_correlation_id: str | None = None + root_correlation_id: str | None = None branch_mode: ConversationBranchMode = ConversationBranchMode.FORK @property @@ -65,6 +70,16 @@ def routing_key(self) -> str: """ return self.parent_correlation_id or self.x_correlation_id + @property + def effective_root_correlation_id(self) -> str: + """Tree root id, defaulting to this session's own x_correlation_id. + + A root session is its own tree root; a child/subchild inherits the + depth-0 root's id, set by the spawning code so per-tree finality + accounting can key on it. + """ + return self.root_correlation_id or self.x_correlation_id + def build_first_turn(self, max_turns: int | None = None) -> TurnToSend: """Build first turn (turn_index=0) from sampled conversation. @@ -80,7 +95,11 @@ def build_first_turn(self, max_turns: int | None = None) -> TurnToSend: num_turns=max_turns or len(self.metadata.turns), agent_depth=self.agent_depth, parent_correlation_id=self.parent_correlation_id, + root_correlation_id=self.root_correlation_id, has_forks=first_meta.has_forks if first_meta is not None else False, + has_branches=bool(first_meta.branch_ids) + if first_meta is not None + else False, branch_mode=self.branch_mode, ) @@ -132,6 +151,7 @@ def start_branch_child( child_conversation_id: str, agent_depth: int, *, + root_correlation_id: str | None = None, branch_mode: ConversationBranchMode = ConversationBranchMode.FORK, ) -> SampledSession: """Build a SampledSession for a DAG child conversation. @@ -143,6 +163,11 @@ def start_branch_child( SPAWN-mode children start with a fresh context, but the sticky pin to the parent's correlation_id is preserved at this layer — routing freedom is enforced upstream by the orchestrator/router. + + ``root_correlation_id`` is the depth-0 root of the spawning parent's + tree; the child inherits it so all descendants of one root share a + single per-tree finality key. Defaults to ``parent_correlation_id`` + when not supplied (a depth-1 child's parent IS the root). """ metadata = self._metadata_lookup[child_conversation_id] return SampledSession( @@ -151,6 +176,7 @@ def start_branch_child( x_correlation_id=str(uuid.uuid4()), agent_depth=agent_depth, parent_correlation_id=parent_correlation_id, + root_correlation_id=root_correlation_id or parent_correlation_id, branch_mode=branch_mode, ) diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index 19dd63746f..c3c00d7e57 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -24,6 +24,7 @@ from aiperf.timing.phase.progress_tracker import PhaseProgressTracker from aiperf.timing.phase.stop_conditions import StopConditionChecker from aiperf.timing.ramping import Ramper, RamperConfig, RampType +from aiperf.timing.session_tree import SessionTreeRegistry from aiperf.timing.strategies.core import RateSettableProtocol from aiperf.timing.url_samplers import URLSelectionStrategyProtocol @@ -132,6 +133,14 @@ def __init__( lifecycle=self._lifecycle, counter=self._progress.counter, ) + # Per-tree finality ledger, shared by the issuer (opens trees + stamps + # finality) and the branch orchestrator (registers descendants + drives + # terminal transitions). Only built for DAG-shaped datasets; None + # elsewhere so finality stays conservative. Must precede the issuer so + # it can be injected at construction. + self._session_tree_registry = self._build_session_tree_registry( + conversation_source + ) self._credit_issuer = self._build_credit_issuer(url_selection_strategy) self._maybe_construct_branch_orchestrator(conversation_source) @@ -141,6 +150,19 @@ def __init__( self._was_cancelled = False self._rampers: list[Ramper] = [] + def _build_session_tree_registry( + self, conversation_source: ConversationSource + ) -> SessionTreeRegistry | None: + """Build the per-tree finality ledger for DAG-shaped datasets only. + + Mirrors the ``_maybe_construct_branch_orchestrator`` gate: a registry is + pointless without DAG fan-out (no descendants to track), and leaving it + None on non-DAG runs keeps ``_finality_for_issue`` conservative. + """ + if not self._is_dag_dataset(conversation_source.dataset_metadata): + return None + return SessionTreeRegistry() + def _build_credit_issuer( self, url_selection_strategy: URLSelectionStrategyProtocol | None ) -> CreditIssuer: @@ -156,6 +178,7 @@ def _build_credit_issuer( cancellation_policy=self._cancellation_policy, lifecycle=self._lifecycle, url_selection_strategy=url_selection_strategy, + session_tree_registry=self._session_tree_registry, ) def _maybe_construct_branch_orchestrator( @@ -177,6 +200,7 @@ def _maybe_construct_branch_orchestrator( conversation_source=conversation_source, credit_issuer=self._credit_issuer, sticky_router=sticky_router, + session_tree_registry=self._session_tree_registry, ) @property diff --git a/src/aiperf/timing/session_tree.py b/src/aiperf/timing/session_tree.py new file mode 100644 index 0000000000..5466423381 --- /dev/null +++ b/src/aiperf/timing/session_tree.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-session-TREE liveness bookkeeping for DAG-lineage finality. + +A DAG session is not a single root conversation -- it is a whole TREE: a +depth-0 root plus every child it spawns, recursively (FORK/SPAWN children and +their subchildren). ``SessionTreeRegistry`` tracks the liveness of each such +tree, keyed by the tree's ``root_correlation_id`` (the depth-0 root's +x_correlation_id), so the credit issuer can stamp two conservative facts on +every emitted ``Credit`` at issue time: + + - ``is_parent_final`` -- the credit's parent had already returned its terminal + turn (v1: determinable only when the parent IS the root). + - ``is_tree_final`` -- this credit is provably the last request the whole tree + will ever send. + +This registry is **finality bookkeeping only**. Unlike the agentic-replay +variant it was ported from, it does NOT own session-slot release: on main the +credit callback handler releases the root's session slot on the root's terminal +turn (``_release_slots_for_return``) exactly as before, and DAG children inherit +the root's slot. The registry never touches the concurrency manager and fires no +drain callback -- it only answers the two finality queries below from live tree +state. + +Wiring (DAG datasets only; the registry is None on non-DAG paths, so +``_finality_for_issue`` returns the conservative ``(None, False)``): + - ``CreditIssuer._open_session_tree`` -> :meth:`open_tree` when a root + session's first credit is issued (turn 0 session-slot acquisition). + - ``BranchOrchestrator`` -> :meth:`register_descendants` when a parent turn + spawns children, :meth:`on_descendant_done` at every descendant-terminal + point, and :meth:`on_root_terminal` from ``intercept`` after the root's + final-turn return is processed (so final-turn spawns register first). + - ``BranchOrchestrator.cleanup`` -> :meth:`release_all` at phase teardown to + drop any still-open trees. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from aiperf.common.aiperf_logger import AIPerfLogger +from aiperf.common.enums import CreditPhase + +_logger = AIPerfLogger(__name__) + + +@dataclass(slots=True) +class _TreeState: + """Liveness of one session tree.""" + + phase: CreditPhase + """Phase the tree was opened under; teardown targets the same phase.""" + root_pending: bool + """True until the root's terminal turn returns.""" + outstanding: int = 0 + """Descendants currently live or registered-pending (any depth).""" + released: bool = False + """Set once the tree has been retired; guards against double retire.""" + + @property + def drained(self) -> bool: + """True when the root is done and no descendant work remains.""" + return not self.root_pending and self.outstanding <= 0 + + +class SessionTreeRegistry: + """Tracks per-session-tree liveness for DAG-lineage finality queries. + + See the module docstring for the invariant and wiring. Single-threaded + asyncio: every method runs to completion between awaits, so the counter + mutations are atomic without locking. Finality bookkeeping only -- this + class never acquires or releases a concurrency slot. + """ + + def __init__(self) -> None: + self._trees: dict[str, _TreeState] = {} + # Descendants registered before their tree is opened. On main's flow a + # root's tree is opened at turn-0 issue, before any child can spawn (a + # spawn only happens on a parent turn RETURN), so this buffer is a + # defensive belt that keeps the branch's state transitions identical. + self._pending_descendants: dict[str, int] = {} + # Roots whose tree already RETIRED (drained), mapped to the phase they + # were opened under. A descendant's final-turn SPAWN registers its + # grandchildren AFTER that descendant's own on_descendant_done drained + # the tree (callback order: the child-completion decrement precedes the + # return-intercept that spawns). Without this ledger those grandchildren + # would buffer into a retired root that nothing drains, and is_tree_final + # could never answer True for them. register_descendants consults it to + # RESURRECT the tree root-terminal instead. Cleared at teardown. + self._retired_roots: dict[str, CreditPhase] = {} + # Peak simultaneously-open trees == peak tree concurrency; logged at + # teardown so any overshoot is visible. + self._peak_open: int = 0 + # A non-zero count means a descendant returned for a tree that was + # already retired -- i.e. the tree drained while that work was still in + # flight (premature drain). Surfaces span-overhang. + self._late_events: int = 0 + + def open_tree( + self, root_corr: str, phase: CreditPhase, *, root_pending: bool + ) -> None: + """Record a newly admitted tree when its root's first credit is issued. + + Args: + root_corr: the tree's ``root_correlation_id`` (depth-0 root's + x_correlation_id). + phase: phase the tree was opened under. + root_pending: True when a root credit is in flight that will reach a + terminal turn (always True on main -- open_tree fires only at a + root session start). + """ + existing = self._trees.get(root_corr) + if existing is not None and not existing.released: + # Duplicate open for a still-live tree: keep the original (its + # outstanding count is authoritative). Should not happen with + # correct wiring; log so a regression is visible. + _logger.warning( + lambda: f"open_tree for already-open tree root_corr={root_corr!r}; " + "ignoring duplicate" + ) + return + state = _TreeState(phase=phase, root_pending=root_pending) + state.outstanding += self._pending_descendants.pop(root_corr, 0) + # A freshly-opened tree supersedes any stale retired record for the same + # id (correlation ids are unique, so this is defensive). + self._retired_roots.pop(root_corr, None) + self._trees[root_corr] = state + if len(self._trees) > self._peak_open: + self._peak_open = len(self._trees) + + @property + def peak_open(self) -> int: + """Maximum simultaneously-open trees seen (== peak tree concurrency).""" + return self._peak_open + + @property + def late_events(self) -> int: + """Count of returns for already-retired trees (premature-drain evidence).""" + return self._late_events + + def has_tree(self, root_corr: str) -> bool: + """True when this registry is tracking ``root_corr`` (engagement gate).""" + return root_corr in self._trees + + def root_terminal(self, root_correlation_id: str) -> bool | None: + """True when the tree's root has returned its terminal turn; None if unknown. + + Queried at credit-issue time while the tree is still live (a descendant + being issued keeps it in ``_trees``). A fully drained tree has been + retired and reads as unknown/None. + """ + state = self._trees.get(root_correlation_id) + if state is None: + return None + return not state.root_pending + + def is_last_tree_request( + self, + root_correlation_id: str, + *, + is_final_turn: bool, + is_root_credit: bool, + has_branches: bool, + ) -> bool: + """Conservative: True only when this credit is provably the tree's last request. + + ``has_branches`` means "this turn declares ANY branch (FORK or SPAWN) + and will therefore spawn descendants on its return". It must NOT be the + FORK-only ``has_forks`` flag: SPAWN children register with this registry + only at return-intercept, AFTER the issuer stamped finality, so a + spawning final turn would otherwise read as last while its children are + pending (wrong-True, violating the conservative contract). + """ + if not is_final_turn or has_branches: + return False + state = self._trees.get(root_correlation_id) + if state is None: + return False + if is_root_credit: + return state.outstanding <= 0 + return (not state.root_pending) and state.outstanding == 1 + + def register_descendants(self, root_corr: str, n: int = 1) -> None: + """Add ``n`` descendants (spawned under one root) to a tree. + + Called when a parent turn spawns children, keyed by the tree's root id. + Three cases when the tree is not live: + - RETIRED root (in ``_retired_roots``): resurrect it root-terminal. + A descendant's final-turn SPAWN registers its grandchildren after + that descendant's own ``on_descendant_done`` drained the tree; the + root was terminal at retire (``drained`` requires ``not + root_pending``), so recreate the tree with ``root_pending=False`` + and these descendants outstanding -- keeping ``is_tree_final`` + answerable and the count coherent (the grandchildren re-drain it + when they finish). + - Unopened root: buffer the count so a later ``open_tree`` folds it in + (defensive -- on main the tree is always open first). + """ + if n <= 0: + return + state = self._trees.get(root_corr) + if state is not None: + state.outstanding += n + return + retired_phase = self._retired_roots.pop(root_corr, None) + if retired_phase is not None: + self._trees[root_corr] = _TreeState( + phase=retired_phase, root_pending=False, outstanding=n + ) + if len(self._trees) > self._peak_open: + self._peak_open = len(self._trees) + return + self._pending_descendants[root_corr] = ( + self._pending_descendants.get(root_corr, 0) + n + ) + + def on_descendant_done(self, root_corr: str) -> bool: + """Account one descendant terminally completing (leaf / error / stopped / + dispatch-rollback). Retires the tree iff it is now drained. + + Returns True if the tree was retired by this call. + """ + state = self._trees.get(root_corr) + if state is None: + pending = self._pending_descendants.get(root_corr, 0) + if pending > 0: + self._pending_descendants[root_corr] = pending - 1 + else: + self._late_events += 1 + return False + if state.outstanding > 0: + state.outstanding -= 1 + return self._retire_if_drained(root_corr, state) + + def on_root_terminal(self, root_corr: str) -> bool: + """Account the root's terminal turn returning (or root cancel/truncate). + + Clears ``root_pending``; retires the tree iff it is now drained. Must be + called AFTER the returning root credit's intercept has run, so any + children spawned on the final turn are already registered. + + Returns True if the tree was retired by this call. False both when the + tree is held (descendants remain) AND when ``root_corr`` is untracked. + """ + state = self._trees.get(root_corr) + if state is None: + return False + state.root_pending = False + return self._retire_if_drained(root_corr, state) + + def _retire_if_drained(self, root_corr: str, state: _TreeState) -> bool: + """Pop a drained tree from tracking. Finality-only: releases no slot.""" + if state.released or not state.drained: + return False + state.released = True + self._trees.pop(root_corr, None) + # Remember the retired root (with its phase) so a late final-turn SPAWN + # can RESURRECT it instead of silently buffering. See register_descendants. + self._retired_roots[root_corr] = state.phase + return True + + def release_all(self, phase: CreditPhase | None = None) -> int: + """Retire every still-open tree at teardown (optionally one phase). + + Finality-only: drops the tracking entries; releases no concurrency slot + (the callback handler owns slot release on main). The registry is + created per-phase, so ``phase=None`` (retire all) is the common + teardown call. Returns the number of trees retired. + """ + to_release = [ + root_corr + for root_corr, state in self._trees.items() + if (phase is None or state.phase == phase) and not state.released + ] + for root_corr in to_release: + state = self._trees.pop(root_corr, None) + if state is None or state.released: + continue + state.released = True + # Drop the transient buffers: with the trees above retired, any + # pending/retired descendant accounting refers to no live tree. On a + # full teardown (``phase is None`` -- the common per-phase call, since + # the registry is created per-phase) clear both; a phase-scoped call + # clears only that phase's retired roots (``_pending_descendants`` + # carries no phase and is left untouched then). + if phase is None: + self._pending_descendants.clear() + self._retired_roots.clear() + else: + self._retired_roots = { + r: p for r, p in self._retired_roots.items() if p != phase + } + return len(to_release) + + def open_count(self, phase: CreditPhase | None = None) -> int: + """Number of trees currently tracked (optionally for one phase).""" + if phase is None: + return len(self._trees) + return sum(1 for state in self._trees.values() if state.phase == phase) diff --git a/src/aiperf/timing/strategies/fixed_schedule.py b/src/aiperf/timing/strategies/fixed_schedule.py index a557130555..49b2d8a243 100644 --- a/src/aiperf/timing/strategies/fixed_schedule.py +++ b/src/aiperf/timing/strategies/fixed_schedule.py @@ -100,6 +100,9 @@ async def setup_phase(self) -> None: x_correlation_id=str(uuid.uuid4()), turn_index=0, num_turns=len(conv.turns), + # Finality stamping must see a spawning turn as + # non-final (any-mode branch flag). + has_branches=bool(conv.turns[0].branch_ids), ), ) ) @@ -151,9 +154,11 @@ async def handle_credit_return( if credit.is_final_turn: return - # This contains the delay_ms or timestamp_ms for the next turn + # This contains the delay_ms or timestamp_ms for the next turn. + # Passed through so has_forks / has_branches derive from the real + # metadata (a spawning final turn must never stamp is_tree_final). next_meta = self._conversation_source.get_next_turn_metadata(credit) - turn = TurnToSend.from_previous_credit(credit) + turn = TurnToSend.from_previous_credit(credit, next_meta) if next_meta.timestamp_ms is not None: self._scheduler.schedule_at_perf_sec( diff --git a/src/aiperf/timing/strategies/user_centric_rate.py b/src/aiperf/timing/strategies/user_centric_rate.py index 360e3770c2..c8dc05805d 100644 --- a/src/aiperf/timing/strategies/user_centric_rate.py +++ b/src/aiperf/timing/strategies/user_centric_rate.py @@ -340,7 +340,11 @@ async def handle_credit_return( raise ValueError( f"User not found for x_correlation_id: {credit.x_correlation_id}" ) - turn = TurnToSend.from_previous_credit(credit) + # Pass the next turn's metadata through so has_forks / has_branches + # derive from the real metadata (a spawning turn must never stamp + # is_tree_final; a branching turn must stay conservative). + next_meta = self._conversation_source.get_next_turn_metadata(credit) + turn = TurnToSend.from_previous_credit(credit, next_meta) # If the next turn time already passed, the max() will # re-align their schedule to account for the delay. diff --git a/src/aiperf/workers/inference_client.py b/src/aiperf/workers/inference_client.py index b376631526..870c42841d 100644 --- a/src/aiperf/workers/inference_client.py +++ b/src/aiperf/workers/inference_client.py @@ -20,6 +20,7 @@ from aiperf.common.redact import redact_headers from aiperf.plugin import plugins from aiperf.plugin.enums import PluginType, TransportType +from aiperf.workers.session_routing import RoutingContext, SessionRoutingBase if TYPE_CHECKING: from aiperf.transports.base_transports import FirstTokenCallback @@ -61,6 +62,21 @@ def __init__(self, model_endpoint: ModelEndpointInfo, service_id: str, **kwargs) self.model_endpoint = model_endpoint self.service_id = service_id + # Session-routing plugin (selected via --session-routing): one instance + # per worker, invoked at the request-serialization chokepoint to stamp + # per-session identity (headers and/or body). None when routing is off. + self._routing: SessionRoutingBase | None = None + self._routing_mode: str | None = None + endpoint_info = model_endpoint.endpoint + if endpoint_info.session_routing is not None: + routing_cls = plugins.get_class( + PluginType.SESSION_ROUTING, endpoint_info.session_routing + ) + self._routing = routing_cls( + routing_cls.Options(**endpoint_info.session_routing_opts) + ) + self._routing_mode = endpoint_info.session_routing + # Detect and set transport type if not explicitly set if not model_endpoint.transport: model_endpoint.transport = TransportType( @@ -78,6 +94,29 @@ def __init__(self, model_endpoint: ModelEndpointInfo, service_id: str, **kwargs) self.transport = TransportClass(model_endpoint=self.model_endpoint) self.attach_child_lifecycle(self.transport) + def notify_session_end(self, x_correlation_id: str) -> None: + """Post-session pass-through to the routing plugin (idempotent hook). + + Called by the worker terminal-eviction paths on ANY terminal outcome of + a session. On this codebase those are: a successful final turn, a + cancellation, and a cancel-before-start (the done-callback path whose + finally block never runs). Idempotency is the plugin's responsibility -- + this hook does not dedupe. No-op when session routing is unset. + + A plugin exception is logged (naming the plugin and session) and + swallowed: this cleanup hook must never break the worker's core + session-eviction lifecycle. + """ + if self._routing is None: + return + try: + self._routing.on_session_end(x_correlation_id) + except Exception as e: + self.warning( + f"session-routing plugin {self._routing_mode!r} on_session_end " + f"failed for session {x_correlation_id!r}; continuing eviction: {e!r}" + ) + async def _send_request_to_transport( self, request_info: RequestInfo, @@ -102,12 +141,48 @@ async def _send_request_to_transport( """ request_info.endpoint_headers = self.endpoint.get_endpoint_headers(request_info) request_info.endpoint_params = self.endpoint.get_endpoint_params(request_info) + + # Session-routing chokepoint: build the per-request routing context once + # and let the plugin stamp its headers now (merged onto the endpoint + # headers). The same context feeds the structured body transform below. + routing_ctx: RoutingContext | None = None + if self._routing is not None: + routing_ctx = RoutingContext( + x_correlation_id=request_info.x_correlation_id, + parent_correlation_id=request_info.parent_correlation_id, + root_correlation_id=request_info.root_correlation_id, + is_final_turn=request_info.is_final_turn, + is_parent_final=request_info.is_parent_final, + is_tree_final=request_info.is_tree_final, + ) + # Attribute a plugin fault to the routing plugin (not the server): + # this raise is caught by _send_request_internal and becomes an error + # record whose message names the plugin instead of the endpoint. + try: + routing_headers = self._routing.headers(routing_ctx) + except Exception as e: + raise RuntimeError( + f"session-routing plugin {self._routing_mode!r} failed in headers(): {e!r}" + ) from e + request_info.endpoint_headers.update(routing_headers) + raw_payload = request_info.turns[-1].raw_payload payload = ( raw_payload if raw_payload is not None else self.endpoint.format_payload(request_info) ) + # Body-based session routing (e.g. Dynamo nvext.session_control): overlay + # 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): + try: + payload = self._routing.transform_body(payload, routing_ctx) + except Exception as e: + raise RuntimeError( + f"session-routing plugin {self._routing_mode!r} failed in transform_body(): {e!r}" + ) from e request_info.payload_bytes = orjson.dumps(payload) return await self.transport.send_request( request_info, diff --git a/src/aiperf/workers/session_routing.py b/src/aiperf/workers/session_routing.py new file mode 100644 index 0000000000..5cd782f943 --- /dev/null +++ b/src/aiperf/workers/session_routing.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Session-routing transforms: per-session identity on the wire. + +One plugin category unifies every mechanism that tells an external router +which session a request belongs to, whether header-based (SGLang Model +Gateway routing keys, Dynamo session headers, generic session-ID headers) or +body-based (Dynamo ``nvext.session_control``). The selected plugin is +instantiated once per worker by ``InferenceClient`` and invoked at the +request-serialization chokepoint. + +Contracts: +- Options instances are canonicalized at config resolution; ``self.options`` + is the plugin's own typed model. +- ``transform_body`` must NEVER mutate its input (the structured path + includes cached ``Turn.raw_payload`` dicts shared with the dataset). +- ``on_session_end`` fires strictly AFTER the session's last worker-side + activity, on every terminal path (final turn, cancellation, and + cancel-before-start via the done-callback path). It MUST be idempotent. +- Stateful plugins key instance state on ``ctx.x_correlation_id`` ONLY: + a session tree deliberately spans workers, so tree-keyed worker state + fragments. Tree-scoped behavior uses the stateless per-request facts + (``root_correlation_id`` + ``is_tree_final``) instead. +""" + +from __future__ import annotations + +from abc import ABC +from dataclasses import dataclass +from typing import Any, ClassVar, Generic, TypeVar + +from pydantic import ConfigDict, Field + +from aiperf.common.models import AIPerfBaseModel + + +class EmptyRoutingOptions(AIPerfBaseModel): + """Options model for parameterless plugins; rejects every opt key.""" + + model_config = ConfigDict(extra="forbid") + + +OptionsT = TypeVar("OptionsT", bound=AIPerfBaseModel) + + +@dataclass(slots=True, frozen=True) +class RoutingContext: + """Per-request identity facts handed to a routing plugin. + + Field naming mirrors ``RequestInfo`` verbatim. ``is_parent_final`` / + ``is_tree_final`` are stamped issuer-side from ``SessionTreeRegistry`` + state and are conservative: ``is_tree_final`` is False whenever + indeterminate, ``is_parent_final`` is None for roots or when unknown. + """ + + x_correlation_id: str + """This session's stable key (same on every turn).""" + parent_correlation_id: str | None + """Immediate parent session's key; None for root sessions.""" + root_correlation_id: str | None + """Session-tree root key, verbatim from RequestInfo (never None on the + dispatch path: the worker passes ``credit.effective_root_correlation_id``).""" + is_final_turn: bool + """True when this is the current session's last request.""" + is_parent_final: bool | None + """True when the parent session had already returned its final turn + at credit-issue time; None for roots or when not determinable.""" + is_tree_final: bool + """Best-effort: True only when this is provably the last request the + whole session tree will send.""" + + +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 + """Protocol metadata: True when ``transform_body`` changes the payload, + marking the mode as body-mutating / incompatible with any verbatim-bytes + request path. NOT currently enforced by any gate on this codebase -- no such + verbatim-bytes send path exists here today (the chokepoint always + re-serializes). It exists so a future pre-serialized-bytes path can refuse a + body-mutating mode instead of silently dropping its transform.""" + + # ClassVar cannot reference the OptionsT type parameter, so the base type + # is kept here; subclasses narrow it via their Generic parameterization. + Options: ClassVar[type[AIPerfBaseModel]] = EmptyRoutingOptions + """Per-plugin options model, populated from --session-routing-opt + key=value pairs. Every Options model must set extra='forbid'.""" + + def __init__(self, options: OptionsT) -> None: + self.options: OptionsT = options + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + """Extra HTTP headers for this request (merged into endpoint headers).""" + return {} + + def transform_body( + self, payload: dict[str, Any], ctx: RoutingContext + ) -> dict[str, Any]: + """Return a (possibly new) payload dict; never mutate the input.""" + return payload + + def on_session_end(self, x_correlation_id: str) -> None: + """Post-session cleanup: no further requests will be sent for this + session by this worker. Idempotent; default no-op.""" + return None + + +class DynamoHeadersRouting(SessionRoutingBase[EmptyRoutingOptions]): + """Dynamo session affinity via X-Dynamo-Session-ID / X-Dynamo-Parent-Session-ID. + + Pair with a Dynamo frontend running ``--router-session-affinity-ttl-secs``. + """ + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + headers = {"X-Dynamo-Session-ID": ctx.x_correlation_id} + if ctx.parent_correlation_id: + headers["X-Dynamo-Parent-Session-ID"] = ctx.parent_correlation_id + return headers + + +class DynamoNvextOptions(AIPerfBaseModel): + """Options for the deprecated-upstream nvext.session_control transport.""" + + model_config = ConfigDict(extra="forbid") + + timeout_seconds: int = Field( + default=300, + ge=1, + description="Dynamo session_control inactivity timeout carried on every bind.", + ) + + +class DynamoNvextRouting(SessionRoutingBase[DynamoNvextOptions]): + """Dynamo session affinity via nvext.session_control request-body metadata. + + Modern contract only: 'bind' on every non-final turn (idempotent on the + router, refreshes the TTL), 'close' on the final turn. Targets Dynamo + builds that implement session_control; current upstream Dynamo main does + not (use dynamo_headers there). + """ + + mutates_body: ClassVar[bool] = True + Options: ClassVar[type[AIPerfBaseModel]] = DynamoNvextOptions + + def transform_body( + self, payload: dict[str, Any], ctx: RoutingContext + ) -> dict[str, Any]: + if ctx.is_final_turn: + session_control: dict[str, Any] = { + "session_id": ctx.x_correlation_id, + "action": "close", + } + else: + session_control = { + "session_id": ctx.x_correlation_id, + "action": "bind", + "timeout": self.options.timeout_seconds, + } + merged = dict(payload) + raw_nvext = merged.get("nvext") + nvext = dict(raw_nvext) if isinstance(raw_nvext, dict) else {} + raw_sc = nvext.get("session_control") + merged_sc = dict(raw_sc) if isinstance(raw_sc, dict) else {} + merged_sc.update(session_control) + nvext["session_control"] = merged_sc + merged["nvext"] = nvext + return merged + + +class SmgRoutingKeyRouting(SessionRoutingBase[EmptyRoutingOptions]): + """SGLang Model Gateway manual-policy stickiness via X-SMG-Routing-Key.""" + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + return {"X-SMG-Routing-Key": ctx.x_correlation_id} + + +class SessionIdHeaderOptions(AIPerfBaseModel): + """Options for the generic additive session-ID header.""" + + model_config = ConfigDict(extra="forbid") + + header_name: str = Field( + default="X-Session-ID", + min_length=1, + description="Header name carrying the per-session correlation ID.", + ) + + +class SessionIdHeaderRouting(SessionRoutingBase[SessionIdHeaderOptions]): + """Generic additive session header for routers expecting a custom name.""" + + Options: ClassVar[type[AIPerfBaseModel]] = SessionIdHeaderOptions + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + return {self.options.header_name: ctx.x_correlation_id} diff --git a/src/aiperf/workers/worker.py b/src/aiperf/workers/worker.py index 79cb7bb40b..dbf1deeef6 100644 --- a/src/aiperf/workers/worker.py +++ b/src/aiperf/workers/worker.py @@ -368,6 +368,14 @@ def _on_credit_drop_message_task_done( self.execute_async(self.credit_return_push_client.send(credit_return)) credit_context.returned = True + # Post-session hook for routing plugins: this cancel-before-start path is a + # terminal disposition the finally block never sees (the credit task was + # cancelled before it started), so notify here too. Idempotent by contract. + if credit_context.credit is not None: + self.inference_client.notify_session_end( + credit_context.credit.x_correlation_id + ) + # Explicitly clear references to help refcounting (GC is disabled on workers) credit_context.credit = None credit_context.error = None @@ -604,6 +612,11 @@ def _release_and_evict_for_terminal( Non-FORK and non-parent sessions evict immediately. """ + # Fire the routing plugin's post-session hook on ANY terminal outcome + # (final turn or cancel), not just a successful final turn, so stateful + # plugins release sessions abandoned mid-conversation. Idempotent by + # contract; no-op when session routing is unset. + self.inference_client.notify_session_end(x_correlation_id) if ( credit.parent_correlation_id is not None and credit.branch_mode == ConversationBranchMode.FORK @@ -662,6 +675,9 @@ def _create_request_info( is_final_turn=credit.is_final_turn, agent_depth=credit.agent_depth, parent_correlation_id=credit.parent_correlation_id, + root_correlation_id=credit.effective_root_correlation_id, + is_parent_final=credit.is_parent_final, + is_tree_final=credit.is_tree_final, # Use session's url_index to ensure all turns hit the same backend url_index=session.url_index, ) diff --git a/tests/integration/test_session_routing_raw_export.py b/tests/integration/test_session_routing_raw_export.py new file mode 100644 index 0000000000..a5d16e7fe5 --- /dev/null +++ b/tests/integration/test_session_routing_raw_export.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end raw-export proof that every ``--session-routing`` mode reaches the wire. + +Runs a real ``aiperf profile`` subprocess against the in-repo mock server with +``--export-level raw`` for each session-routing mode, then reads the exported +wire payloads / request headers and asserts the per-mode contract: + +- ``dynamo_headers``: ``X-Dynamo-Session-ID`` on every request equals the + session's ``x_correlation_id``; no parent header on root sessions; body is + untouched (``"nvext" not in payload``). +- ``dynamo_nvext``: ``nvext.session_control`` carries ``bind`` (+ timeout) on + every non-final turn and ``close`` (no timeout) on the final turn, with one + stable ``session_id == x_correlation_id`` across the session. +- ``smg_routing_key``: ``X-SMG-Routing-Key`` equals ``x_correlation_id``. +- ``session_id_header`` (``--session-routing-opt header_name=X-Affinity``): the + configured header equals ``x_correlation_id``. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable, Sequence + +import pytest +from pytest import param + +from aiperf.common.models import RawRecordInfo +from tests.harness.utils import AIPerfCLI, AIPerfMockServer + +_TOKENIZER = "openai/gpt-oss-120b" # pre-cached + offline in integration conftest + +_NUM_SESSIONS = 2 +_TURNS_PER_SESSION = 3 +# Non-default (plugin default is 300) so the assertions prove the CLI -> +# config -> worker numeric plumb actually carries this value to the wire, +# rather than coinciding with the default. +_TIMEOUT_SECONDS = 123 + + +def _build_cmd(url: str, *, mode: str, opts: Sequence[str]) -> str: + opt_flags = " ".join(f"--session-routing-opt {kv}" for kv in opts) + return f""" + aiperf profile \ + --model {_TOKENIZER} \ + --url {url} \ + --endpoint-type chat \ + --num-sessions {_NUM_SESSIONS} \ + --session-turns-mean {_TURNS_PER_SESSION} \ + --session-turns-stddev 0 \ + --random-seed 42 \ + --workers-max 1 \ + --session-routing {mode} \ + {opt_flags} \ + --export-level raw \ + --ui simple + """ + + +async def _records_by_session( + cli: AIPerfCLI, url: str, *, mode: str, opts: Sequence[str] +) -> dict[str, list[RawRecordInfo]]: + """Run a benchmark and return each session's raw records ordered by + 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 []) + assert records, f"no raw records\n{(result.log or '')[-1500:]}" + + grouped: dict[str, list[RawRecordInfo]] = defaultdict(list) + for rec in records: + grouped[rec.metadata.x_correlation_id].append(rec) + + assert len(grouped) == _NUM_SESSIONS, ( + f"expected {_NUM_SESSIONS} sessions, got {len(grouped)}" + ) + out: dict[str, list[RawRecordInfo]] = {} + for xcorr, recs in grouped.items(): + assert len(recs) == _TURNS_PER_SESSION, ( + f"session {xcorr}: expected {_TURNS_PER_SESSION} turns, got {len(recs)}" + ) + recs.sort(key=lambda r: r.metadata.turn_index) + out[xcorr] = recs + return out + + +def _verify_dynamo_headers(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + assert headers.get("X-Dynamo-Session-ID") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + # Root sessions have no parent, so the parent header must be absent. + assert "X-Dynamo-Parent-Session-ID" not in headers, ( + f"session {xcorr}: root session emitted a parent header; {headers}" + ) + assert "nvext" not in rec.payload, ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_smg_routing_key(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + assert headers.get("X-SMG-Routing-Key") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + assert "nvext" not in rec.payload, ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_session_id_header(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + assert headers.get("X-Affinity") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + assert "nvext" not in rec.payload, ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_dynamo_nvext(by_session: dict[str, list[RawRecordInfo]]) -> None: + for xcorr, recs in by_session.items(): + scs = [] + for rec in recs: + sc = rec.payload.get("nvext", {}).get("session_control") + assert sc is not None, ( + f"session {xcorr} turn {rec.metadata.turn_index}: " + "every request must carry nvext.session_control" + ) + scs.append(sc) + + actions = [sc.get("action") for sc in scs] + assert all(a == "bind" for a in actions[:-1]), ( + f"session {xcorr}: non-final turns must bind; {actions}" + ) + assert actions[-1] == "close", f"session {xcorr}: actions={actions}" + assert "open" not in actions, f"session {xcorr}: emitted open; {actions}" + + # Every non-final 'bind' carries the timeout; 'close' does not. + for sc in scs[:-1]: + assert sc["timeout"] == _TIMEOUT_SECONDS, f"session {xcorr}: {sc}" + assert "timeout" not in scs[-1], ( + f"session {xcorr}: close carried timeout; {scs[-1]}" + ) + + # One stable session_id == the X-Correlation-ID, on every turn. + assert {sc["session_id"] for sc in scs} == {xcorr}, ( + f"session {xcorr}: session_id drift; {[sc.get('session_id') for sc in scs]}" + ) + + +_VERIFIERS: dict[str, Callable[[dict[str, list[RawRecordInfo]]], None]] = { + "dynamo_headers": _verify_dynamo_headers, + "dynamo_nvext": _verify_dynamo_nvext, + "smg_routing_key": _verify_smg_routing_key, + "session_id_header": _verify_session_id_header, +} + + +@pytest.mark.integration +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode, opts", + [ + param("dynamo_headers", (), id="dynamo_headers"), + param("dynamo_nvext", (f"timeout_seconds={_TIMEOUT_SECONDS}",), id="dynamo_nvext"), + param("smg_routing_key", (), id="smg_routing_key"), + param("session_id_header", ("header_name=X-Affinity",), id="session_id_header"), + ], +) # fmt: skip +async def test_session_routing_mode_reaches_wire( + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + mode: str, + opts: tuple[str, ...], +): + """Each session-routing mode stamps its per-session identity on the wire and + it survives to the raw export exactly as the plugin specifies.""" + by_session = await _records_by_session( + cli, aiperf_mock_server.url, mode=mode, opts=opts + ) + _VERIFIERS[mode](by_session) diff --git a/tests/unit/config/test_session_routing_config.py b/tests/unit/config/test_session_routing_config.py new file mode 100644 index 0000000000..83c3f22c7a --- /dev/null +++ b/tests/unit/config/test_session_routing_config.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pytest import param + +from aiperf.config.endpoint import EndpointConfig +from aiperf.config.flags._converter_endpoint import _parse_routing_opts + + +def _config(**kwargs) -> EndpointConfig: + return EndpointConfig(urls=["http://localhost:8000"], **kwargs) + + +def test_defaults_off(): + config = _config() + assert config.session_routing is None + assert config.session_routing_opts == {} + + +def test_mode_with_valid_opts(): + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "600"}, + ) + assert str(config.session_routing) == "dynamo_nvext" + + +def test_opts_canonicalized_to_typed_values(): + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "600"}, + ) + assert config.session_routing_opts == {"timeout_seconds": 600} + assert isinstance(config.session_routing_opts["timeout_seconds"], int) + + +def test_canonicalization_idempotent(): + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": 600}, + ) + assert config.session_routing_opts == {"timeout_seconds": 600} + + +def test_opts_without_mode_rejected(): + with pytest.raises(ValueError, match="session-routing-opt"): + _config(session_routing_opts={"header_name": "X-A"}) + + +def test_unknown_opt_key_rejected(): + with pytest.raises(ValueError, match="timeout_secs"): + _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_secs": "600"}, + ) + + +def test_invalid_opt_value_rejected(): + with pytest.raises(ValueError): + _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "0"}, + ) + + +def test_parameterless_mode_rejects_any_opt(): + with pytest.raises(ValueError): + _config( + session_routing="smg_routing_key", + session_routing_opts={"anything": "x"}, + ) + + +def test_parse_routing_opts_duplicate_key_rejected(): + with pytest.raises(ValueError, match="Duplicate"): + _parse_routing_opts(["header_name=X-A", "header_name=X-B"]) + + +@pytest.mark.parametrize( + "item", + [ + param("noequals", id="no_separator"), + param("key=", id="empty_value"), + ], +) # fmt: skip +def test_parse_routing_opts_malformed_pair_rejected(item): + with pytest.raises(ValueError, match="expected non-empty key=value"): + _parse_routing_opts([item]) diff --git a/tests/unit/credit/test_issuer_finality.py b/tests/unit/credit/test_issuer_finality.py new file mode 100644 index 0000000000..4306b0006e --- /dev/null +++ b/tests/unit/credit/test_issuer_finality.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Guards the credit issuer's lineage-finality stamp. + +Touch #1 (the ``Credit`` struct fields) and touch #3 (worker -> ``RequestInfo``, +``test_create_request_info_plumbs_finality_from_credit``) are covered elsewhere. +This file covers the issuer touch: ``CreditIssuer`` reading per-tree state from +a REAL ``SessionTreeRegistry`` (``_finality_for_issue``) AND stamping the result +onto the emitted ``Credit`` at its sole construction site +(``_issue_credit_internal``). + +Both the registry and the emitted ``Credit`` are real objects here on purpose -- +a MagicMock would auto-create ``is_parent_final`` / ``is_tree_final`` and pass +even if the ``Credit(...)`` kwargs were deleted, defeating the guard. +""" + +import time +from unittest.mock import MagicMock + +import pytest + +from aiperf.common.enums import CreditPhase +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.structs import Credit, TurnToSend +from aiperf.timing.session_tree import SessionTreeRegistry + + +class _FakeConcurrency: + """Slots always granted; releases are no-ops. + + The main ``SessionTreeRegistry`` needs no concurrency manager at all; the + issuer's acquire path needs the two coroutines. Neither the registry nor the + emitted ``Credit`` is a mock -- that is the point of this file. + """ + + async def acquire_session_slot(self, phase: CreditPhase, can_proceed) -> bool: + return True + + async def acquire_prefill_slot(self, phase: CreditPhase, can_proceed) -> bool: + return True + + def release_session_slot(self, phase: CreditPhase) -> None: + pass + + +class _CapturingRouter: + """Captures the emitted ``Credit`` so tests can assert its stamped finality.""" + + def __init__(self) -> None: + self.sent: list[Credit] = [] + + async def send_credit(self, *, credit: Credit) -> None: + self.sent.append(credit) + + +def _make_registry() -> SessionTreeRegistry: + return SessionTreeRegistry() + + +def _make_issuer( + registry: SessionTreeRegistry | None, +) -> tuple[CreditIssuer, _CapturingRouter]: + """Minimal real issuer: mocked scalars/lifecycle, REAL registry + router. + + ``session_tree_registry_enabled=True`` engages ``registry`` regardless of + phase; passing ``None`` for the registry yields the non-DAG path. + """ + progress = MagicMock() + progress.increment_sent = MagicMock(return_value=(1, False)) + progress.freeze_sent_counts = MagicMock() + progress.all_credits_sent_event = MagicMock() + + stop_checker = MagicMock() + stop_checker.can_send_any_turn = MagicMock(return_value=True) + stop_checker.can_start_new_session = MagicMock(return_value=True) + stop_checker.can_send_dag_child_turn = MagicMock(return_value=True) + + cancellation = MagicMock() + cancellation.next_cancellation_delay_ns = MagicMock(return_value=None) + + lifecycle = MagicMock() + lifecycle.started_at_ns = time.time_ns() + lifecycle.started_at_perf_ns = time.perf_counter_ns() + + router = _CapturingRouter() + issuer = CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=stop_checker, + progress=progress, + concurrency_manager=_FakeConcurrency(), + credit_router=router, + cancellation_policy=cancellation, + lifecycle=lifecycle, + session_tree_registry=registry, + session_tree_registry_enabled=True, + ) + return issuer, router + + +def _root_turn(root_id: str = "root-1") -> TurnToSend: + """Depth-0 root, single-turn (final), no forks.""" + return TurnToSend( + conversation_id="conv-1", + x_correlation_id=root_id, + turn_index=0, + num_turns=1, + ) + + +def _child_turn(root_id: str = "root-1", child_id: str = "child-1") -> TurnToSend: + """Child whose parent IS the root, single-turn (final), no forks.""" + return TurnToSend( + conversation_id="conv-1", + x_correlation_id=child_id, + turn_index=0, + num_turns=1, + agent_depth=1, + parent_correlation_id=root_id, + root_correlation_id=root_id, + ) + + +# ============================================================================= +# _finality_for_issue: reads REAL SessionTreeRegistry state +# ============================================================================= + + +def test_finality_root_final_turn_no_descendants_is_tree_final(): + """Scenario 1: root, final turn, no descendants, no forks.""" + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + issuer, _ = _make_issuer(registry) + + is_parent_final, is_tree_final = issuer._finality_for_issue(_root_turn()) + + assert is_parent_final is None + assert is_tree_final is True + + +def test_finality_root_with_outstanding_descendant_not_tree_final(): + """Scenario 2: root, final turn, one outstanding descendant -> not last.""" + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + registry.register_descendants("root-1", n=1) + issuer, _ = _make_issuer(registry) + + is_parent_final, is_tree_final = issuer._finality_for_issue(_root_turn()) + + assert is_parent_final is None + assert is_tree_final is False + + +def test_finality_sole_child_after_root_terminal_is_both_final(): + """Scenario 3: child whose parent is the root; root terminal; sole + outstanding child on its final turn -> both facts True.""" + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + registry.register_descendants("root-1", n=1) + registry.on_root_terminal("root-1") # root_pending cleared; tree still live + issuer, _ = _make_issuer(registry) + + is_parent_final, is_tree_final = issuer._finality_for_issue(_child_turn()) + + assert is_parent_final is True + assert is_tree_final is True + + +def test_finality_no_registry_is_conservative_none_false(): + """Scenario 4: no registry engaged -> conservative ``(None, False)``.""" + issuer, _ = _make_issuer(None) + + assert issuer._finality_for_issue(_root_turn()) == (None, False) + + +def test_finality_spawning_final_turn_never_tree_final(): + """Scenario 5 (regression): a final turn declaring ANY branch is never + tree-final, even with nothing outstanding in the registry. + + SPAWN children register only at return-intercept, AFTER this issue-time + stamp, so ``has_branches`` (any-mode) must gate the query -- the FORK-only + ``has_forks`` flag stays False for a SPAWN-declaring turn and previously + produced a wrong ``is_tree_final=True``. + """ + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + issuer, _ = _make_issuer(registry) + + # Same registry state as scenario 1 (which yields True) but the turn + # declares a branch: SPAWN-shaped, so has_forks stays False. + spawning_root_final = TurnToSend( + conversation_id="conv-1", + x_correlation_id="root-1", + turn_index=0, + num_turns=1, + has_forks=False, + has_branches=True, + ) + is_parent_final, is_tree_final = issuer._finality_for_issue(spawning_root_final) + + assert is_parent_final is None + assert is_tree_final is False + + +# ============================================================================= +# GUARD: the Credit(...) construction site must pass the helper's results through +# ============================================================================= + + +@pytest.mark.asyncio +async def test_issue_credit_stamps_finality_onto_emitted_credit(): + """RED if either ``is_parent_final=`` / ``is_tree_final=`` kwarg is removed + from the ``Credit(...)`` construction in ``_issue_credit_internal``. + + Uses scenario 3 (both facts True) so the stamped values differ from the + struct defaults (``None`` / ``False``): a dropped kwarg reverts the emitted + ``Credit`` to the default and this assertion fails. + """ + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + registry.register_descendants("root-1", n=1) + registry.on_root_terminal("root-1") + issuer, router = _make_issuer(registry) + + # Child inherits the root's slot; dispatch_child_turn is the wire path that + # runs _issue_credit_internal for a DAG child (issue_credit's session-slot + # path is for roots). + await issuer.dispatch_child_turn(_child_turn()) + + assert len(router.sent) == 1 + credit = router.sent[0] + assert credit.is_parent_final is True + assert credit.is_tree_final is True diff --git a/tests/unit/dataset/loader/test_exgentic.py b/tests/unit/dataset/loader/test_exgentic.py index b6c4ddb86f..f2b6513848 100644 --- a/tests/unit/dataset/loader/test_exgentic.py +++ b/tests/unit/dataset/loader/test_exgentic.py @@ -217,10 +217,6 @@ async def test_convert_preserves_snapshots_tools_osl_order_and_delays() -> None: }, } ] - assert all( - turn.extra_headers == {"x-dynamo-session-id": "session-1"} - for turn in conversation.turns - ) assert all( "recorded output" not in orjson.dumps(turn.raw_messages).decode() for turn in conversation.turns @@ -265,10 +261,6 @@ async def test_fixed_schedule_preserves_overlapping_start_times() -> None: ] assert all(len(conversation.turns) == 1 for conversation in conversations) assert all(conversation.turns[0].delay is None for conversation in conversations) - assert all( - conversation.turns[0].extra_headers == {"x-dynamo-session-id": "session-1"} - for conversation in conversations - ) @pytest.mark.asyncio diff --git a/tests/unit/plugin/test_session_routing_registry.py b/tests/unit/plugin/test_session_routing_registry.py new file mode 100644 index 0000000000..1d2a2a809c --- /dev/null +++ b/tests/unit/plugin/test_session_routing_registry.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pytest import param + +from aiperf.plugin import plugins +from aiperf.plugin.enums import PluginType +from aiperf.workers.session_routing import SessionRoutingBase + + +@pytest.mark.parametrize( + "name", + [ + param("dynamo_headers", id="dynamo_headers"), + param("dynamo_nvext", id="dynamo_nvext"), + param("smg_routing_key", id="smg_routing_key"), + param("session_id_header", id="session_id_header"), + ], +) # fmt: skip +def test_session_routing_plugins_resolve(name): + cls = plugins.get_class(PluginType.SESSION_ROUTING, name) + assert issubclass(cls, SessionRoutingBase) + + +def test_session_routing_enum_generated(): + from aiperf.plugin.enums import SessionRoutingType + + assert SessionRoutingType.DYNAMO_HEADERS == "dynamo_headers" + assert SessionRoutingType.SESSION_ID_HEADER == "session_id_header" diff --git a/tests/unit/timing/strategies/test_user_centric_rate.py b/tests/unit/timing/strategies/test_user_centric_rate.py index e0b86d3eb2..ae484a4e74 100644 --- a/tests/unit/timing/strategies/test_user_centric_rate.py +++ b/tests/unit/timing/strategies/test_user_centric_rate.py @@ -3,8 +3,11 @@ from unittest.mock import MagicMock import pytest +from pytest import param from aiperf.common.enums import CreditPhase +from aiperf.common.models import TurnMetadata +from aiperf.credit.structs import Credit, TurnToSend from aiperf.plugin.enums import TimingMode from aiperf.timing.config import CreditPhaseConfig from aiperf.timing.strategies.user_centric_rate import User, UserCentricStrategy @@ -169,6 +172,75 @@ async def test_kv_cache_benchmark(self, create_orchestrator_harness) -> None: assert len(h.sent_credits) >= 100 +@pytest.mark.asyncio +class TestContinuationBranchPlumb: + """Regression: the continuation turn built in ``handle_credit_return`` must + carry the next turn's branch/fork facts. The strategy has to pass the next + turn's metadata into ``TurnToSend.from_previous_credit`` -- dropping it + silently zeroes ``has_branches`` (and ``has_forks``), so a branching turn + would be wrongly treated as a leaf by downstream finality stamping. + """ + + def _make_strategy( + self, next_meta: TurnMetadata + ) -> tuple[UserCentricStrategy, MagicMock]: + cfg = CreditPhaseConfig( + phase=CreditPhase.PROFILING, + timing_mode=TimingMode.USER_CENTRIC_RATE, + request_rate=10.0, + num_users=5, + total_expected_requests=10, + ) + conversation_source = MagicMock() + conversation_source.get_next_turn_metadata.return_value = next_meta + credit_issuer = MagicMock() + strategy = UserCentricStrategy( + config=cfg, + conversation_source=conversation_source, + scheduler=MagicMock(), + stop_checker=MagicMock(), + credit_issuer=credit_issuer, + lifecycle=MagicMock(), + ) + sampled = MagicMock() + sampled.x_correlation_id = "sess-1" + strategy._session_to_user["sess-1"] = User(user_id=1, sampled=sampled) + return strategy, credit_issuer + + @pytest.mark.parametrize( + "branch_ids, expected", + [ + param(["b1"], True, id="declares-branch"), + param([], False, id="no-branch"), + ], + ) # fmt: skip + async def test_continuation_turn_carries_has_branches( + self, branch_ids, expected + ) -> None: + next_meta = TurnMetadata(branch_ids=branch_ids) + strategy, credit_issuer = self._make_strategy(next_meta) + credit = Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="conv-1", + x_correlation_id="sess-1", + turn_index=0, + num_turns=3, + issued_at_ns=0, + ) + + await strategy.handle_credit_return(credit) + + strategy._conversation_source.get_next_turn_metadata.assert_called_once_with( + credit + ) + credit_issuer.issue_credit.assert_called_once() + turn = credit_issuer.issue_credit.call_args.args[0] + assert isinstance(turn, TurnToSend) + assert turn.turn_index == 1 + assert turn.has_branches is expected + + class TestUserClass: def test_x_correlation_id_delegates_to_sampled(self) -> None: """User.x_correlation_id should delegate to sampled session.""" diff --git a/tests/unit/timing/test_session_tree_finality.py b/tests/unit/timing/test_session_tree_finality.py new file mode 100644 index 0000000000..b2f25844df --- /dev/null +++ b/tests/unit/timing/test_session_tree_finality.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Finality-query unit tests for the trimmed main ``SessionTreeRegistry``. + +The main registry is finality bookkeeping only: it holds no concurrency +manager and releases no slot (unlike the agentic-replay variant it was ported +from). These tests exercise the state map + the two finality queries directly. +""" + +from aiperf.common.enums import CreditPhase +from aiperf.timing.session_tree import SessionTreeRegistry + + +def _make_registry() -> SessionTreeRegistry: + return SessionTreeRegistry() + + +def _registry_with_tree( + root: str = "root-1", descendants: int = 0 +) -> SessionTreeRegistry: + registry = _make_registry() + registry.open_tree(root, phase=CreditPhase.PROFILING, root_pending=True) + if descendants: + registry.register_descendants(root, n=descendants) + return registry + + +def test_root_terminal_unknown_tree_is_none(): + assert _make_registry().root_terminal("nope") is None + + +def test_root_terminal_false_while_pending_true_after(): + # A live descendant keeps the tree open past on_root_terminal so the + # post-terminal state is observable; a drained tree would be retired + # (popped) and root_terminal would then read as unknown/None. + registry = _registry_with_tree(descendants=1) + assert registry.root_terminal("root-1") is False + registry.on_root_terminal("root-1") + assert registry.root_terminal("root-1") is True + + +def test_last_tree_request_root_with_no_descendants(): + registry = _registry_with_tree() + assert registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=True, has_branches=False + ) + + +def test_not_last_when_descendants_outstanding_or_branches_pending(): + registry = _registry_with_tree(descendants=1) + assert not registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=True, has_branches=False + ) + registry_no_desc = _registry_with_tree(root="root-2") + assert not registry_no_desc.is_last_tree_request( + "root-2", is_final_turn=True, is_root_credit=True, has_branches=True + ) + assert not registry_no_desc.is_last_tree_request( + "root-2", is_final_turn=False, is_root_credit=True, has_branches=False + ) + + +def test_last_tree_request_final_child_after_root_done(): + registry = _registry_with_tree(descendants=1) + registry.on_root_terminal("root-1") + assert registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=False, has_branches=False + ) + + +def test_unknown_tree_is_conservative_false(): + assert not _make_registry().is_last_tree_request( + "nope", is_final_turn=True, is_root_credit=True, has_branches=False + ) + + +def test_descendant_done_retires_drained_tree(): + # Sole child completes after root terminal -> tree drains and is retired. + registry = _registry_with_tree(descendants=1) + registry.on_root_terminal("root-1") + assert registry.has_tree("root-1") + retired = registry.on_descendant_done("root-1") + assert retired is True + assert not registry.has_tree("root-1") + assert registry.root_terminal("root-1") is None + + +def test_release_all_retires_open_trees_without_slot_release(): + registry = _make_registry() + registry.open_tree("a", CreditPhase.PROFILING, root_pending=True) + registry.open_tree("b", CreditPhase.PROFILING, root_pending=True) + assert registry.open_count() == 2 + assert registry.release_all() == 2 + assert registry.open_count() == 0 + + +def test_final_turn_spawn_resurrects_retired_tree_for_grandchild_finality(): + """Regression: root done -> last-outstanding child C's final turn declares a + SPAWN grandchild. The callback order retires the tree on C's own + on_descendant_done (step 4b) BEFORE the return-intercept registers C's + grandchildren (step 5); register_descendants must RESURRECT the retired tree + so the grandchild's genuinely-last credit can still stamp is_tree_final=True. + """ + registry = _registry_with_tree(descendants=1) # root + one live child C + registry.on_root_terminal("root-1") # root's terminal turn returns; C still live + + # C is the last outstanding descendant; its final-turn return decrements it, + # draining and retiring the tree (step 4b, before C's spawn intercept). + assert registry.on_descendant_done("root-1") is True + assert not registry.has_tree("root-1") + + # Step 5: C's final-turn SPAWN registers the grandchild AFTER that retire. + # Old behavior buffered it into a retired root nothing drains; now it + # resurrects the tree root-terminal with the grandchild outstanding. + registry.register_descendants("root-1", n=1) + assert registry.has_tree("root-1") + assert registry.root_terminal("root-1") is True + + # The grandchild's genuinely-last (non-branching final) credit CAN now be + # stamped tree-final -- root terminal, sole outstanding descendant. + assert registry.is_last_tree_request( + "root-1", is_final_turn=True, is_root_credit=False, has_branches=False + ) + + # The grandchild finishing re-drains and re-retires the tree coherently. + assert registry.on_descendant_done("root-1") is True + assert not registry.has_tree("root-1") + assert registry.late_events == 0 + + +def test_release_all_drains_pending_descendants(): + """release_all must clear the pre-open descendant buffer, not just _trees.""" + registry = _make_registry() + # No open_tree first: register_descendants buffers into _pending_descendants + # (the defensive pre-open path). + registry.register_descendants("orphan-root", n=2) + assert registry._pending_descendants # buffered, not yet folded into a tree + + registry.release_all() + + assert registry._pending_descendants == {} + assert registry._retired_roots == {} + + +def test_finality_flows_credit_to_request_info(): + """Schema guard: the three lineage-finality fields exist on BOTH the Credit + struct and the RequestInfo model, so the worker has fields to copy between. + + This asserts field-NAME presence only -- it does NOT verify a value is + actually copied (deleting the plumb kwargs in ``worker._create_request_info`` + keeps this green). The value-level plumb guard is + ``test_worker.py::test_create_request_info_plumbs_finality_from_credit``, + which stamps a REAL Credit and asserts the values surface on the RequestInfo. + """ + from aiperf.common.models.record_models import RequestInfo + from aiperf.credit.structs import Credit + + credit_fields = {"is_parent_final", "is_tree_final", "root_correlation_id"} + assert credit_fields <= set(Credit.__struct_fields__) + assert credit_fields <= set(RequestInfo.model_fields) diff --git a/tests/unit/timing/test_session_tree_wiring.py b/tests/unit/timing/test_session_tree_wiring.py new file mode 100644 index 0000000000..9eb12cd990 --- /dev/null +++ b/tests/unit/timing/test_session_tree_wiring.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end wiring test at the DAG orchestrator seam. + +Drives the REAL CreditIssuer + REAL BranchOrchestrator + REAL +SessionTreeRegistry + REAL ConversationSource through the full tree lifecycle +(open -> register -> root terminal -> descendant done) and asserts the +transitions produce correct lineage finality on the ISSUED credits. No +MagicMock for the registry or any credit -- only the phase scalars +(progress/lifecycle/stop_checker) and the concurrency slots are faked. +""" + +import time +from unittest.mock import MagicMock + +import pytest + +from aiperf.common.enums import ConversationBranchMode, CreditPhase +from aiperf.common.models import ConversationMetadata, DatasetMetadata, TurnMetadata +from aiperf.common.models.branch import ConversationBranchInfo +from aiperf.credit.issuer import CreditIssuer +from aiperf.credit.structs import Credit, TurnToSend +from aiperf.plugin import plugins +from aiperf.plugin.enums import DatasetSamplingStrategy, PluginType +from aiperf.timing.branch_orchestrator import BranchOrchestrator +from aiperf.timing.conversation_source import ConversationSource +from aiperf.timing.session_tree import SessionTreeRegistry + + +class _FakeConcurrency: + """Slots always granted; releases are no-ops.""" + + async def acquire_session_slot(self, phase, can_proceed) -> bool: + return True + + async def acquire_prefill_slot(self, phase, can_proceed) -> bool: + return True + + def release_session_slot(self, phase) -> None: + pass + + +class _CapturingRouter: + def __init__(self) -> None: + self.sent: list[Credit] = [] + + async def send_credit(self, *, credit: Credit) -> None: + self.sent.append(credit) + + +def _mk_source() -> ConversationSource: + """Root (2 turns, SPAWN branch on turn 0) -> one child (2 turns).""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[ + TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:0"]), + TurnMetadata(timestamp_ms=0.0), + ], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="child-conv", + turns=[TurnMetadata(timestamp_ms=0.0), TurnMetadata(timestamp_ms=0.0)], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +def _make_issuer( + registry: SessionTreeRegistry, router: _CapturingRouter +) -> CreditIssuer: + progress = MagicMock() + progress.increment_sent = MagicMock(return_value=(1, False)) + progress.freeze_sent_counts = MagicMock() + progress.all_credits_sent_event = MagicMock() + + stop_checker = MagicMock() + stop_checker.can_send_any_turn = MagicMock(return_value=True) + stop_checker.can_start_new_session = MagicMock(return_value=True) + stop_checker.can_send_dag_child_turn = MagicMock(return_value=True) + + cancellation = MagicMock() + cancellation.next_cancellation_delay_ns = MagicMock(return_value=None) + + lifecycle = MagicMock() + lifecycle.started_at_ns = time.time_ns() + lifecycle.started_at_perf_ns = time.perf_counter_ns() + + return CreditIssuer( + phase=CreditPhase.PROFILING, + stop_checker=stop_checker, + progress=progress, + concurrency_manager=_FakeConcurrency(), + credit_router=router, + cancellation_policy=cancellation, + lifecycle=lifecycle, + session_tree_registry=registry, + ) + + +@pytest.mark.asyncio +async def test_orchestrator_seam_drives_tree_finality_transitions(): + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + # 1. OPEN: issuing the root's turn-0 credit opens the tree (issuer seam). + root_turn0 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-x", + turn_index=0, + num_turns=2, + ) + await issuer.issue_credit(root_turn0) + assert registry.has_tree("root-x") + root_credit0 = router.sent[0] + assert root_credit0.agent_depth == 0 + # Turn 0 of 2 is not final -> conservative False. + assert root_credit0.is_tree_final is False + + # 2. REGISTER: intercepting the root turn-0 return spawns the child and + # registers it as a descendant (orchestrator seam). + await orch.intercept(root_credit0) + child_credits = [c for c in router.sent if c.agent_depth == 1] + assert len(child_credits) == 1 + child_credit0 = child_credits[0] + assert child_credit0.root_correlation_id == "root-x" + # Root still pending + descendant outstanding -> child not tree-final, and + # parent (== root) not yet final. + assert child_credit0.is_parent_final is False + assert child_credit0.is_tree_final is False + assert registry.open_count() == 1 + child_x = child_credit0.x_correlation_id + + # 3. ROOT TERMINAL: issue the root's final turn, then intercept it so the + # orchestrator clears root_pending (orchestrator seam via intercept). + root_turn1 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-x", + turn_index=1, + num_turns=2, + ) + await issuer.issue_credit(root_turn1) + root_credit1 = router.sent[-1] + # Child still outstanding when the root's final turn is issued. + assert root_credit1.is_tree_final is False + await orch.intercept(root_credit1) + assert registry.root_terminal("root-x") is True + + # 4. FINALITY ON ISSUE: the child's final continuation, issued after the + # root terminal with the child as the sole outstanding descendant, is + # provably the tree's last request AND its parent (the root) is final. + child_turn1 = TurnToSend( + conversation_id="child-conv", + x_correlation_id=child_x, + turn_index=1, + num_turns=2, + agent_depth=1, + parent_correlation_id="root-x", + root_correlation_id="root-x", + ) + await issuer.dispatch_child_turn(child_turn1) + child_credit1 = router.sent[-1] + assert child_credit1.x_correlation_id == child_x + assert child_credit1.is_parent_final is True + assert child_credit1.is_tree_final is True + + # 5. DESCENDANT DONE: the child's terminal return drains + retires the tree + # (orchestrator seam). + await orch.on_child_leaf_reached(child_x) + assert not registry.has_tree("root-x") + assert registry.open_count() == 0 + + +# ============================================================================= +# Conservative-contract regressions: SPAWN branches must gate finality +# ============================================================================= + + +def _mk_single_turn_spawn_source() -> ConversationSource: + """Single-turn root whose ONLY turn declares a terminal SPAWN branch.""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:0"])], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="child-conv", + turns=[TurnMetadata(timestamp_ms=0.0)], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +@pytest.mark.asyncio +async def test_single_turn_root_with_terminal_spawn_branch_not_tree_final(): + """(a) A single-turn root declaring a terminal SPAWN branch must stamp + is_tree_final=False on the root credit: its children spawn at + return-intercept, AFTER issue-time stamping, so the registry shows nothing + outstanding yet. Previously stamped a wrong True (has_forks is FORK-only).""" + from aiperf.timing.conversation_source import SampledSession + + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_single_turn_spawn_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + session = SampledSession( + conversation_id="root-conv", + metadata=source.get_metadata("root-conv"), + x_correlation_id="root-a", + ) + root_turn0 = session.build_first_turn() + # End-to-end stamp: SPAWN-only branch -> has_branches True, has_forks False. + assert root_turn0.has_branches is True + assert root_turn0.has_forks is False + + await issuer.issue_credit(root_turn0) + root_credit = router.sent[0] + assert root_credit.is_final_turn is True + assert root_credit.is_tree_final is False + + # The spawn then registers + the tree drains cleanly through the child. + await orch.intercept(root_credit) + child_credits = [c for c in router.sent if c.agent_depth == 1] + assert len(child_credits) == 1 + await orch.on_child_leaf_reached(child_credits[0].x_correlation_id) + assert not registry.has_tree("root-a") + assert registry.late_events == 0 + + +def _mk_grandchild_spawn_source() -> ConversationSource: + """Single-turn root -> child (2 turns; final turn declares a SPAWN + branch to a grandchild).""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:0"])], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="child-conv", + turns=[ + TurnMetadata(timestamp_ms=0.0), + TurnMetadata(timestamp_ms=0.0, branch_ids=["child-conv:1"]), + ], + branches=[ + ConversationBranchInfo( + branch_id="child-conv:1", + child_conversation_ids=["grandchild-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="post", + ), + ], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ConversationMetadata( + conversation_id="grandchild-conv", + turns=[TurnMetadata(timestamp_ms=0.0)], + agent_depth=2, + parent_conversation_id="child-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +@pytest.mark.asyncio +async def test_child_final_turn_spawning_grandchild_not_tree_final(): + """(b) A child's final turn that declares a SPAWN branch (grandchild + pending) must stamp is_tree_final=False even when the child is the sole + outstanding descendant and the root is already terminal -- previously the + exact wrong-True scenario.""" + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_grandchild_spawn_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + # Root turn 0 (single turn, spawning): opens tree; intercept spawns the + # child, registers it, then marks the root terminal. + root_turn0 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-b", + turn_index=0, + num_turns=1, + has_branches=True, + ) + await issuer.issue_credit(root_turn0) + root_credit = router.sent[0] + assert root_credit.is_tree_final is False + await orch.intercept(root_credit) + assert registry.root_terminal("root-b") is True + child_credit0 = next(c for c in router.sent if c.agent_depth == 1) + child_x = child_credit0.x_correlation_id + + # Child's final turn, built from real metadata (branch_ids on turn 1), + # issued while the child is the sole outstanding descendant. + next_meta = source.get_next_turn_metadata(child_credit0) + child_turn1 = TurnToSend.from_previous_credit(child_credit0, next_meta) + assert child_turn1.has_branches is True + assert child_turn1.has_forks is False + await issuer.dispatch_child_turn(child_turn1) + child_credit1 = router.sent[-1] + assert child_credit1.is_final_turn is True + assert child_credit1.is_parent_final is True + assert child_credit1.is_tree_final is False # grandchild pending + + # Production return order: leaf-reached fires BEFORE intercept spawns the + # grandchild (callback handler step 4b before step 5). The child is the last + # outstanding descendant, so its leaf-reached decrement drains and retires + # the tree; intercept's register_descendants then RESURRECTS it (root already + # terminal). The grandchild -- a single-turn, sole-remaining, genuinely-last + # request -- is therefore correctly stamped tree-final (previously the + # resurrect was missing and it was wrongly under-fired to False). + await orch.on_child_leaf_reached(child_x) + await orch.intercept(child_credit1) + grandchild_credits = [c for c in router.sent if c.agent_depth == 2] + assert len(grandchild_credits) == 1 + assert grandchild_credits[0].root_correlation_id == "root-b" + assert grandchild_credits[0].is_tree_final is True + await orch.on_child_leaf_reached(grandchild_credits[0].x_correlation_id) + assert not registry.has_tree("root-b") + assert registry.late_events == 0 + + +def _mk_pre_session_source() -> ConversationSource: + """Two-turn root with a pre-session SPAWN branch attached to turn 0.""" + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="root-conv", + turns=[ + TurnMetadata(timestamp_ms=0.0, branch_ids=["root-conv:pre"]), + TurnMetadata(timestamp_ms=0.0), + ], + branches=[ + ConversationBranchInfo( + branch_id="root-conv:pre", + child_conversation_ids=["pre-child-conv"], + mode=ConversationBranchMode.SPAWN, + dispatch_timing="pre", + ), + ], + agent_depth=0, + ), + ConversationMetadata( + conversation_id="pre-child-conv", + turns=[TurnMetadata(timestamp_ms=0.0)], + agent_depth=1, + parent_conversation_id="root-conv", + is_root=False, + ), + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + SamplerClass = plugins.get_class(PluginType.DATASET_SAMPLER, ds.sampling_strategy) + sampler = SamplerClass( + conversation_ids=[c.conversation_id for c in ds.conversations] + ) + return ConversationSource(ds, sampler) + + +@pytest.mark.asyncio +async def test_pre_session_child_live_blocks_root_final_turn_tree_final(): + """(c) A live pre-session SPAWN child must keep the root's final turn from + stamping is_tree_final=True. Pre-dispatch runs before sampling (no root id + exists yet), so the orchestrator folds the live pre children into the root + instance's tree at its turn-0 return -- before any final-turn issue.""" + registry = SessionTreeRegistry() + router = _CapturingRouter() + source = _mk_pre_session_source() + issuer = _make_issuer(registry, router) + orch = BranchOrchestrator( + conversation_source=source, + credit_issuer=issuer, + session_tree_registry=registry, + ) + + # Pre-dispatch fires the background child before any root credit exists. + await orch.dispatch_pre_session_branches() + pre_credits = [c for c in router.sent if c.agent_depth == 1] + assert len(pre_credits) == 1 + pre_child_x = pre_credits[0].x_correlation_id + assert pre_credits[0].is_tree_final is False + + # Root turn 0: opens the tree; its return-intercept folds the live + # pre-session child into this root's tree. + root_turn0 = TurnToSend( + conversation_id="root-conv", + x_correlation_id="root-c", + turn_index=0, + num_turns=2, + has_branches=True, # turn 0 declares the pre branch + ) + await issuer.issue_credit(root_turn0) + root_credit0 = [c for c in router.sent if c.agent_depth == 0][0] + await orch.intercept(root_credit0) + + # Root's FINAL turn (declares no branches) with the pre child still live: + # must stamp False. Previously wrong-True (pre children were never + # registered anywhere). + next_meta = source.get_next_turn_metadata(root_credit0) + root_turn1 = TurnToSend.from_previous_credit(root_credit0, next_meta) + assert root_turn1.has_branches is False + await issuer.issue_credit(root_turn1) + root_credit1 = router.sent[-1] + assert root_credit1.is_final_turn is True + assert root_credit1.is_tree_final is False # pre-session child still live + + # Root terminal, then the pre child's terminal drains + retires the tree. + await orch.intercept(root_credit1) + assert registry.root_terminal("root-c") is True + assert registry.has_tree("root-c") + await orch.on_child_leaf_reached(pre_child_x) + assert not registry.has_tree("root-c") + assert registry.late_events == 0 diff --git a/tests/unit/workers/test_inference_client.py b/tests/unit/workers/test_inference_client.py index cee9cb695e..7a78e14535 100644 --- a/tests/unit/workers/test_inference_client.py +++ b/tests/unit/workers/test_inference_client.py @@ -20,6 +20,19 @@ from aiperf.common.redact import REDACTED_VALUE from aiperf.plugin.enums import EndpointType, TransportType from aiperf.workers.inference_client import InferenceClient, detect_transport_from_url +from aiperf.workers.session_routing import ( + DynamoHeadersRouting, + DynamoNvextRouting, + SessionIdHeaderRouting, + SmgRoutingKeyRouting, +) + +_ROUTING_CLASSES = { + "dynamo_headers": DynamoHeadersRouting, + "dynamo_nvext": DynamoNvextRouting, + "smg_routing_key": SmgRoutingKeyRouting, + "session_id_header": SessionIdHeaderRouting, +} @pytest.fixture @@ -435,3 +448,298 @@ def mock_get_class(protocol, name): assert not pydantic_warnings, ( f"Unexpected Pydantic serialization warnings for {base_url!r}: {pydantic_warnings}" ) + + +class TestSessionRouting: + """Session-routing plugins wired through the InferenceClient chokepoint. + + The endpoint/transport plugins are mocked as before; the session_routing + protocol resolves to the real routing classes so the chokepoint exercises + genuine header/body transforms and the notify_session_end pass-through. + """ + + def _build_client( + self, + mock_http_transport_entry, + *, + session_routing: str | None, + session_routing_opts: dict | None = None, + ) -> InferenceClient: + model_endpoint = ModelEndpointInfo( + models=ModelListInfo( + models=[ModelInfo(name="test-model")], + model_selection_strategy=ModelSelectionStrategy.ROUND_ROBIN, + ), + endpoint=EndpointInfo( + type=EndpointType.CHAT, + base_url="http://localhost:8000/v1/test", + session_routing=session_routing, + session_routing_opts=session_routing_opts or {}, + ), + ) + mock_transport = MagicMock() + mock_endpoint = MagicMock() + mock_endpoint.get_endpoint_headers.return_value = {} + mock_endpoint.get_endpoint_params.return_value = {} + mock_endpoint.format_payload.return_value = { + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + } + + def mock_get_class(protocol, name): + if protocol == "endpoint": + return lambda **kwargs: mock_endpoint + if protocol == "transport": + return lambda **kwargs: mock_transport + if protocol == "session_routing": + return _ROUTING_CLASSES[name] + raise ValueError(f"Unknown protocol: {protocol}") + + with ( + patch( + "aiperf.workers.inference_client.plugins.get_class", + side_effect=mock_get_class, + ), + patch( + "aiperf.workers.inference_client.plugins.list_entries", + return_value=[mock_http_transport_entry], + ), + ): + client = InferenceClient( + model_endpoint=model_endpoint, service_id="test-service-id" + ) + client.transport.send_request = AsyncMock(return_value=RequestRecord()) + return client + + def _request_info( + self, + client: InferenceClient, + *, + x_correlation_id: str = "corr-1", + parent_correlation_id: str | None = None, + is_final_turn: bool = False, + raw_payload: dict | None = None, + ) -> RequestInfo: + turn = ( + Turn(role="user", raw_payload=raw_payload) + if raw_payload is not None + else Turn(role="user", texts=[Text(contents=["hello"])]) + ) + return RequestInfo( + model_endpoint=client.model_endpoint, + turns=[turn], + turn_index=0, + credit_num=0, + credit_phase=CreditPhase.PROFILING, + x_request_id="req-1", + x_correlation_id=x_correlation_id, + parent_correlation_id=parent_correlation_id, + conversation_id="conv-template", + is_final_turn=is_final_turn, + ) + + def _sent_payload(self, client: InferenceClient) -> dict: + return client.transport.send_request.call_args.kwargs["payload"] + + @pytest.mark.asyncio + async def test_dynamo_headers_mode_emits_headers_and_leaves_body( + self, mock_http_transport_entry + ): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + request_info = self._request_info( + client, parent_correlation_id="parent-corr", is_final_turn=False + ) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Dynamo-Session-ID"] == "corr-1" + assert ( + request_info.endpoint_headers["X-Dynamo-Parent-Session-ID"] == "parent-corr" + ) + assert "nvext" not in self._sent_payload(client) + + @pytest.mark.asyncio + async def test_dynamo_headers_root_omits_parent_header( + self, mock_http_transport_entry + ): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + request_info = self._request_info(client, parent_correlation_id=None) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Dynamo-Session-ID"] == "corr-1" + assert "X-Dynamo-Parent-Session-ID" not in request_info.endpoint_headers + + @pytest.mark.asyncio + async def test_dynamo_nvext_mode_binds_then_closes(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": "123"}, + ) + + non_final = self._request_info(client, is_final_turn=False) + await client._send_request_to_transport(non_final) + assert self._sent_payload(client)["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "bind", + "timeout": 123, + } + + final = self._request_info(client, is_final_turn=True) + await client._send_request_to_transport(final) + assert self._sent_payload(client)["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "close", + } + + @pytest.mark.asyncio + async def test_dynamo_nvext_on_raw_payload_turn_does_not_mutate_cache( + self, mock_http_transport_entry + ): + """The nvext transform runs on a cached raw_payload dict; the copy-on-write + contract must leave the shared Turn.raw_payload untouched.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" + ) + raw_payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + request_info = self._request_info(client, raw_payload=raw_payload) + + await client._send_request_to_transport(request_info) + + # The wire payload carries the injected session_control ... + assert ( + self._sent_payload(client)["nvext"]["session_control"]["action"] == "bind" + ) + # ... but the cached raw_payload dict is unchanged (no nvext leaked in). + assert "nvext" not in raw_payload + assert request_info.turns[0].raw_payload == { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + } + + @pytest.mark.asyncio + async def test_smg_routing_key_mode_emits_header(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, session_routing="smg_routing_key" + ) + request_info = self._request_info(client) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-SMG-Routing-Key"] == "corr-1" + assert "nvext" not in self._sent_payload(client) + + @pytest.mark.asyncio + async def test_session_id_header_custom_name(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, + session_routing="session_id_header", + session_routing_opts={"header_name": "X-Affinity"}, + ) + request_info = self._request_info(client) + + await client._send_request_to_transport(request_info) + + assert request_info.endpoint_headers["X-Affinity"] == "corr-1" + assert "nvext" not in self._sent_payload(client) + + @pytest.mark.asyncio + async def test_routing_unset_no_headers_no_body_change( + self, mock_http_transport_entry + ): + client = self._build_client(mock_http_transport_entry, session_routing=None) + assert client._routing is None + request_info = self._request_info(client, parent_correlation_id="parent-corr") + + await client._send_request_to_transport(request_info) + + payload = self._sent_payload(client) + assert "nvext" not in payload + assert "X-Dynamo-Session-ID" not in request_info.endpoint_headers + assert "X-Dynamo-Parent-Session-ID" not in request_info.endpoint_headers + + @pytest.mark.asyncio + async def test_notify_session_end_reaches_plugin(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + client._routing.on_session_end = MagicMock() + + # Pass-through must not dedupe: idempotency is the plugin's job. + client.notify_session_end("corr-1") + client.notify_session_end("corr-1") + + assert client._routing.on_session_end.call_count == 2 + client._routing.on_session_end.assert_called_with("corr-1") + + def test_notify_session_end_noop_when_routing_unset( + self, mock_http_transport_entry + ): + client = self._build_client(mock_http_transport_entry, session_routing=None) + # No routing plugin: the hook is a safe no-op (never raises). + client.notify_session_end("corr-1") + + def test_notify_session_end_swallows_plugin_error_and_warns( + self, mock_http_transport_entry + ): + """A raising on_session_end must NOT propagate (core eviction must + proceed); the failure is logged with the plugin + session named.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + client._routing.on_session_end = MagicMock(side_effect=RuntimeError("boom")) + + with patch.object(client, "warning") as warn: + # Must not raise. + client.notify_session_end("corr-err") + + client._routing.on_session_end.assert_called_once_with("corr-err") + warn.assert_called_once() + msg = warn.call_args.args[0] + assert "dynamo_headers" in msg and "corr-err" in msg + + @pytest.mark.asyncio + async def test_raising_headers_produces_plugin_attributed_error_record( + self, mock_http_transport_entry + ): + """A plugin exception in headers() surfaces as an error record whose + message names the routing plugin, not the inference server.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + client._routing.headers = MagicMock( + side_effect=RuntimeError("bad header build") + ) + request_info = self._request_info(client) + + record = await client._send_request_internal(request_info) + + assert record.error is not None + assert "dynamo_headers" in record.error.message + assert "headers()" in record.error.message + # The transport was never reached (the fault is pre-send). + client.transport.send_request.assert_not_called() + + @pytest.mark.asyncio + async def test_raising_transform_body_produces_plugin_attributed_error_record( + self, mock_http_transport_entry + ): + """A plugin exception in transform_body() is attributed to the plugin.""" + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" + ) + client._routing.transform_body = MagicMock( + side_effect=RuntimeError("bad body transform") + ) + request_info = self._request_info(client) + + record = await client._send_request_internal(request_info) + + assert record.error is not None + assert "dynamo_nvext" in record.error.message + assert "transform_body()" in record.error.message diff --git a/tests/unit/workers/test_session_routing.py b/tests/unit/workers/test_session_routing.py new file mode 100644 index 0000000000..94045640e5 --- /dev/null +++ b/tests/unit/workers/test_session_routing.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pydantic import ValidationError + +from aiperf.workers.session_routing import ( + DynamoHeadersRouting, + DynamoNvextOptions, + DynamoNvextRouting, + RoutingContext, + SessionIdHeaderOptions, + SessionIdHeaderRouting, + SessionRoutingBase, + SmgRoutingKeyRouting, +) + + +def _ctx(**overrides) -> RoutingContext: + defaults = dict( + x_correlation_id="corr-1", + parent_correlation_id=None, + root_correlation_id="corr-1", + is_final_turn=False, + is_parent_final=None, + is_tree_final=False, + ) + defaults.update(overrides) + return RoutingContext(**defaults) + + +class TestDynamoHeadersRouting: + def test_root_emits_session_header_only(self): + plugin = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + assert plugin.headers(_ctx()) == {"X-Dynamo-Session-ID": "corr-1"} + assert plugin.mutates_body is False + + def test_child_emits_parent_header(self): + plugin = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + headers = plugin.headers(_ctx(parent_correlation_id="parent-1")) + assert headers == { + "X-Dynamo-Session-ID": "corr-1", + "X-Dynamo-Parent-Session-ID": "parent-1", + } + + def test_body_untouched(self): + plugin = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + payload = {"messages": []} + assert plugin.transform_body(payload, _ctx()) is payload + + +class TestDynamoNvextRouting: + def test_non_final_turn_binds_with_timeout(self): + plugin = DynamoNvextRouting(DynamoNvextOptions(timeout_seconds=123)) + merged = plugin.transform_body({"messages": []}, _ctx()) + assert merged["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "bind", + "timeout": 123, + } + assert plugin.mutates_body is True + + def test_final_turn_closes_without_timeout(self): + plugin = DynamoNvextRouting(DynamoNvextOptions()) + merged = plugin.transform_body({}, _ctx(is_final_turn=True)) + assert merged["nvext"]["session_control"] == { + "session_id": "corr-1", + "action": "close", + } + + def test_never_mutates_input_payload(self): + nested_sc = {"existing": "keep"} + nvext = {"trace": "keep", "session_control": nested_sc} + payload = {"nvext": nvext} + plugin = DynamoNvextRouting(DynamoNvextOptions()) + merged = plugin.transform_body(payload, _ctx()) + assert payload == { + "nvext": {"trace": "keep", "session_control": {"existing": "keep"}} + } + assert nvext == {"trace": "keep", "session_control": {"existing": "keep"}} + assert merged is not payload + assert merged["nvext"]["session_control"]["existing"] == "keep" + + def test_options_default_and_bounds(self): + assert DynamoNvextOptions().timeout_seconds == 300 + with pytest.raises(ValidationError): + DynamoNvextOptions(timeout_seconds=0) + + def test_options_reject_unknown_keys(self): + with pytest.raises(ValidationError): + DynamoNvextOptions(timeout_secs=5) + + def test_typed_options_access(self): + plugin = DynamoNvextRouting(DynamoNvextOptions(timeout_seconds=42)) + assert plugin.options.timeout_seconds == 42 + + +class TestSmgRoutingKeyRouting: + def test_emits_routing_key(self): + plugin = SmgRoutingKeyRouting(SmgRoutingKeyRouting.Options()) + assert plugin.headers(_ctx()) == {"X-SMG-Routing-Key": "corr-1"} + + def test_rejects_any_opt(self): + with pytest.raises(ValidationError): + SmgRoutingKeyRouting.Options(anything="x") + + +class TestSessionIdHeaderRouting: + def test_default_header_name(self): + plugin = SessionIdHeaderRouting(SessionIdHeaderOptions()) + assert plugin.headers(_ctx()) == {"X-Session-ID": "corr-1"} + + def test_custom_header_name(self): + plugin = SessionIdHeaderRouting( + SessionIdHeaderOptions(header_name="X-Affinity") + ) + assert plugin.headers(_ctx()) == {"X-Affinity": "corr-1"} + + +class TestBaseDefaults: + def test_on_session_end_default_noop_and_idempotent(self): + class Passthrough(SessionRoutingBase): + pass + + plugin = Passthrough(Passthrough.Options()) + plugin.on_session_end("corr-1") + plugin.on_session_end("corr-1") + + def test_stateful_open_once_lifecycle_expressible(self): + """The legacy-nvext shape: open-once instance state, bounded by on_session_end.""" + + class OpenOnce(SessionRoutingBase): + mutates_body = True + + def __init__(self, options): + super().__init__(options) + self.opened: set[str] = set() + + def transform_body(self, payload, ctx): + merged = dict(payload) + if ctx.x_correlation_id not in self.opened: + self.opened.add(ctx.x_correlation_id) + merged["action"] = "open" + return merged + + def on_session_end(self, x_correlation_id): + self.opened.discard(x_correlation_id) + + plugin = OpenOnce(OpenOnce.Options()) + assert plugin.transform_body({}, _ctx())["action"] == "open" + assert "action" not in plugin.transform_body({}, _ctx()) + plugin.on_session_end("corr-1") + plugin.on_session_end("corr-1") # idempotent + assert plugin.opened == set() diff --git a/tests/unit/workers/test_worker.py b/tests/unit/workers/test_worker.py index 244443d2f5..1afde57a65 100644 --- a/tests/unit/workers/test_worker.py +++ b/tests/unit/workers/test_worker.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import asyncio -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, MagicMock, Mock import pytest @@ -288,3 +288,184 @@ async def test_falls_back_to_dataset_manager_when_no_client_and_not_stopping( assert result == expected_conversation mock_fallback.assert_called_once_with("test-conv-123", sample_credit_context) + + +# --- Session-Routing Terminal Hook Tests --- + + +@pytest.mark.asyncio +class TestReleaseAndEvictForTerminal: + """The terminal-eviction path fires the routing plugin's post-session hook + (``InferenceClient.notify_session_end``) so stateful plugins release the + session on ANY terminal outcome, even one abandoned before its final turn + (e.g. cancellation).""" + + async def test_release_and_evict_for_terminal_notifies_session_end( + self, mock_worker, sample_credit_context + ): + credit = sample_credit_context.credit + mock_worker.inference_client.notify_session_end = Mock() + + mock_worker._release_and_evict_for_terminal(credit, credit.x_correlation_id) + + mock_worker.inference_client.notify_session_end.assert_called_once_with( + credit.x_correlation_id + ) + + async def test_raising_on_session_end_does_not_block_eviction( + self, mock_worker, sample_credit_context + ): + """A routing plugin whose on_session_end raises must not abort the + worker's eviction: the InferenceClient hook swallows + warns, so the + session still gets evicted.""" + from aiperf.workers.inference_client import InferenceClient + + credit = sample_credit_context.credit + fake_client = MagicMock() + fake_client._routing.on_session_end.side_effect = RuntimeError("plugin boom") + fake_client._routing_mode = "dynamo_headers" + # Drive the REAL hook logic (try/except-log) bound to the fake client. + fake_client.notify_session_end = InferenceClient.notify_session_end.__get__( + fake_client + ) + mock_worker.inference_client = fake_client + mock_worker.session_manager.get = Mock(return_value=None) + mock_worker.session_manager.evict = Mock() + + # Must not raise despite the plugin fault. + mock_worker._release_and_evict_for_terminal(credit, credit.x_correlation_id) + + mock_worker.session_manager.evict.assert_called_once_with( + credit.x_correlation_id + ) + fake_client._routing.on_session_end.assert_called_once_with( + credit.x_correlation_id + ) + fake_client.warning.assert_called_once() + + +class TestSessionRoutingTerminalHooks: + """Every terminal disposition this codebase has reaches + ``InferenceClient.notify_session_end`` so routing plugins get their + idempotent post-session hook: the final-turn / cancellation eviction path + and the done-callback cancel-before-start branch (which the finally block + never sees).""" + + def _done_callback(self, *, returned: bool) -> MagicMock: + """Drive the real done-callback on a mock worker with a not-yet-returned + (or already-returned) credit context.""" + worker = MagicMock() + worker.inference_client = MagicMock() + worker.service_id = "worker-1" + credit = MagicMock() + credit.id = 7 + credit.x_correlation_id = "conv-done" + ctx = MagicMock() + ctx.returned = returned + ctx.credit = credit + ctx.error = None + task = MagicMock() + task.cancelled.return_value = True + Worker._on_credit_drop_message_task_done.__get__(worker)(task, ctx) + return worker + + def test_done_callback_not_returned_branch_notifies_session_end(self) -> None: + """The cancel-before-start path (task done, credit never returned) is a + terminal disposition the finally block never sees, so the done-callback + must fire the hook too.""" + worker = self._done_callback(returned=False) + worker.inference_client.notify_session_end.assert_called_once_with("conv-done") + + def test_done_callback_already_returned_does_not_double_notify(self) -> None: + """When the credit already returned, the finally block already notified; + the done-callback short-circuits without a second hook call.""" + worker = self._done_callback(returned=True) + worker.inference_client.notify_session_end.assert_not_called() + + +@pytest.mark.asyncio +class TestCreateRequestInfoFinality: + """Touch #3: worker -> RequestInfo lineage-finality plumb. + + The Credit is a REAL struct (not a MagicMock, which would auto-create the + attributes and mask a missed plumb) carrying both finality facts plus a + tree root id. + """ + + async def test_create_request_info_plumbs_finality_from_credit(self, mock_worker): + from aiperf.common.models import Turn + + credit_context = CreditContext( + credit=Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="test-conv", + x_correlation_id="child-x", + turn_index=0, + num_turns=1, + issued_at_ns=0, + agent_depth=1, + parent_correlation_id="root-x", + root_correlation_id="root-x", + is_parent_final=True, + is_tree_final=True, + ), + drop_perf_ns=0, + ) + + # Session fields the builder reads, set to real values (the guard is on + # the Credit, so a lightweight session stand-in is fine here). + session = MagicMock() + session.x_correlation_id = "child-x" + session.conversation.session_id = "test-conv" + session.turn_index = 0 + session.turn_list = [Turn()] + session.url_index = None + + request_info = mock_worker._create_request_info( + x_request_id="request-id", + session=session, + credit_context=credit_context, + ) + + assert request_info.is_parent_final is True + assert request_info.is_tree_final is True + # effective_root_correlation_id resolves to the stamped root id. + assert request_info.root_correlation_id == "root-x" + + async def test_create_request_info_root_defaults_to_own_correlation( + self, mock_worker + ): + """A root credit (no root_correlation_id) surfaces its own + x_correlation_id as the record's tree root, with conservative finality. + """ + from aiperf.common.models import Turn + + credit_context = CreditContext( + credit=Credit( + id=2, + phase=CreditPhase.PROFILING, + conversation_id="test-conv", + x_correlation_id="root-x", + turn_index=0, + num_turns=1, + issued_at_ns=0, + ), + drop_perf_ns=0, + ) + session = MagicMock() + session.x_correlation_id = "root-x" + session.conversation.session_id = "test-conv" + session.turn_index = 0 + session.turn_list = [Turn()] + session.url_index = None + + request_info = mock_worker._create_request_info( + x_request_id="request-id", + session=session, + credit_context=credit_context, + ) + + assert request_info.root_correlation_id == "root-x" + assert request_info.is_parent_final is None + assert request_info.is_tree_final is False