From 205b4dd9525ba2a50ca27a0ef19ea23cb7e834fe Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Wed, 8 Jul 2026 22:22:02 -0700 Subject: [PATCH 1/2] feat(routing): session-routing plugin layer for per-session router affinity One plugin category (session_routing) + one CLI flag (--session-routing , with repeatable --session-routing-opt key=value) covers every mechanism for telling an external router which session a request belongs to, so all of a conversation's turns re-land on the replica holding its KV prefix. "Session" (live instance keyed by x_correlation_id), not "conversation" (dataset template), is the affinity identity. Exactly one mode runs per invocation; the selected plugin is instantiated once per worker by InferenceClient and invoked at the request-serialization chokepoint. Built-in modes: - dynamo_headers (preset): X-Dynamo-Session-ID plus X-Dynamo-Parent-Session-ID on subagent children, for a Dynamo frontend running --router-session-affinity-ttl-secs. - dynamo_nvext: nvext.session_control request-body metadata -- bind (with timeout, option timeout_seconds, default 300) on every non-final turn, close on the final turn -- for Dynamo builds that implement session_control. Plugin values win over dataset-shipped session_control keys; other nvext content is preserved. - smg_routing_key (preset): X-SMG-Routing-Key for the SGLang Model Gateway manual policy. - session_id_header (preset): additive X-Session-ID, parameterless. - identity_headers: fully generic tiered identity headers emitting exactly what is configured, nothing by default. Options (each a comma-separable name list): session (this session's correlation ID), parent (immediate parent's ID, omitted for roots), root (session-tree root's ID -- whole-tree affinity, e.g. pinning an entire agent tree to one replica for prefix-cache locality). Requires at least one name across tiers; names are validated as RFC 9110 tokens at config time and must be case-insensitively unique across all tiers combined. Each header preset is documented with its identity_headers equivalent. Removed (superseded): --use-dynamo-conv-aware-routing / --use-dynamo-session-control, --use-legacy-dynamo-session-control, --dynamo-session-timeout-seconds, the dynamo_session_control module, AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID, AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID, and the base_transports env-toggle header blocks. Co-Authored-By: Claude Fable 5 Signed-off-by: Anthony Casagrande --- README.md | 2 +- docs/cli-options.md | 20 +- docs/environment-variables.md | 2 - docs/plugins/plugin-system.md | 55 +- src/aiperf/common/config/config_defaults.py | 3 - src/aiperf/common/config/endpoint_config.py | 143 +++-- src/aiperf/common/environment.py | 10 - .../common/models/model_endpoint_info.py | 31 +- src/aiperf/common/models/record_models.py | 11 + src/aiperf/credit/issuer.py | 32 ++ src/aiperf/credit/structs.py | 20 +- src/aiperf/dataset/dataset_manager.py | 147 +++-- .../exporters/console_api_error_exporter.py | 30 +- src/aiperf/plugin/categories.yaml | 16 + src/aiperf/plugin/enums.py | 6 +- src/aiperf/plugin/plugins.py | 7 +- src/aiperf/plugin/plugins.yaml | 53 ++ src/aiperf/plugin/schema/plugins.schema.json | 49 ++ src/aiperf/timing/branch_orchestrator.py | 11 + src/aiperf/timing/conversation_source.py | 4 + src/aiperf/timing/session_tree.py | 38 ++ .../timing/strategies/agentic_replay.py | 5 +- .../timing/strategies/fixed_schedule.py | 3 + src/aiperf/transports/base_transports.py | 5 - src/aiperf/workers/dynamo_session_control.py | 113 ---- src/aiperf/workers/inference_client.py | 170 ++++-- src/aiperf/workers/session_manager.py | 9 + src/aiperf/workers/session_routing.py | 327 +++++++++++ src/aiperf/workers/worker.py | 67 ++- .../test_dynamo_session_control_raw_export.py | 143 ----- .../test_session_routing_raw_export.py | 224 ++++++++ .../common/config/test_endpoint_config.py | 72 ++- .../unit/common/models/test_endpoint_info.py | 23 +- .../config/test_session_routing_config.py | 121 ++++ .../test_credit_issuer_join_adversarial.py | 30 + tests/unit/credit/test_issuer_finality.py | 297 ++++++++++ tests/unit/credit/test_structs.py | 21 + tests/unit/dataset/test_dataset_manager.py | 281 +++++++++- .../dataset/test_dataset_manager_cache.py | 4 +- ...dataset_manager_inputs_json_adversarial.py | 6 +- .../test_console_api_error_exporter.py | 42 +- .../plugin/test_session_routing_registry.py | 30 + .../unit/timing/test_session_tree_finality.py | 112 ++++ tests/unit/transports/test_base_transport.py | 23 - .../workers/test_dynamo_session_control.py | 164 ------ tests/unit/workers/test_inference_client.py | 518 ++++++++++++------ tests/unit/workers/test_session_routing.py | 302 ++++++++++ tests/unit/workers/test_worker.py | 229 +++++++- 48 files changed, 3166 insertions(+), 865 deletions(-) delete mode 100644 src/aiperf/workers/dynamo_session_control.py create mode 100644 src/aiperf/workers/session_routing.py delete mode 100644 tests/integration/test_dynamo_session_control_raw_export.py create mode 100644 tests/integration/test_session_routing_raw_export.py create mode 100644 tests/unit/config/test_session_routing_config.py create mode 100644 tests/unit/credit/test_issuer_finality.py create mode 100644 tests/unit/plugin/test_session_routing_registry.py create mode 100644 tests/unit/timing/test_session_tree_finality.py delete mode 100644 tests/unit/workers/test_dynamo_session_control.py create mode 100644 tests/unit/workers/test_session_routing.py diff --git a/README.md b/README.md index 76013db374..ea91418646 100644 --- a/README.md +++ b/README.md @@ -200,7 +200,7 @@ Log File: /home/user/aiperf/artifacts/granite4:350m-openai-chat-concurrency1/log | [CLI Options](docs/cli-options.md) | Complete command and option reference | | [Metrics Reference](docs/metrics-reference.md) | All metric definitions, formulas, and requirements | | [Environment Variables](docs/environment-variables.md) | All `AIPERF_*` configuration variables | -| [Plugin System](docs/plugins/plugin-system.md) | Plugin architecture, 27 categories, creation guide | +| [Plugin System](docs/plugins/plugin-system.md) | Plugin architecture, 31 categories, creation guide | | [Creating Plugins](docs/plugins/creating-your-first-plugin.md) | Step-by-step plugin tutorial | | [Accuracy Benchmarks](docs/accuracy/accuracy_stubs.md) | Accuracy evaluation stubs and datasets | | [Benchmark Modes](docs/benchmark-modes/trace-replay.md) | Trace replay and timing modes | diff --git a/docs/cli-options.md b/docs/cli-options.md index 940c821a06..1afeeed320 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -266,21 +266,15 @@ Use the legacy 'max_tokens' field instead of 'max_completion_tokens' in request Use server-reported token counts from API usage fields instead of client-side tokenization. When enabled, tokenizers are still loaded (needed for dataset generation) but tokenizer.encode() is not called for computing metrics. Token count fields will be None if the server does not provide usage information. For OpenAI-compatible streaming endpoints (chat/completions), stream_options.include_usage is automatically configured when this flag is enabled. Recommended whenever the AIPerf tokenizer can disagree with the server's tokenizer (e.g. unmatched tokenizer revision, vendor-specific BPE merges, or chat templates that differ from the server) — this most often shows up as an output sequence length (OSL) mismatch even when the server is honoring the request (e.g. with ignore_eos=true).
_Flag (no value required)_ -#### `--use-dynamo-conv-aware-routing`, `--use-dynamo-session-control` +#### `--session-routing` `` -Emit Dynamo nvext.session_control in OpenAI-compatible request bodies so Dynamo can bind all turns from the same replayed conversation lineage to the same backend worker. This is only intended for Dynamo frontends that implement session_control. -
_Flag (no value required)_ - -#### `--use-legacy-dynamo-session-control` - -Emit the legacy Dynamo nvext.session_control lifecycle that released Dynamo (v1.2.x) understands: action 'open' on the first turn, session_id only on intermediate turns, and action 'close' on the final turn. Use this when the target Dynamo predates the 'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected with an HTTP 400. Requires --use-dynamo-conv-aware-routing, and the Dynamo deployment must expose a worker session_control endpoint for 'open' to take effect. -
_Flag (no value required)_ +Session-aware routing mode: stamps per-session identity on every request for router affinity. Presets: dynamo_headers (X-Dynamo-Session-ID + parent header), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (additive X-Session-ID). Generic: identity_headers (any name(s) per identity tier -- session/parent/root -- via --session-routing-opt). Body-based: dynamo_nvext (nvext.session_control bind/close request-body metadata; --session-routing-opt timeout_seconds=N). +
_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`, `identity_headers`]_ -#### `--dynamo-session-timeout-seconds` `` +#### `--session-routing-opt` `` -Dynamo nvext.session_control timeout in seconds when --use-dynamo-conv-aware-routing is enabled. -
_Constraints: ≥ 1_ -
_Default: `300`_ +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. Commas inside the value are passed through to the plugin (repeat the flag for multiple opts). +
_Default: `[]`_ #### `--connection-reuse-strategy` `` @@ -1306,7 +1300,7 @@ Explore AIPerf plugins: aiperf plugins [category] [type] #### `--category` `` Category to explore. -
_Choices: [`accumulator`, `accuracy_benchmark`, `accuracy_grader`, `analyzer`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `plot`, `public_dataset_loader`, `ramp`, `record_processor`, `service`, `service_manager`, `stream_exporter`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_ +
_Choices: [`accumulator`, `accuracy_benchmark`, `accuracy_grader`, `analyzer`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `plot`, `public_dataset_loader`, `ramp`, `record_processor`, `service`, `service_manager`, `session_routing`, `stream_exporter`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_ #### `--name` `` diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 3e1be99b38..65857eea38 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -135,8 +135,6 @@ HTTP client socket and connection configuration. Controls low-level socket optio | `AIPERF_HTTP_REQUEST_CANCELLATION_SEND_TIMEOUT` | `300.0` | ≥ 10.0, ≤ 3600.0 | Safety net timeout in seconds for waiting for HTTP request to be fully sent when request cancellation is enabled. Used as fallback when no explicit timeout is configured to prevent hanging indefinitely while waiting for the request to be written to the socket. | | `AIPERF_HTTP_IP_VERSION` | `'4'` | — | IP version for HTTP socket connections. Options: '4' (AF_INET, default), '6' (AF_INET6), or 'auto' (AF_UNSPEC, system chooses). | | `AIPERF_HTTP_TRUST_ENV` | `False` | — | Trust environment variables for HTTP client configuration. When enabled, aiohttp will read proxy settings from HTTP_PROXY, HTTPS_PROXY, and NO_PROXY environment variables. | -| `AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID` | `False` | — | Also send X-Session-ID with the stable X-Correlation-ID value. Use this when an external router requires a session-affinity header. | -| `AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID` | `False` | — | Also send X-SMG-Routing-Key with the stable X-Correlation-ID value. Use this with the SGLang Model Gateway manual routing policy. | | `AIPERF_HTTP_VIDEO_POLL_INTERVAL` | `0.1` | ≥ 0.001, ≤ 10.0 | Interval in seconds between status polls for async video generation jobs. Lower values provide faster completion detection but increase server load. Applies to the aiohttp transport. | ## LOGGING diff --git a/docs/plugins/plugin-system.md b/docs/plugins/plugin-system.md index a22b003c9b..b5af198608 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 27 plugin categories organized by function: +AIPerf supports 31 plugin categories organized by function: ### Timing Categories @@ -128,6 +128,59 @@ AIPerf supports 27 plugin categories organized by function: | `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` | Preset: additive `X-Session-ID` header carrying the session's correlation ID. No options. | +| `identity_headers` | `IdentityHeadersRouting` | Fully generic tiered identity headers: emits exactly what you configure, nothing by default. Options (each a comma-separable name list): `session` (this session's ID), `parent` (immediate parent's ID; omitted for roots), `root` (session-tree root's ID, for whole-tree affinity). Requires at least one name across tiers; names must be unique (case-insensitive) across all tiers combined, including within a single tier. | + +The three header presets are expressible via `identity_headers`: `session_id_header` ≡ +`session=X-Session-ID`, `smg_routing_key` ≡ `session=X-SMG-Routing-Key`, and `dynamo_headers` ≡ +`session=X-Dynamo-Session-ID` + `parent=X-Dynamo-Parent-Session-ID`. Reach for `identity_headers` +when a router needs custom names, several headers stamped with the same value (layered routers), +or tree-scoped affinity; reach for a preset when it already says exactly what you mean. + +**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. Commas inside a value are passed through +to the plugin (repeat the flag for multiple opts). `--session-routing-opt` without +`--session-routing` is an error, and parameterless modes (`dynamo_headers`, `smg_routing_key`, +`session_id_header`) reject every opt key. + +**`mutates_body` and the PAYLOAD_BYTES fast path.** A plugin sets the `mutates_body` class var to +`True` when `transform_body` changes the payload (only `dynamo_nvext` does). Body-mutating modes are +incompatible with the verbatim PAYLOAD_BYTES mmap fast path and are gated off it at three points: +dataset build, cache hit, and runtime. Header-only modes leave the body untouched and keep the fast +path. `transform_body` must never mutate its input — the structured 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, terminal context overflow, 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/config/config_defaults.py b/src/aiperf/common/config/config_defaults.py index cad890feb4..3b16b7e2ad 100644 --- a/src/aiperf/common/config/config_defaults.py +++ b/src/aiperf/common/config/config_defaults.py @@ -48,9 +48,6 @@ class EndpointDefaults: CONNECTION_REUSE_STRATEGY = ConnectionReuseStrategy.POOLED DOWNLOAD_VIDEO_CONTENT = False REQUEST_CONTENT_TYPE = None - USE_DYNAMO_CONV_AWARE_ROUTING = False - USE_LEGACY_DYNAMO_SESSION_CONTROL = False - DYNAMO_SESSION_TIMEOUT_SECONDS = 300 # Readiness probe defaults. Timeout 0 disables the probe (the default); # any positive value enables it. Interval is only consulted when the # probe is enabled but is validated positive so mis-configuration diff --git a/src/aiperf/common/config/endpoint_config.py b/src/aiperf/common/config/endpoint_config.py index 4543f52157..4d3cba2496 100644 --- a/src/aiperf/common/config/endpoint_config.py +++ b/src/aiperf/common/config/endpoint_config.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Annotated, Literal +from typing import Annotated, Any, Literal from pydantic import ( BeforeValidator, @@ -14,7 +14,7 @@ from aiperf.common.aiperf_logger import AIPerfLogger from aiperf.common.config.base_config import BaseConfig -from aiperf.common.config.cli_parameter import CLIParameter +from aiperf.common.config.cli_parameter import CLIParameter, DisableCLI from aiperf.common.config.config_defaults import EndpointDefaults from aiperf.common.config.config_validators import parse_str_or_list from aiperf.common.config.groups import Groups @@ -26,6 +26,7 @@ from aiperf.common.redact import REDACTED_VALUE from aiperf.plugin.enums import ( EndpointType, + SessionRoutingType, TransportType, URLSelectionStrategy, ) @@ -33,6 +34,37 @@ _logger = AIPerfLogger(__name__) +def _one_opt_or_list(value: Any) -> Any: + """Wrap a bare string opt into a single-item list WITHOUT comma-splitting. + + Unlike ``parse_str_or_list``, commas are preserved: they are legal inside + opt VALUES (e.g. ``session=X-Session-ID,X-SMG-Routing-Key``) and are + interpreted by the selected plugin's Options model, not the CLI layer. + Repeat the flag to pass multiple opts. + """ + if isinstance(value, str): + return [value] + return value + + +def _parse_session_routing_opts(values: list[str]) -> dict[str, str]: + """Parse repeatable ``key=value`` pairs into a dict, rejecting malformed + or duplicate entries with an actionable error. + """ + 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 + + class EndpointConfig(BaseConfig): """ A configuration class for defining endpoint related settings. @@ -85,20 +117,34 @@ def validate_wait_for_model_coherent(self) -> Self: return self @model_validator(mode="after") - def validate_dynamo_session_control_coherent(self) -> Self: - """Reject --use-legacy-dynamo-session-control unless conversation-aware - routing is enabled, since the legacy flag only selects the wire contract - for the session_control that --use-dynamo-conv-aware-routing emits. + def validate_session_routing(self) -> Self: + """Fail fast: opts require a mode; opts must satisfy the plugin's Options. + + Parses the repeatable ``--session-routing-opt key=value`` pairs (merged + over any ``session_routing_opts`` set directly in a config file) and + canonicalizes them to the plugin's Options model types, so downstream + consumers (including the pickled UserConfig that reaches workers) carry + coerced values (e.g. ``{"timeout_seconds": 600}``, not ``"600"``). """ - if ( - self.use_legacy_dynamo_session_control - and not self.use_dynamo_conv_aware_routing - ): - raise ValueError( - "--use-legacy-dynamo-session-control has no effect unless " - "--use-dynamo-conv-aware-routing is enabled. Enable conversation-" - "aware routing, or drop the legacy flag." - ) + opts = { + **self.session_routing_opts, + **_parse_session_routing_opts(self.session_routing_opt), + } + if self.session_routing is None: + if opts: + raise ValueError( + "--session-routing-opt requires --session-routing to select a mode." + ) + return self + # Lazy import to avoid circular dependency + 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(**opts) + self.session_routing_opts = options.model_dump(mode="json", exclude_unset=True) return self model_names: Annotated[ @@ -339,59 +385,56 @@ def url(self) -> str: ), ] = EndpointDefaults.USE_SERVER_TOKEN_COUNT - use_dynamo_conv_aware_routing: Annotated[ - bool, + session_routing: Annotated[ + SessionRoutingType | None, Field( description=( - "Emit Dynamo nvext.session_control in OpenAI-compatible request " - "bodies so Dynamo can bind all turns from the same replayed " - "conversation lineage to the same backend worker. This is only " - "intended for Dynamo frontends that implement session_control." + "Session-aware routing mode: stamps per-session identity on " + "every request for router affinity. Presets: dynamo_headers " + "(X-Dynamo-Session-ID + parent header), smg_routing_key " + "(X-SMG-Routing-Key for the SGLang Model Gateway manual " + "policy), session_id_header (additive X-Session-ID). Generic: " + "identity_headers (any name(s) per identity tier -- session/" + "parent/root -- via --session-routing-opt). Body-based: " + "dynamo_nvext (nvext.session_control bind/close request-body " + "metadata; --session-routing-opt timeout_seconds=N)." ), ), CLIParameter( - name=( - "--use-dynamo-conv-aware-routing", - "--use-dynamo-session-control", - ), + name=("--session-routing",), group=Groups.ENDPOINT, ), - ] = EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING + ] = None - use_legacy_dynamo_session_control: Annotated[ - bool, + session_routing_opt: Annotated[ + list[str], Field( description=( - "Emit the legacy Dynamo nvext.session_control lifecycle that " - "released Dynamo (v1.2.x) understands: action 'open' on the first " - "turn, session_id only on intermediate turns, and action 'close' " - "on the final turn. Use this when the target Dynamo predates the " - "'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected " - "with an HTTP 400. Requires --use-dynamo-conv-aware-routing, and " - "the Dynamo deployment must expose a worker session_control " - "endpoint for 'open' to take effect." + "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. Commas inside " + "the value are passed through to the plugin (repeat the flag " + "for multiple opts)." ), ), + BeforeValidator(_one_opt_or_list), CLIParameter( - name=("--use-legacy-dynamo-session-control",), + name=("--session-routing-opt",), + consume_multiple=True, group=Groups.ENDPOINT, ), - ] = EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL + ] = [] - dynamo_session_timeout_seconds: Annotated[ - int, + session_routing_opts: Annotated[ + dict[str, Any], Field( - description=( - "Dynamo nvext.session_control timeout in seconds when " - "--use-dynamo-conv-aware-routing is enabled." - ), - ge=1, - ), - CLIParameter( - name=("--dynamo-session-timeout-seconds",), - group=Groups.ENDPOINT, + description="Runtime-canonicalized options dict for --session-routing, " + "parsed from --session-routing-opt key=value pairs and coerced to the " + "selected plugin's Options model types. Not user-settable on the CLI.", + json_schema_extra={"add_to_template": False}, ), - ] = EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS + DisableCLI(reason="Runtime-stamped from --session-routing-opt"), + ] = {} connection_reuse_strategy: Annotated[ ConnectionReuseStrategy, diff --git a/src/aiperf/common/environment.py b/src/aiperf/common/environment.py index a7c6510188..aa17b19c02 100644 --- a/src/aiperf/common/environment.py +++ b/src/aiperf/common/environment.py @@ -675,16 +675,6 @@ class _HTTPSettings(BaseSettings): "When enabled, aiohttp will read proxy settings from HTTP_PROXY, HTTPS_PROXY, " "and NO_PROXY environment variables.", ) - X_SESSION_ID_FROM_CORRELATION_ID: bool = Field( - default=False, - description="Also send X-Session-ID with the stable X-Correlation-ID value. " - "Use this when an external router requires a session-affinity header.", - ) - X_SMG_ROUTING_KEY_FROM_CORRELATION_ID: bool = Field( - default=False, - description="Also send X-SMG-Routing-Key with the stable X-Correlation-ID value. " - "Use this with the SGLang Model Gateway manual routing policy.", - ) VIDEO_POLL_INTERVAL: float = Field( ge=0.001, le=10.0, diff --git a/src/aiperf/common/models/model_endpoint_info.py b/src/aiperf/common/models/model_endpoint_info.py index cd034f4b32..3653bce9c1 100644 --- a/src/aiperf/common/models/model_endpoint_info.py +++ b/src/aiperf/common/models/model_endpoint_info.py @@ -114,19 +114,13 @@ class EndpointInfo(AIPerfBaseModel): default=EndpointDefaults.USE_SERVER_TOKEN_COUNT, description="Use server-reported token counts from API usage fields instead of client-side tokenization.", ) - use_dynamo_conv_aware_routing: bool = Field( - default=EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING, - description="Emit Dynamo nvext.session_control for conversation-aware routing.", - ) - use_legacy_dynamo_session_control: bool = Field( - default=EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL, - description="Emit the v1.2.x-compatible open/close session_control lifecycle " - "instead of the 'bind' action (which only exists in Dynamo >= v1.3.0-dev).", + session_routing: str | None = Field( + default=None, + description="Selected session-routing plugin name (None = off).", ) - dynamo_session_timeout_seconds: int = Field( - default=EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS, - ge=1, - description="Timeout in seconds for Dynamo nvext.session_control sessions.", + 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, @@ -178,15 +172,12 @@ def from_user_config(cls, user_config: UserConfig) -> "EndpointInfo": api_key=user_config.endpoint.api_key, use_legacy_max_tokens=user_config.endpoint.use_legacy_max_tokens, use_server_token_count=user_config.endpoint.use_server_token_count, - use_dynamo_conv_aware_routing=( - user_config.endpoint.use_dynamo_conv_aware_routing - ), - use_legacy_dynamo_session_control=( - user_config.endpoint.use_legacy_dynamo_session_control - ), - dynamo_session_timeout_seconds=( - user_config.endpoint.dynamo_session_timeout_seconds + session_routing=( + str(user_config.endpoint.session_routing) + if user_config.endpoint.session_routing is not None + else None ), + session_routing_opts=dict(user_config.endpoint.session_routing_opts), connection_reuse_strategy=user_config.endpoint.connection_reuse_strategy, download_video_content=user_config.endpoint.download_video_content, request_content_type=user_config.endpoint.request_content_type, diff --git a/src/aiperf/common/models/record_models.py b/src/aiperf/common/models/record_models.py index ed92256fc2..c118d047c7 100644 --- a/src/aiperf/common/models/record_models.py +++ b/src/aiperf/common/models/record_models.py @@ -852,6 +852,17 @@ class RequestInfo(RecordContext): description="Whether this is the final turn in the conversation. " "Used by per-conversation connection strategy to release the connection lease.", ) + 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/credit/issuer.py b/src/aiperf/credit/issuer.py index 875295d27c..4920aa50bb 100644 --- a/src/aiperf/credit/issuer.py +++ b/src/aiperf/credit/issuer.py @@ -192,6 +192,33 @@ def _open_session_tree(self, turn: TurnToSend) -> None: 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-agentic 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 release_lane_credit(self) -> None: """Release a session slot directly (legacy / non-registry path). @@ -359,6 +386,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, @@ -374,6 +403,8 @@ async def _issue_credit_internal(self, turn: TurnToSend) -> bool: root_correlation_id=turn.root_correlation_id, counts_toward_phase_target=turn.counts_toward_phase_target, has_forks=turn.has_forks, + is_parent_final=is_parent_final, + is_tree_final=is_tree_final, branch_mode=turn.branch_mode, cache_bust_marker=turn.cache_bust_marker, cache_bust_target=turn.cache_bust_target, @@ -487,6 +518,7 @@ async def dispatch_join_turn(self, pending: PendingBranchJoin) -> bool: # the sampled root plan and counts. counts_toward_phase_target=pending.parent_agent_depth == 0, has_forks=pending.parent_has_forks_on_gated_turn, + has_branches=pending.parent_has_branches_on_gated_turn, branch_mode=pending.parent_branch_mode, cache_bust_marker=pending.parent_cache_bust_marker, cache_bust_target=pending.parent_cache_bust_target, diff --git a/src/aiperf/credit/structs.py b/src/aiperf/credit/structs.py index 2fa8a8125a..b851bea781 100644 --- a/src/aiperf/credit/structs.py +++ b/src/aiperf/credit/structs.py @@ -72,6 +72,12 @@ class Credit( ``agent_depth``. """ 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 = @@ -169,6 +175,14 @@ class TurnToSend(Struct, frozen=True): session start regardless of this flag. A mid-trace session start can only legitimately occur during a phase's initial dispatch (execute_phase).""" 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 cache_bust_marker: str | None = None @@ -200,7 +214,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, @@ -212,6 +229,7 @@ def from_previous_credit( root_correlation_id=credit.root_correlation_id, counts_toward_phase_target=credit.counts_toward_phase_target, 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, cache_bust_marker=credit.cache_bust_marker, cache_bust_target=credit.cache_bust_target, diff --git a/src/aiperf/dataset/dataset_manager.py b/src/aiperf/dataset/dataset_manager.py index a30c08262c..68d353d7d9 100644 --- a/src/aiperf/dataset/dataset_manager.py +++ b/src/aiperf/dataset/dataset_manager.py @@ -264,12 +264,13 @@ def _lookup_under_lock(self) -> mmap_cache.CacheHit | None: """Re-check the cache for a HIT after the populate lock is held.""" assert self._cache_key_for_run is not None try: - return mmap_cache.lookup( + hit = mmap_cache.lookup( self._cache_key_for_run, compressed=self._compress_only ) except (OSError, ValueError) as e: self.warning(f"Cache re-lookup under lock failed: {e!r}") return None + return self._downgrade_body_mutator_cache_hit(hit) async def _configure_dataset_client_and_free_memory(self) -> None: """Configure the dataset client for serving fallback requests, then free memory.""" @@ -414,13 +415,15 @@ def _preformat_payloads(self, conversations: list[Conversation]) -> None: if self.user_config is None: return - # Cache-bust dispatch (worker.py `_process_credit_with_session`) mutates - # `session_message`/`raw_messages` per credit; the PAYLOAD_BYTES fast - # path early-returns before that dispatch, sending pre-encoded mmap - # bytes to the wire verbatim. Pre-formatting under cache-bust would - # silently no-op the marker injection. Bail to the structured-turns - # path whenever cache-bust is enabled. - if self.user_config.input.prompt.cache_bust.target != CacheBustTarget.NONE: + # Body mutators (cache-bust marker injection, a body-mutating + # session-routing mode like dynamo_nvext) rewrite the request body per + # credit; the PAYLOAD_BYTES fast path early-returns before that + # dispatch, sending pre-encoded mmap bytes to the wire verbatim. + # Pre-formatting would either silently no-op the mutation or trip the + # _select_mmap_format refusal on a dataset that is perfectly usable as + # structured turns. Bail to the structured-turns path whenever any + # body-mutating feature is active. + if self._body_mutating_feature() is not None: return # DAG datasets (any FORK/SPAWN branch) are delta-compressed and @@ -476,8 +479,51 @@ def _preformat_payloads(self, conversations: list[Conversation]) -> None: self.info(f"Pre-formatted {count} payloads for payload mmap fast path") + def _body_mutating_feature(self) -> str | None: + """Name of the active body-mutating feature, or None. + + Both PAYLOAD_BYTES gates (build-path format selection and cache-hit + adoption) key off this: the mmap cache key deliberately excludes + cache-bust and routing settings, so the cache-hit gate cannot be + skipped (a feature-free run's PAYLOAD_BYTES entry is a valid hit for + a feature-enabled run). + """ + if self.user_config is None: + return None + mode = self.user_config.endpoint.session_routing + if mode is not None: + routing_cls = plugins.get_class(PluginType.SESSION_ROUTING, str(mode)) + if routing_cls.mutates_body: + return f"session-routing mode {str(mode)!r}" + if self.user_config.input.prompt.cache_bust.target != CacheBustTarget.NONE: + return "cache-bust" + return None + + def _reject_body_mutators_for_payload_bytes( + self, mmap_format: MemoryMapFormat | str + ) -> None: + """Refuse a cached PAYLOAD_BYTES dataset when a body-mutator is active. + + Defense-in-depth behind ``_downgrade_body_mutator_cache_hit``: promoted + entries are already treated as a MISS at lookup, so a HIT that reaches + here carries source-loaded payload bytes (the dataset itself ships + pre-encoded bodies), which a rebuild cannot fix. Guarded at the shared + restore path so both cache-hit call sites are covered. + """ + if MemoryMapFormat(mmap_format) != MemoryMapFormat.PAYLOAD_BYTES: + return + feature = self._body_mutating_feature() + if feature is not None: + raise ValueError( + f"{feature} must mutate request bodies and is incompatible with " + "this cached PAYLOAD_BYTES dataset (its payload bytes are shipped " + "verbatim by the dataset itself). Choose a headers-based routing " + "mode / disable cache-bust, or use a dataset type that produces " + "structured turns (e.g. single_turn / multi_turn / dag_jsonl)." + ) + def _select_mmap_format(self, conversations: list[Conversation]) -> MemoryMapFormat: - """Pick the dataset mmap format and refuse PAYLOAD_BYTES under cache-bust. + """Pick the dataset mmap format and refuse PAYLOAD_BYTES for body-mutators. This is the earliest authoritative point in the loader where the run's ``MemoryMapFormat`` is finalized -- it runs once after dataset @@ -485,13 +531,14 @@ def _select_mmap_format(self, conversations: list[Conversation]) -> MemoryMapFor per-fork / per-dataset preformat decisions have happened yet. PAYLOAD_BYTES is the mmap fast path: workers stream pre-encoded - bytes verbatim and skip the cache-bust dispatch in + bytes verbatim and skip the per-credit dispatch in ``_process_credit_with_session``. Loaders that natively populate ``Turn.raw_payload`` (RawPayloadDatasetLoader, InputsJsonPayloadLoader, and MooncakeTraceDatasetLoader entries with a ``payload`` field) - would otherwise silently bypass the marker injection. Refuse here - with a clear, actionable error rather than letting the worker - discover the conflict at runtime. + would otherwise silently bypass cache-bust marker injection or a + body-mutating routing mode (e.g. Dynamo nvext.session_control). + Refuse here with a clear, actionable error rather than letting the + worker discover the conflict at runtime. """ has_payload_bytes = any( turn.raw_payload is not None @@ -507,33 +554,13 @@ def _select_mmap_format(self, conversations: list[Conversation]) -> MemoryMapFor "Mixed raw_payload state: all turns must have raw_payload " "when any turn does (PAYLOAD_BYTES format requires uniformity)" ) - if ( - has_payload_bytes - and self.user_config is not None - and self.user_config.input.prompt.cache_bust.target != CacheBustTarget.NONE - ): - raise ValueError( - "--cache-bust is incompatible with the PAYLOAD_BYTES mmap " - "fast path. The selected dataset (raw_payload / inputs_json " - "/ mooncake_trace with payload field) ships pre-encoded bytes " - "verbatim and bypasses the per-credit cache-bust marker " - "injection. Either remove --cache-bust, or use a dataset " - "type that produces structured turns " - "(e.g. single_turn / multi_turn / dag_jsonl)." - ) - if ( - has_payload_bytes - and self.user_config is not None - and self.user_config.endpoint.use_dynamo_conv_aware_routing - ): + feature = self._body_mutating_feature() + if has_payload_bytes and feature is not None: raise ValueError( - "--use-dynamo-conv-aware-routing is incompatible with the " - "PAYLOAD_BYTES mmap fast path. The selected dataset (raw_payload " - "/ inputs_json / mooncake_trace with payload field) ships " - "pre-encoded bytes verbatim, so nvext.session_control cannot be " - "injected. Either disable Dynamo conversation-aware routing, or " - "use a dataset type that produces structured turns " - "(e.g. single_turn / multi_turn / dag_jsonl)." + f"{feature} must mutate request bodies and is incompatible with the " + "verbatim PAYLOAD_BYTES mmap fast path. Choose a headers-based " + "routing mode / disable cache-bust, or use a dataset type that " + "produces structured turns (e.g. single_turn / multi_turn / dag_jsonl)." ) return ( MemoryMapFormat.PAYLOAD_BYTES @@ -821,10 +848,40 @@ def _try_cache_lookup(self) -> mmap_cache.CacheHit | None: return None self._cache_key_for_run = key try: - return mmap_cache.lookup(key, compressed=self._compress_only) + hit = mmap_cache.lookup(key, compressed=self._compress_only) except Exception as e: self.warning(f"Skipping mmap cache lookup: {e!r}") return None + return self._downgrade_body_mutator_cache_hit(hit) + + def _downgrade_body_mutator_cache_hit( + self, hit: mmap_cache.CacheHit | None + ) -> mmap_cache.CacheHit | None: + """Treat a preformat-promoted PAYLOAD_BYTES HIT as a MISS under a body-mutator. + + The cache key deliberately excludes routing/cache-bust settings, so a + feature-free run's PAYLOAD_BYTES entry is a valid HIT for a + feature-enabled run. Entries promoted by ``_preformat_payloads`` + (``all_turns_source_loaded_payloads=False``) can simply be rebuilt: with + the feature active the preformatter bails, so the rebuild stays on the + structured CONVERSATION path and succeeds. Entries whose payload bytes + came from the dataset itself keep the hard-fail in + ``_reject_body_mutators_for_payload_bytes`` (a rebuild cannot help). + """ + if hit is None: + return None + if MemoryMapFormat(hit.manifest.mmap_format) != MemoryMapFormat.PAYLOAD_BYTES: + return hit + if hit.manifest.all_turns_source_loaded_payloads: + return hit + feature = self._body_mutating_feature() + if feature is None: + return hit + self.info( + f"Ignoring PAYLOAD_BYTES mmap cache entry built without {feature}; " + "rebuilding on the structured-turns path." + ) + return None async def _configure_from_cache_hit(self, hit: mmap_cache.CacheHit) -> None: """Restore mmap files + metadata from a cache HIT, then init backing store. @@ -833,6 +890,8 @@ async def _configure_from_cache_hit(self, hit: mmap_cache.CacheHit) -> None: rest of the pipeline (backing-store cleanup, worker mmap reads, k8s download) sees byte-identical files to a non-cached run. """ + self._reject_body_mutators_for_payload_bytes(hit.manifest.mmap_format) + run_data_path, run_index_path = self._run_mmap_paths() mmap_cache.restore_to_run_dir(hit, run_data_path, run_index_path) @@ -895,6 +954,14 @@ def _populate_cache_after_run(self) -> None: """Write the just-finalized run's mmap files into the cache.""" if self._cache_hit_used: return + # The cache key deliberately excludes routing/cache-bust settings, so + # an entry written by THIS run is a valid HIT for feature-free runs of + # the same dataset. A body-mutator run built CONVERSATION only because + # the preformat bailed -- caching that would pin the shared key to the + # slow path and silently demote every later feature-free run off + # PAYLOAD_BYTES. Leave the key for a feature-free run to populate. + if self._body_mutating_feature() is not None: + return if self._cache_key_for_run is None or self._backing_store is None: return if self.dataset_metadata is None: diff --git a/src/aiperf/exporters/console_api_error_exporter.py b/src/aiperf/exporters/console_api_error_exporter.py index 2a323f14c2..4e8ad1e6da 100644 --- a/src/aiperf/exporters/console_api_error_exporter.py +++ b/src/aiperf/exporters/console_api_error_exporter.py @@ -112,7 +112,7 @@ def detect(error_summary: list[ErrorDetailsCount]) -> ErrorInsight | None: "v1.3.0-dev / upstream commit d97c889ba)." ), causes=[ - "--use-dynamo-conv-aware-routing emits action='bind' on every non-final turn.", + "--session-routing dynamo_nvext emits action='bind' on every non-final turn.", "The target Dynamo server predates the 'bind' action (e.g. v1.2.1).", ], investigation=[ @@ -121,8 +121,32 @@ def detect(error_summary: list[ErrorDetailsCount]) -> ErrorInsight | None: ], fixes=[ "Upgrade Dynamo to a build that supports action='bind' (>= v1.3.0-dev, upstream commit d97c889ba).", - "Or run with --use-legacy-dynamo-session-control to emit the v1.2.x-compatible open/close lifecycle (requires the worker to expose a session_control endpoint).", - "Or disable --use-dynamo-conv-aware-routing.", + "Or switch to --session-routing dynamo_headers (header-based affinity; requires --router-session-affinity-ttl-secs on the Dynamo frontend).", + ], + ) + + # Current upstream Dynamo removed session_control and its NvExt + # deserializer denies unknown fields, so the whole field is rejected. + if "unknown field" in error_blob and "session_control" in error_blob: + return ErrorInsight( + title="Dynamo build does not implement nvext.session_control", + problem=( + "The Dynamo frontend rejected nvext.session_control as an unknown " + "field. Current upstream Dynamo removed session_control; its NvExt " + "deserializer denies unknown fields, so every request fails." + ), + causes=[ + "--session-routing dynamo_nvext emits nvext.session_control on every turn.", + "The target Dynamo build has no session_control support (e.g. current main).", + ], + investigation=[ + "Check the Dynamo frontend version/build for session_control support.", + "Inspect request payloads in profile_export.jsonl -> nvext.session_control.", + ], + fixes=[ + "Switch to --session-routing dynamo_headers and run the Dynamo " + "frontend with --router-session-affinity-ttl-secs.", + "Or drop --session-routing for an untagged baseline.", ], ) diff --git a/src/aiperf/plugin/categories.yaml b/src/aiperf/plugin/categories.yaml index 19a1e4401a..92c35def76 100644 --- a/src/aiperf/plugin/categories.yaml +++ b/src/aiperf/plugin/categories.yaml @@ -140,6 +140,22 @@ 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. Body-mutating plugins (mutates_body=True) are rejected + against the verbatim PAYLOAD_BYTES mmap fast path. + # ============================================================================= # Processing Categories # ============================================================================= diff --git a/src/aiperf/plugin/enums.py b/src/aiperf/plugin/enums.py index bc45939c96..af0696f467 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", "AccumulatorType", "AccumulatorTypeStr", "AccuracyBenchmarkType", "AccuracyBenchmarkTypeStr", "AccuracyGraderType", "AccuracyGraderTypeStr", "AnalyzerType", "AnalyzerTypeStr", "ArrivalPattern", "ArrivalPatternStr", "CommClientType", "CommClientTypeStr", "CommunicationBackend", "CommunicationBackendStr", "ComposerType", "ComposerTypeStr", "ConsoleExporterType", "ConsoleExporterTypeStr", "CustomDatasetType", "CustomDatasetTypeStr", "DataExporterType", "DataExporterTypeStr", "DatasetBackingStoreType", "DatasetBackingStoreTypeStr", "DatasetClientStoreType", "DatasetClientStoreTypeStr", "DatasetSamplingStrategy", "DatasetSamplingStrategyStr", "EndpointType", "EndpointTypeStr", "GPUTelemetryCollectorType", "GPUTelemetryCollectorTypeStr", "PlotType", "PlotTypeStr", "PluginType", "PluginTypeStr", "PublicDatasetType", "PublicDatasetTypeStr", "RampType", "RampTypeStr", "RecordProcessorType", "RecordProcessorTypeStr", "ServiceRunType", "ServiceRunTypeStr", "ServiceType", "ServiceTypeStr", "StreamExporterType", "StreamExporterTypeStr", "TimingMode", "TimingModeStr", "TransportType", "TransportTypeStr", "UIType", "UITypeStr", "URLSelectionStrategy", "URLSelectionStrategyStr", "ZMQProxyType", "ZMQProxyTypeStr"] +__all__ = ["APIRouterType", "APIRouterTypeStr", "AccumulatorType", "AccumulatorTypeStr", "AccuracyBenchmarkType", "AccuracyBenchmarkTypeStr", "AccuracyGraderType", "AccuracyGraderTypeStr", "AnalyzerType", "AnalyzerTypeStr", "ArrivalPattern", "ArrivalPatternStr", "CommClientType", "CommClientTypeStr", "CommunicationBackend", "CommunicationBackendStr", "ComposerType", "ComposerTypeStr", "ConsoleExporterType", "ConsoleExporterTypeStr", "CustomDatasetType", "CustomDatasetTypeStr", "DataExporterType", "DataExporterTypeStr", "DatasetBackingStoreType", "DatasetBackingStoreTypeStr", "DatasetClientStoreType", "DatasetClientStoreTypeStr", "DatasetSamplingStrategy", "DatasetSamplingStrategyStr", "EndpointType", "EndpointTypeStr", "GPUTelemetryCollectorType", "GPUTelemetryCollectorTypeStr", "PlotType", "PlotTypeStr", "PluginType", "PluginTypeStr", "PublicDatasetType", "PublicDatasetTypeStr", "RampType", "RampTypeStr", "RecordProcessorType", "RecordProcessorTypeStr", "ServiceRunType", "ServiceRunTypeStr", "ServiceType", "ServiceTypeStr", "SessionRoutingType", "SessionRoutingTypeStr", "StreamExporterType", "StreamExporterTypeStr", "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.IDENTITY_HEADERS, 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.METRIC_RECORD, RecordProcessorType.RAW_RECORD_WRITER""" diff --git a/src/aiperf/plugin/plugins.py b/src/aiperf/plugin/plugins.py index e9c33bf6df..3226cafbf7 100644 --- a/src/aiperf/plugin/plugins.py +++ b/src/aiperf/plugin/plugins.py @@ -935,7 +935,7 @@ def _load_package_metadata( from aiperf.exporters.protocols import ArtifactPublisherProtocol, ConsoleExporterProtocol, DataExporterProtocol from aiperf.gpu_telemetry.protocols import GPUTelemetryCollectorProtocol from aiperf.plot.core.plot_type_handlers import PlotTypeHandlerProtocol - from aiperf.plugin.enums import APIRouterType, AccumulatorType, AccuracyBenchmarkType, AccuracyGraderType, AnalyzerType, ArrivalPattern, ArtifactPublisherType, CommClientType, CommunicationBackend, ComposerType, ConsoleExporterType, CustomDatasetType, DataExporterType, DatasetBackingStoreType, DatasetClientStoreType, DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, PlotType, PluginType, PluginTypeStr, PublicDatasetType, RampType, RecordProcessorType, ServiceRunType, ServiceType, StreamExporterType, TimingMode, TransportType, UIType, URLSelectionStrategy, ZMQProxyType + from aiperf.plugin.enums import APIRouterType, AccumulatorType, AccuracyBenchmarkType, AccuracyGraderType, AnalyzerType, ArrivalPattern, ArtifactPublisherType, CommClientType, CommunicationBackend, ComposerType, ConsoleExporterType, CustomDatasetType, DataExporterType, DatasetBackingStoreType, DatasetClientStoreType, DatasetSamplingStrategy, EndpointType, GPUTelemetryCollectorType, PlotType, PluginType, PluginTypeStr, PublicDatasetType, RampType, RecordProcessorType, ServiceRunType, ServiceType, SessionRoutingType, StreamExporterType, TimingMode, TransportType, UIType, URLSelectionStrategy, ZMQProxyType from aiperf.post_processors.protocols import RecordProcessorProtocol from aiperf.timing.intervals import IntervalGeneratorProtocol from aiperf.timing.ramping import RampStrategyProtocol @@ -943,6 +943,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 # @@ -996,6 +997,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 52b48ba063..6c3c55474e 100644 --- a/src/aiperf/plugin/plugins.yaml +++ b/src/aiperf/plugin/plugins.yaml @@ -426,6 +426,59 @@ 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). Mutates request + bodies, so it is incompatible with the PAYLOAD_BYTES mmap fast path. + 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: | + Preset: additive X-Session-ID header carrying the session's + correlation ID. No options; equivalent to identity_headers with + session=X-Session-ID (use identity_headers for custom names or + parent/root tiers). + + identity_headers: + class: aiperf.workers.session_routing:IdentityHeadersRouting + description: | + Fully generic tiered identity headers: emits exactly the headers + configured, nothing by default. Options (each a comma-separable + list of names): session (this session's correlation ID), parent + (immediate parent's ID, omitted for roots), root (session-tree + root's ID, for whole-tree affinity). At least one name across + tiers is required; names must be unique (case-insensitive) across + all tiers combined. The named + presets are expressible via this mode, e.g. dynamo_headers is + session=X-Dynamo-Session-ID + parent=X-Dynamo-Parent-Session-ID. + # ============================================================================= # Dataset Composers # ============================================================================= diff --git a/src/aiperf/plugin/schema/plugins.schema.json b/src/aiperf/plugin/schema/plugins.schema.json index d683c841dd..de84eb21d0 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. Body-mutating plugins (mutates_body=True) are rejected\nagainst the verbatim PAYLOAD_BYTES mmap fast path.", + "additionalProperties": { + "$ref": "#/$defs/SessionRoutingPlugin" + } + }, "record_processor": { "title": "Record Processor Plugins", "type": "object", @@ -1047,6 +1055,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. Body-mutating plugins (mutates_body=True) are rejected\nagainst the verbatim PAYLOAD_BYTES mmap fast path." + }, "RecordProcessorPlugin": { "type": "object", "properties": { diff --git a/src/aiperf/timing/branch_orchestrator.py b/src/aiperf/timing/branch_orchestrator.py index 24e964e538..230875c16b 100644 --- a/src/aiperf/timing/branch_orchestrator.py +++ b/src/aiperf/timing/branch_orchestrator.py @@ -184,6 +184,11 @@ class PendingBranchJoin: 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) # Cache-bust state captured from the credit that suspends the parent so @@ -677,8 +682,10 @@ def _ensure_seeded_join( return future 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, @@ -689,6 +696,7 @@ def _ensure_seeded_join( gated_turn_index=gated_idx, parent_branch_mode=parent_state.branch_mode, parent_has_forks_on_gated_turn=has_forks, + parent_has_branches_on_gated_turn=has_branches, parent_cache_bust_marker=cache_bust_marker, parent_cache_bust_target=self._cache_bust_target, ) @@ -1174,10 +1182,12 @@ 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, @@ -1189,6 +1199,7 @@ def _ensure_future_join( credit, "branch_mode", ConversationBranchMode.FORK ), parent_has_forks_on_gated_turn=has_forks, + parent_has_branches_on_gated_turn=has_branches, # Capture parent's cache-bust state from the suspending # credit so the join turn (k+1) inherits the same marker # as turns 0..k. The credit always has these fields diff --git a/src/aiperf/timing/conversation_source.py b/src/aiperf/timing/conversation_source.py index e18c29e73b..280ec69921 100644 --- a/src/aiperf/timing/conversation_source.py +++ b/src/aiperf/timing/conversation_source.py @@ -89,6 +89,9 @@ def build_first_turn(self, max_turns: int | None = None) -> TurnToSend: root_correlation_id=self.root_correlation_id, is_session_start=True, 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, cache_bust_marker=self.cache_bust_marker, cache_bust_target=self.cache_bust_target, @@ -125,6 +128,7 @@ def build_turn_at_index(self, turn_index: int) -> TurnToSend: # the resumed root acquires a session slot + counts even at k_i > 0. is_session_start=True, has_forks=meta.has_forks if meta is not None else False, + has_branches=bool(meta.branch_ids) if meta is not None else False, branch_mode=self.branch_mode, ) diff --git a/src/aiperf/timing/session_tree.py b/src/aiperf/timing/session_tree.py index 485518d69e..540d8f77c3 100644 --- a/src/aiperf/timing/session_tree.py +++ b/src/aiperf/timing/session_tree.py @@ -159,6 +159,44 @@ 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 + popped 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 or snapshot-seeded) to a tree. diff --git a/src/aiperf/timing/strategies/agentic_replay.py b/src/aiperf/timing/strategies/agentic_replay.py index 462ef54cb9..bb66c6b7f0 100644 --- a/src/aiperf/timing/strategies/agentic_replay.py +++ b/src/aiperf/timing/strategies/agentic_replay.py @@ -1488,8 +1488,9 @@ def _get_snapshot(self, trajectory: Trajectory) -> TrajectorySnapshot: ``TrajectorySource`` constructs each timestamped lane once and is shared across WARMUP and PROFILING. Reusing that realized graph keeps - every continuing root and subagent on the same ``X-Session-ID`` across - the phase boundary. + every continuing root and subagent on the same ``x_correlation_id`` + across the phase boundary (and therefore on the same affinity header + when a ``--session-routing`` header mode is active). """ assert trajectory.snapshot is not None return trajectory.snapshot diff --git a/src/aiperf/timing/strategies/fixed_schedule.py b/src/aiperf/timing/strategies/fixed_schedule.py index 759c03e1a4..3bd0d2b677 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), ), ) ) diff --git a/src/aiperf/transports/base_transports.py b/src/aiperf/transports/base_transports.py index 1b57f453ce..c4a1493567 100644 --- a/src/aiperf/transports/base_transports.py +++ b/src/aiperf/transports/base_transports.py @@ -8,7 +8,6 @@ from typing import Protocol, runtime_checkable from urllib.parse import parse_qs, urlencode, urlparse, urlunparse -from aiperf.common.environment import Environment from aiperf.common.mixins import AIPerfLifecycleMixin from aiperf.common.models import ( RequestInfo, @@ -117,10 +116,6 @@ def build_headers(self, request_info: RequestInfo) -> dict[str, str]: headers["X-Request-ID"] = request_info.x_request_id if request_info.x_correlation_id: headers["X-Correlation-ID"] = request_info.x_correlation_id - if Environment.HTTP.X_SESSION_ID_FROM_CORRELATION_ID: - headers["X-Session-ID"] = request_info.x_correlation_id - if Environment.HTTP.X_SMG_ROUTING_KEY_FROM_CORRELATION_ID: - headers["X-SMG-Routing-Key"] = request_info.x_correlation_id headers.update(request_info.endpoint_headers) headers.update(self.get_transport_headers(request_info)) diff --git a/src/aiperf/workers/dynamo_session_control.py b/src/aiperf/workers/dynamo_session_control.py deleted file mode 100644 index d3a45429c8..0000000000 --- a/src/aiperf/workers/dynamo_session_control.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Dynamo conversation-aware routing via ``nvext.session_control``. - -Opt-in helpers that shape the ``nvext.session_control`` block Dynamo frontends -read to pin every turn of a replayed conversation to the same backend worker -(reusing its prefix KV cache). Kept separate from request building so the -worker stays unaware of wire-body internals: the policy (which lifecycle action -to emit) lives in :func:`build_session_control`, and the mechanism (overlaying -it onto the structured request body) in :func:`merge_session_control`. The -single caller is the serialization chokepoint -:meth:`aiperf.workers.inference_client.InferenceClient._send_request_to_transport`, -right after the endpoint formats the request dict. - -The verbatim PAYLOAD_BYTES mmap fast path (raw_payload / inputs_json / -mooncake-with-payload datasets) is refused against this feature at dataset load -(:meth:`aiperf.dataset.dataset_manager.DatasetManager._select_mmap_format`), so -this module only ever sees a freshly built request dict -- mirroring how -``--cache-bust`` is handled. -""" - -from __future__ import annotations - -from typing import Any - - -def build_session_control( - *, - session_id: str, - is_final_turn: bool, - timeout_seconds: int, - legacy: bool = False, - already_opened: bool = False, -) -> dict[str, Any]: - """Build the ``nvext.session_control`` block for a single request. - - Each conversation instance is its own sticky session keyed by its - X-Correlation-ID (``session_id``). Dynamo's router co-locates turns purely - by ``session_id``. Two wire contracts are produced depending on ``legacy``. - - Modern (``legacy=False``, the default) -- targets Dynamo builds that - implement the ``bind`` action (>= v1.3.0-dev / upstream commit d97c889ba): - - - non-final turn -> ``action: "bind"`` (router-only sticky affinity). - Re-bind is idempotent on the router and refreshes the inactivity TTL, so - long inter-turn replay delays cannot silently expire affinity - mid-conversation. That is why every non-final turn re-binds rather than - only the first -- it also removes any need for per-worker bind tracking. - - final turn -> ``action: "close"`` to release the router affinity (and any - worker-side session) immediately instead of leaking it until the TTL - reaper fires -- material on high-cardinality runs (millions of sessions). - - Legacy (``legacy=True``) -- targets released Dynamo (v1.2.x), whose - ``SessionAction`` enum only accepts ``open`` and ``close`` (``bind`` does not - exist there and is rejected with an HTTP 400 ``unknown variant`` error): - - - first request the worker sends for a session (``already_opened`` False) - -> ``action: "open"`` to create the worker session and bind router - affinity. ``open`` is NOT idempotent, so the caller tracks which sessions - it has opened and passes ``already_opened`` accordingly. The trigger is - "first request the worker sends", NOT ``turn_index == 0``: under agentic - replay the first request is the WARMUP turn (k_i), and profiling resumes - mid-trace at k_i+1, so no profiling request ever carries ``turn_index 0``. - ``open`` also requires the Dynamo deployment to expose a worker - ``session_control`` endpoint; without one Dynamo silently skips session - lifecycle (no 400, but no affinity). - - subsequent turns (``already_opened`` True) -> ``session_id`` only, which - keeps affinity on the sticky-router path. - - final turn -> ``action: "close"`` (same as modern). - - Fork/spawn children use their own ``session_id`` rather than a shared - lineage key: ``close`` is keyed on ``session_id``, so a shared key would let - a child's ``close`` tear down the parent's still-live affinity. Cross-branch - KV reuse survives anyway because the child's first (still-unbound) request is - routed by Dynamo's prefix-overlap router onto the worker already holding the - shared prefix. - """ - session_control: dict[str, Any] = {"session_id": session_id} - if is_final_turn: - session_control["action"] = "close" - elif legacy: - if not already_opened: - session_control["action"] = "open" - session_control["timeout"] = timeout_seconds - else: - session_control["action"] = "bind" - session_control["timeout"] = timeout_seconds - return session_control - - -def merge_session_control( - payload: dict[str, Any], - session_control: dict[str, Any], -) -> dict[str, Any]: - """Return a copy of ``payload`` with ``session_control`` overlaid under ``nvext``. - - Never mutates ``payload`` nor its nested ``nvext`` / ``session_control`` - dicts, so it is safe on a cached ``Turn.raw_payload`` or on an ``extra_body`` - reference shared with the dataset. Existing ``nvext`` keys and any - pre-existing ``session_control`` fields are preserved; the computed fields - overlay them. - """ - merged = dict(payload) - raw_nvext = merged.get("nvext") - nvext = dict(raw_nvext) if isinstance(raw_nvext, dict) else {} - raw_session_control = nvext.get("session_control") - merged_session_control = ( - dict(raw_session_control) if isinstance(raw_session_control, dict) else {} - ) - merged_session_control.update(session_control) - nvext["session_control"] = merged_session_control - merged["nvext"] = nvext - return merged diff --git a/src/aiperf/workers/inference_client.py b/src/aiperf/workers/inference_client.py index 49b6c7be2c..c76bc44084 100644 --- a/src/aiperf/workers/inference_client.py +++ b/src/aiperf/workers/inference_client.py @@ -20,10 +20,7 @@ from aiperf.common.redact import redact_headers from aiperf.plugin import plugins from aiperf.plugin.enums import PluginType -from aiperf.workers.dynamo_session_control import ( - build_session_control, - merge_session_control, -) +from aiperf.workers.session_routing import RoutingContext, SessionRoutingBase if TYPE_CHECKING: from aiperf.transports.base_transports import FirstTokenCallback @@ -76,14 +73,54 @@ def __init__( # Resolved by the worker via record payload-retention auto-detection. self.strip_record_payload_bytes = strip_record_payload_bytes - # Legacy Dynamo session_control only: session_ids this worker has already - # sent an 'open' for. 'open' is not idempotent and must be sent exactly - # once on the first request the worker makes for a session -- which under - # agentic replay is the WARMUP turn (k_i), not turn_index 0. The - # StickyCreditRouter pins every turn of a session (warmup + profiling) to - # one worker, so this per-process set sees them all. Entries are dropped - # on 'close' to bound the set to in-flight sessions. - self._dynamo_opened_sessions: set[str] = set() + # 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 + # Per-request capability flags resolved once at init: a header-only + # plugin inherits the base transform_body and a body-only plugin + # inherits the base headers -- calling the no-op (plus the empty-dict + # merge) costs ~75ns/request, skipped via these flags (measured). + self._routing_stamps_headers = False + self._routing_transforms_body = False + 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 + # Class-level override OR instance-level binding (a plugin may + # select a strategy by assigning self.headers/self.transform_body + # in __init__) both count as capabilities. + instance_attrs = vars(self._routing) + self._routing_stamps_headers = ( + "headers" in instance_attrs + or type(self._routing).headers is not SessionRoutingBase.headers + ) + self._routing_transforms_body = ( + "transform_body" in instance_attrs + or type(self._routing).transform_body + is not SessionRoutingBase.transform_body + ) + # Fail fast on a split-brain plugin: every PAYLOAD_BYTES gate keys + # off the declarative mutates_body ClassVar, while the transform is + # invoked via this structural capability flag. If they disagree, a + # transform silently never fires on fast-path datasets (or runs are + # forced off the fast path for nothing) -- reject at worker init + # instead of shipping unrouted bodies. + if self._routing_transforms_body != self._routing.mutates_body: + raise ValueError( + f"session-routing plugin {self._routing_mode!r} is " + f"inconsistent: transform_body " + f"{'overridden' if self._routing_transforms_body else 'not overridden'} " + f"but mutates_body={self._routing.mutates_body}. A plugin " + "must declare mutates_body=True if and only if it " + "implements transform_body." + ) # Detect and set transport type if not explicitly set if not model_endpoint.transport: @@ -102,6 +139,28 @@ def __init__( 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 path on ANY terminal outcome + (final turn, cancellation, terminal context overflow, cancel-before- + start). 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, @@ -126,11 +185,49 @@ 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. + # The capability flags skip the no-op base headers()/transform_body() + # calls entirely for plugins that don't override them. + routing_ctx: RoutingContext | None = None + if self._routing_stamps_headers or self._routing_transforms_body: + 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, + ) + if self._routing_stamps_headers: + # 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) + if request_info.payload_bytes is not None: # PAYLOAD_BYTES fast path: bytes were validated at dataset-load time - # by the mmap loader / DatasetManager, and body-mutating features - # (cache-bust, Dynamo session_control) are refused against this - # verbatim-bytes path at dataset load, so nothing is injected here. + # by the mmap loader / DatasetManager. A body-mutating routing plugin + # cannot rewrite opaque bytes without a reparse/redump that defeats + # the fast path, so it is refused here (the raise is converted to an + # error record by _send_request_internal). Header-based routing is + # compatible and was already applied above. + if self._routing is not None and self._routing.mutates_body: + raise ValueError( + f"session-routing mode " + f"{self.model_endpoint.endpoint.session_routing!r} mutates " + "request bodies and is incompatible with the verbatim PAYLOAD_BYTES " + "fast path; choose a headers-based mode or a structured-turn dataset." + ) formatted_payload = request_info.payload_bytes else: current_turn = request_info.turns[-1] if request_info.turns else None @@ -138,32 +235,23 @@ async def _send_request_to_transport( formatted_payload = current_turn.raw_payload else: formatted_payload = self.endpoint.format_payload(request_info) - # Dynamo conversation-aware routing (opt-in): overlay - # nvext.session_control onto the structured request body. Done here, - # after the endpoint built the dict, so it is endpoint-agnostic and - # never mutates a cached Turn. The verbatim PAYLOAD_BYTES path is - # excluded by the dataset-load guard, so it is not handled here. - endpoint = self.model_endpoint.endpoint - if endpoint.use_dynamo_conv_aware_routing: - session_id = request_info.x_correlation_id - legacy = endpoint.use_legacy_dynamo_session_control - session_control = build_session_control( - session_id=session_id, - is_final_turn=request_info.is_final_turn, - timeout_seconds=endpoint.dynamo_session_timeout_seconds, - legacy=legacy, - already_opened=session_id in self._dynamo_opened_sessions, - ) - # Track the open/close lifecycle so legacy 'open' is sent exactly - # once per session (modern 'bind' is stateless and ignores this). - if legacy: - if session_control.get("action") == "open": - self._dynamo_opened_sessions.add(session_id) - elif request_info.is_final_turn: - self._dynamo_opened_sessions.discard(session_id) - formatted_payload = merge_session_control( - formatted_payload, session_control - ) + # Body-based session routing (e.g. Dynamo nvext.session_control): + # overlay onto the structured body after the endpoint built the dict, + # so it is endpoint-agnostic and never mutates a cached Turn + # (transform_body returns a copy). + if ( + routing_ctx is not None + and self._routing_transforms_body + and isinstance(formatted_payload, dict) + ): + try: + formatted_payload = self._routing.transform_body( + formatted_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 # Canonicalise to bytes and stash on request_info. Two wins: (1) the # transport skips its own orjson.dumps on the dict path, (2) the # record processor can drop request_info.turns before the ZMQ hop diff --git a/src/aiperf/workers/session_manager.py b/src/aiperf/workers/session_manager.py index f80befcb36..de502ba826 100644 --- a/src/aiperf/workers/session_manager.py +++ b/src/aiperf/workers/session_manager.py @@ -314,6 +314,15 @@ def get(self, x_correlation_id: str) -> UserSession | None: """ return self._cache.get(x_correlation_id) + def live_session_ids(self) -> list[str]: + """Correlation IDs of every session still resident in the cache. + + Snapshot (not a view) so callers may evict/notify while iterating. + Used by the worker's shutdown sweep to fire the routing plugin's + post-session hook for sessions abandoned mid-conversation at run end. + """ + return list(self._cache.keys()) + def evict(self, x_correlation_id: str) -> None: """ Evict user session. diff --git a/src/aiperf/workers/session_routing.py b/src/aiperf/workers/session_routing.py new file mode 100644 index 0000000000..c20748ecb0 --- /dev/null +++ b/src/aiperf/workers/session_routing.py @@ -0,0 +1,327 @@ +# 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 tiered identity +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, terminal + context overflow, cancel-before-start). 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 + +import re +from abc import ABC +from typing import Annotated, Any, ClassVar, Generic, TypeVar + +from msgspec import Struct +from pydantic import BeforeValidator, ConfigDict, Field, model_validator + +from aiperf.common.models import AIPerfBaseModel + +# RFC 9110 token: the only characters legal in an HTTP field name. Rejecting +# anything else at config time turns a guaranteed runtime wire failure (or +# header-injection vector via CR/LF) into an immediate, actionable error. +_HEADER_TOKEN_RE = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") + + +class EmptyRoutingOptions(AIPerfBaseModel): + """Options model for parameterless plugins; rejects every opt key.""" + + model_config = ConfigDict(extra="forbid") + + +OptionsT = TypeVar("OptionsT", bound=AIPerfBaseModel) + + +class RoutingContext(Struct, frozen=True, gc=False): + """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. + + msgspec.Struct (not a frozen dataclass): constructed once per request on + the worker hot path, and frozen-dataclass __init__ goes through + object.__setattr__ per field (~4x slower to create, measured). gc=False + is safe: every field is a str/bool/None, so no reference cycles. + """ + + 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 + """True when ``transform_body`` changes the payload. Gates the plugin off + the verbatim PAYLOAD_BYTES mmap fast path at dataset build, cache hit, + and runtime. CONTRACT: a plugin overrides ``transform_body`` if and only + if it declares ``mutates_body = True`` -- the two must agree, and + ``InferenceClient`` enforces this at worker init (fail-fast) because a + transform without the declaration would be silently skipped on + PAYLOAD_BYTES datasets, and a declaration without a transform would + force runs off the fast path for nothing.""" + + # 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 __init__(self, options: DynamoNvextOptions) -> None: + super().__init__(options) + # Hot path: one attribute hop per request instead of two through the + # pydantic options model. + self._timeout = options.timeout_seconds + + 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._timeout, + } + merged = dict(payload) + raw_nvext = merged.get("nvext") + if raw_nvext is None: + # Common case: dataset payloads carry no nvext. Skip the two + # defensive dict copies below (~40% faster, measured). + merged["nvext"] = {"session_control": session_control} + return merged + nvext = dict(raw_nvext) if isinstance(raw_nvext, dict) else {} + raw_sc = nvext.get("session_control") + if isinstance(raw_sc, dict): + session_control = {**raw_sc, **session_control} + nvext["session_control"] = session_control + 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 SessionIdHeaderRouting(SessionRoutingBase[EmptyRoutingOptions]): + """Preset: additive X-Session-ID header carrying the session's correlation ID. + + Equivalent to ``identity_headers`` with ``session=X-Session-ID``; use + ``identity_headers`` for custom names or parent/root tiers. + """ + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + return {"X-Session-ID": ctx.x_correlation_id} + + +def _split_header_names(value: Any) -> Any: + """Comma-split a raw string opt value into a list of stripped names. + + HTTP header names are RFC 9110 tokens and can never contain a comma, so + the split is unambiguous. List inputs (canonicalized opts, YAML lists) + keep their item boundaries but are stripped the same way, so the + validator's non-empty/duplicate/token checks see identical canonical + names regardless of input form. + """ + if isinstance(value, str): + return [item.strip() for item in value.split(",")] + if isinstance(value, list): + return [item.strip() if isinstance(item, str) else item for item in value] + return value + + +_HeaderNames = Annotated[list[str], BeforeValidator(_split_header_names)] + + +class IdentityHeadersOptions(AIPerfBaseModel): + """Options for the fully generic tiered identity headers. + + Every tier defaults EMPTY: the plugin emits exactly the headers configured + and nothing else (unlike the presets, which exist to have opinions). + """ + + model_config = ConfigDict(extra="forbid") + + session: _HeaderNames = Field( + default_factory=list, + description="Header name(s) carrying this session's correlation ID " + "(same value on every turn). Comma-separate for multiple names.", + ) + parent: _HeaderNames = Field( + default_factory=list, + description="Header name(s) carrying the immediate parent session's " + "correlation ID; omitted on requests whose session has no parent. " + "Comma-separate for multiple names.", + ) + root: _HeaderNames = Field( + default_factory=list, + description="Header name(s) carrying the session-tree root's " + "correlation ID (the session's own ID for root sessions), keying " + "whole-tree affinity. Comma-separate for multiple names.", + ) + + @model_validator(mode="after") + def validate_header_names(self) -> IdentityHeadersOptions: + """Require >=1 valid RFC 9110 token overall; reject duplicates. + + Uniqueness is case-insensitive over ALL configured names combined + (within a tier and across tiers): a repeated name means one value + silently overwriting another -- the exact affinity corruption this + feature exists to prevent. + """ + all_names = [*self.session, *self.parent, *self.root] + if not all_names: + raise ValueError( + "identity_headers requires at least one header name across " + "the session/parent/root tiers." + ) + for name in all_names: + if not name: + raise ValueError("header names must be non-empty") + if not _HEADER_TOKEN_RE.match(name): + raise ValueError( + f"invalid header name {name!r}: HTTP field names are " + "RFC 9110 tokens (letters, digits, and !#$%&'*+-.^_`|~; " + "no spaces, colons, or control characters)" + ) + lowered = [name.lower() for name in all_names] + if len(set(lowered)) != len(lowered): + raise ValueError( + f"duplicate header name (names must be unique across all " + f"tiers combined) in session={self.session!r} " + f"parent={self.parent!r} root={self.root!r}" + ) + return self + + +class IdentityHeadersRouting(SessionRoutingBase[IdentityHeadersOptions]): + """Fully generic tiered identity headers: any name(s) per identity tier. + + Stamps each configured header with its tier's correlation ID: ``session`` + (this session), ``parent`` (immediate parent; omitted for roots), and + ``root`` (session-tree root -- whole-tree affinity, e.g. pinning an entire + agent tree to one replica for prefix-cache locality). The named presets + are expressible in terms of this mode: ``session_id_header`` is + ``session=X-Session-ID``, ``smg_routing_key`` is + ``session=X-SMG-Routing-Key``, and ``dynamo_headers`` is + ``session=X-Dynamo-Session-ID`` + ``parent=X-Dynamo-Parent-Session-ID``. + """ + + Options: ClassVar[type[AIPerfBaseModel]] = IdentityHeadersOptions + + def __init__(self, options: IdentityHeadersOptions) -> None: + super().__init__(options) + # Hot path: immutable snapshots read once per request, skipping the + # two-hop pydantic options attribute walk (~17% faster on the common + # single-tier shape, measured). + self._session_names = tuple(options.session) + self._parent_names = tuple(options.parent) + self._root_names = tuple(options.root) + + def headers(self, ctx: RoutingContext) -> dict[str, str]: + xcorr = ctx.x_correlation_id + headers = {name: xcorr for name in self._session_names} + parent = ctx.parent_correlation_id + if parent and self._parent_names: + for name in self._parent_names: + headers[name] = parent + root = ctx.root_correlation_id + if root and self._root_names: + for name in self._root_names: + headers[name] = root + return headers diff --git a/src/aiperf/workers/worker.py b/src/aiperf/workers/worker.py index 27184f6948..d98c52c6de 100644 --- a/src/aiperf/workers/worker.py +++ b/src/aiperf/workers/worker.py @@ -63,6 +63,7 @@ RequestClientProtocol, StreamingDealerClientProtocol, ) +from aiperf.common.scenario import is_context_overflow_response from aiperf.credit.messages import ( CancelCredits, CreditReturn, @@ -74,7 +75,7 @@ from aiperf.credit.structs import Credit, CreditContext from aiperf.dataset.protocols import DatasetClientStoreProtocol from aiperf.plugin import plugins -from aiperf.plugin.enums import PluginType +from aiperf.plugin.enums import PluginType, TimingMode from aiperf.records.payload_retention import resolve_strip_record_payload_bytes from aiperf.workers.inference_client import InferenceClient from aiperf.workers.session_manager import UserSession, UserSessionManager @@ -86,6 +87,23 @@ _logger = AIPerfLogger(__name__) +def _is_terminal_context_overflow( + credit: Credit, credit_context: CreditContext +) -> bool: + """Whether this credit return ends the conversation via context overflow. + + A context-overflow error on a non-final, non-cancelled turn is terminal: + agentic_replay recycles the lane and issues no further (final/cancel) credit + for the session, so the worker must treat it as a terminal disposition. + Mirrors the strategy's own ``terminal_overflow`` classification, applied to + the same error body the worker returns on the credit. + """ + if credit.is_final_turn or credit_context.cancelled: + return False + error = credit_context.error + return error is not None and is_context_overflow_response(body=str(error)) + + def _apply_cache_bust_to_system_message( system_message: str | None, marker: str, target: CacheBustTarget ) -> str | None: @@ -505,6 +523,15 @@ def __init__( self.model_endpoint = ModelEndpointInfo.from_user_config(self.user_config) + # A mid-stream context overflow ends the conversation ONLY under + # agentic_replay (the strategy recycles the lane and issues no further + # credit for the session). Other strategies keep issuing the session's + # remaining turns after an overflow error, so treating it as terminal + # there would fire the routing plugin's post-session hook mid-session. + self._overflow_ends_session = ( + self.user_config.timing_mode == TimingMode.AGENTIC_REPLAY + ) + self.inference_client: InferenceClient = InferenceClient( model_endpoint=self.model_endpoint, service_id=self.service_id, @@ -712,6 +739,18 @@ def _on_credit_drop_message_task_done( self.execute_async(self.credit_dealer_client.send(credit_return)) credit_context.returned = True + # Post-session hook for routing plugins: the cancel-before-start path + # is a terminal disposition the finally block in _process_credit never + # sees (the task died before it could run), so notify here too -- + # gated on cancelled because the OTHER way to reach this branch (the + # finally block's CreditReturn send raised) is not terminal for a + # non-final turn, and for a final/cancelled turn the finally block + # already notified. Idempotent by contract. + if credit_context.cancelled and 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 @@ -905,7 +944,22 @@ async def _process_credit(self, credit_context: CreditContext) -> None: self.exception(f"Error processing credit: {e!r}") finally: if credit_context.credit.is_final_turn or credit_context.cancelled: + # 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) self.session_manager.evict(x_correlation_id) + elif self._overflow_ends_session and _is_terminal_context_overflow( + credit_context.credit, credit_context + ): + # Under agentic_replay a mid-stream context overflow ends the + # conversation without a final/cancel credit (the lane is + # recycled), so the routing plugin's post-session hook must + # fire here explicitly. Other strategies keep issuing the + # session's turns after an overflow, so this is gated off. + self.inference_client.notify_session_end(x_correlation_id) def _make_first_token_callback( self, credit_context: CreditContext @@ -1111,6 +1165,8 @@ def _create_request_info( 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, cache_bust_marker=credit.cache_bust_marker, cache_bust_target=credit.cache_bust_target if credit.cache_bust_marker is not None @@ -1300,6 +1356,15 @@ async def _on_profile_configure_command(self, message: CommandMessage) -> None: @on_stop async def _worker_stop(self) -> None: + # Sessions still resident at shutdown (mid-conversation when the phase + # ended: sitting in think-time delay or a strategy continuation queue, + # so no in-flight credit existed for cancel_all_credits to cancel) + # never hit any per-credit terminal path. Fire the routing plugin's + # post-session hook for each so stateful plugins release them; + # idempotent by contract and a no-op when routing is unset. + for x_correlation_id in self.session_manager.live_session_ids(): + self.inference_client.notify_session_end(x_correlation_id) + # Clean up dataset client resources using protocol lifecycle if self._dataset_client is not None: dataset_client = self._dataset_client diff --git a/tests/integration/test_dynamo_session_control_raw_export.py b/tests/integration/test_dynamo_session_control_raw_export.py deleted file mode 100644 index 385a12d36b..0000000000 --- a/tests/integration/test_dynamo_session_control_raw_export.py +++ /dev/null @@ -1,143 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""End-to-end raw-export checks of the Dynamo nvext.session_control lifecycle. - -Runs a real ``aiperf profile`` subprocess against the in-repo mock server with -``--use-dynamo-conv-aware-routing`` and ``--export-level raw``, then reads the -exported wire payloads and asserts the per-session session_control lifecycle: - -- modern (default): ``bind`` on every non-final turn, ``close`` on the last. -- legacy (``--use-legacy-dynamo-session-control``): ``open`` on the first turn, - ``session_id`` only on intermediate turns, ``close`` on the last -- and never - ``bind`` (which released Dynamo v1.2.x rejects with HTTP 400). -""" - -from __future__ import annotations - -import json -from collections import defaultdict - -import pytest - -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 -_TIMEOUT_SECONDS = 300 - - -def _payload_dict(record) -> dict: - if record.payload is not None: - return record.payload - if record.payload_bytes is not None: - return json.loads(record.payload_bytes) - return {} - - -def _build_cmd(url: str, *, legacy: bool) -> str: - legacy_flag = "--use-legacy-dynamo-session-control" if legacy else "" - 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 \ - --use-dynamo-conv-aware-routing \ - {legacy_flag} \ - --dynamo-session-timeout-seconds {_TIMEOUT_SECONDS} \ - --export-level raw \ - --ui simple - """ - - -async def _session_controls_by_session( - cli: AIPerfCLI, url: str, *, legacy: bool -) -> dict[str, list[dict]]: - """Run a benchmark and return each session's session_control blocks, ordered - by turn_index, keyed by X-Correlation-ID.""" - result = await cli.run(_build_cmd(url, legacy=legacy), timeout=300.0) - - records = list(result.raw_records or []) - assert records, f"no raw records\n{(result.log or '')[-1500:]}" - - grouped: dict[str, list[tuple[int, dict]]] = defaultdict(list) - for rec in records: - sc = _payload_dict(rec).get("nvext", {}).get("session_control") - assert sc is not None, "every request must carry nvext.session_control" - grouped[rec.metadata.x_correlation_id].append((rec.metadata.turn_index, sc)) - - assert len(grouped) == _NUM_SESSIONS, ( - f"expected {_NUM_SESSIONS} sessions, got {len(grouped)}" - ) - out: dict[str, list[dict]] = {} - for xcorr, turns in grouped.items(): - assert len(turns) == _TURNS_PER_SESSION, ( - f"session {xcorr}: expected {_TURNS_PER_SESSION} turns, got {len(turns)}" - ) - turns.sort(key=lambda t: t[0]) - out[xcorr] = [sc for _, sc in turns] - return out - - -@pytest.mark.integration -@pytest.mark.asyncio -async def test_modern_session_control_bind_close_lifecycle_on_wire( - cli: AIPerfCLI, - aiperf_mock_server: AIPerfMockServer, -): - """Default mode: every non-final turn re-binds (with timeout), the final - turn closes, every turn shares one session_id, and 'open' never appears.""" - by_session = await _session_controls_by_session( - cli, aiperf_mock_server.url, legacy=False - ) - - for xcorr, scs in by_session.items(): - 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 - assert "timeout" not in scs[-1] - - # One stable session_id == the X-Correlation-ID, on every turn. - assert {sc["session_id"] for sc in scs} == {xcorr} - - -@pytest.mark.integration -@pytest.mark.asyncio -async def test_legacy_session_control_open_close_lifecycle_on_wire( - cli: AIPerfCLI, - aiperf_mock_server: AIPerfMockServer, -): - """Legacy mode: open on the first request, session_id only on intermediate - turns, close on the final turn -- never bind.""" - by_session = await _session_controls_by_session( - cli, aiperf_mock_server.url, legacy=True - ) - - for xcorr, scs in by_session.items(): - actions = [sc.get("action") for sc in scs] - assert actions[0] == "open", f"session {xcorr}: actions={actions}" - assert actions[-1] == "close", f"session {xcorr}: actions={actions}" - assert all(a is None for a in actions[1:-1]), ( - f"session {xcorr}: intermediate turns must carry no action; {actions}" - ) - assert "bind" not in actions, f"session {xcorr}: emitted bind; {actions}" - - # One stable session_id == the X-Correlation-ID, on every turn. - assert {sc["session_id"] for sc in scs} == {xcorr} - - # 'open' carries the timeout; 'close' / bare turns do not. - assert scs[0]["timeout"] == _TIMEOUT_SECONDS - assert "timeout" not in scs[-1] 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..e422375c39 --- /dev/null +++ b/tests/integration/test_session_routing_raw_export.py @@ -0,0 +1,224 @@ +# 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``: the preset's ``X-Session-ID`` equals ``x_correlation_id``. +- ``identity_headers`` (``session=X-Affinity,X-SMG-Routing-Key`` + + ``root=X-Tree-ID``): every configured header equals ``x_correlation_id`` + (flat sessions are their own tree roots). +""" + +from __future__ import annotations + +import json +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 _payload_dict(record: RawRecordInfo) -> dict: + """Decode the exported wire payload regardless of which field carries it.""" + if record.payload is not None: + return record.payload + if record.payload_bytes is not None: + return json.loads(record.payload_bytes) + return {} + + +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 _payload_dict(rec), ( + 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 _payload_dict(rec), ( + 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-Session-ID") == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: headers={headers}" + ) + assert "nvext" not in _payload_dict(rec), ( + f"session {xcorr}: header-mode must not mutate the body" + ) + + +def _verify_identity_headers(by_session: dict[str, list[RawRecordInfo]]) -> None: + """Two session-tier names plus a root-tier name; flat sessions are their + own tree roots, so every header carries the session's X-Correlation-ID.""" + for xcorr, recs in by_session.items(): + for rec in recs: + headers = rec.request_headers or {} + for name in ("X-Affinity", "X-SMG-Routing-Key", "X-Tree-ID"): + assert headers.get(name) == xcorr, ( + f"session {xcorr} turn {rec.metadata.turn_index}: " + f"missing/mismatched {name}; headers={headers}" + ) + assert "nvext" not in _payload_dict(rec), ( + 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 = _payload_dict(rec).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, + "identity_headers": _verify_identity_headers, +} + + +@pytest.mark.integration +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode, opts, verifier", + [ + param("dynamo_headers", (), "dynamo_headers", id="dynamo_headers"), + param("dynamo_nvext", (f"timeout_seconds={_TIMEOUT_SECONDS}",), "dynamo_nvext", id="dynamo_nvext"), + param("smg_routing_key", (), "smg_routing_key", id="smg_routing_key"), + param("session_id_header", (), "session_id_header", id="session_id_header"), + param( + "identity_headers", + ("session=X-Affinity,X-SMG-Routing-Key", "root=X-Tree-ID"), + "identity_headers", + id="identity_headers", + ), + ], +) # fmt: skip +async def test_session_routing_mode_reaches_wire( + cli: AIPerfCLI, + aiperf_mock_server: AIPerfMockServer, + mode: str, + opts: tuple[str, ...], + verifier: 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[verifier](by_session) diff --git a/tests/unit/common/config/test_endpoint_config.py b/tests/unit/common/config/test_endpoint_config.py index 25aac41266..0a2c709455 100644 --- a/tests/unit/common/config/test_endpoint_config.py +++ b/tests/unit/common/config/test_endpoint_config.py @@ -26,18 +26,9 @@ def test_endpoint_config_defaults(): assert config.custom_endpoint == EndpointDefaults.CUSTOM_ENDPOINT assert config.streaming == EndpointDefaults.STREAMING assert config.url == EndpointDefaults.URL - assert ( - config.use_dynamo_conv_aware_routing - == EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING - ) - assert ( - config.use_legacy_dynamo_session_control - == EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL - ) - assert ( - config.dynamo_session_timeout_seconds - == EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS - ) + assert config.session_routing is None + assert config.session_routing_opt == [] + assert config.session_routing_opts == {} def test_endpoint_config_custom_values(): @@ -61,9 +52,6 @@ def test_endpoint_config_custom_values(): "urls": ["http://custom-url"], "timeout_seconds": 10, "api_key": "custom_api_key", - "use_dynamo_conv_aware_routing": True, - "use_legacy_dynamo_session_control": True, - "dynamo_session_timeout_seconds": 123, } config = EndpointConfig(**custom_values) for key, value in custom_values.items(): @@ -226,29 +214,55 @@ def test_invalid_mode_rejected(self): ) -class TestDynamoSessionControlValidation: - """Coherence between the conv-aware-routing flag and its legacy modifier.""" +class TestSessionRoutingValidation: + """Fail-fast and canonicalization of --session-routing / --session-routing-opt.""" - def test_legacy_without_master_flag_rejected(self): - """The legacy modifier has no meaning unless routing is enabled.""" - with pytest.raises(ValueError, match="use-legacy-dynamo-session-control"): + def test_opts_without_mode_rejected(self): + """Opts have no meaning unless a routing mode is selected.""" + with pytest.raises(ValueError, match="--session-routing-opt requires"): EndpointConfig( model_names=["gpt2"], - use_legacy_dynamo_session_control=True, + session_routing_opt=["timeout_seconds=600"], ) - def test_legacy_with_master_flag_allowed(self): + def test_mode_with_opts_canonicalized(self): + """Opt values are coerced to the plugin Options model types.""" config = EndpointConfig( model_names=["gpt2"], - use_dynamo_conv_aware_routing=True, - use_legacy_dynamo_session_control=True, + session_routing="dynamo_nvext", + session_routing_opt=["timeout_seconds=600"], ) - assert config.use_legacy_dynamo_session_control + assert config.session_routing_opts == {"timeout_seconds": 600} - def test_master_flag_alone_allowed(self): - """Conv-aware routing without the legacy modifier is the modern default.""" + def test_mode_alone_allowed(self): + """A parameterless mode needs no opts.""" config = EndpointConfig( model_names=["gpt2"], - use_dynamo_conv_aware_routing=True, + session_routing="dynamo_headers", ) - assert not config.use_legacy_dynamo_session_control + assert config.session_routing_opts == {} + + def test_unknown_opt_key_rejected(self): + """Unknown keys are rejected by the plugin's extra='forbid' Options model.""" + with pytest.raises(ValueError): + EndpointConfig( + model_names=["gpt2"], + session_routing="dynamo_headers", + session_routing_opt=["bogus=1"], + ) + + def test_malformed_opt_rejected(self): + with pytest.raises(ValueError, match="expected non-empty key=value"): + EndpointConfig( + model_names=["gpt2"], + session_routing="dynamo_nvext", + session_routing_opt=["timeout_seconds"], + ) + + def test_duplicate_opt_key_rejected(self): + with pytest.raises(ValueError, match="Duplicate --session-routing-opt"): + EndpointConfig( + model_names=["gpt2"], + session_routing="dynamo_nvext", + session_routing_opt=["timeout_seconds=1", "timeout_seconds=2"], + ) diff --git a/tests/unit/common/models/test_endpoint_info.py b/tests/unit/common/models/test_endpoint_info.py index 30a723f74d..5a99c59961 100644 --- a/tests/unit/common/models/test_endpoint_info.py +++ b/tests/unit/common/models/test_endpoint_info.py @@ -16,14 +16,8 @@ def test_single_url_default(self): info = EndpointInfo() assert info.base_urls == [EndpointDefaults.URL] assert info.base_url == EndpointDefaults.URL - assert ( - info.use_dynamo_conv_aware_routing - == EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING - ) - assert ( - info.dynamo_session_timeout_seconds - == EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS - ) + assert info.session_routing is None + assert info.session_routing_opts == {} def test_single_url_custom(self): """Custom single URL should work.""" @@ -43,18 +37,19 @@ def test_base_urls_must_have_at_least_one(self): with pytest.raises(ValueError): EndpointInfo(base_urls=[]) - def test_dynamo_session_control_from_user_config(self): - """Dynamo session-control fields should flow into runtime endpoint info.""" + def test_session_routing_from_user_config(self): + """Session-routing fields should flow into runtime endpoint info, + with opts canonicalized to the plugin's Options model types.""" user_config = UserConfig( endpoint=EndpointConfig( model_names=["test-model"], - use_dynamo_conv_aware_routing=True, - dynamo_session_timeout_seconds=123, + session_routing="dynamo_nvext", + session_routing_opt=["timeout_seconds=123"], ) ) info = ModelEndpointInfo.from_user_config(user_config).endpoint - assert info.use_dynamo_conv_aware_routing is True - assert info.dynamo_session_timeout_seconds == 123 + assert info.session_routing == "dynamo_nvext" + assert info.session_routing_opts == {"timeout_seconds": 123} class TestEndpointInfoGetUrl: 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..89e978b661 --- /dev/null +++ b/tests/unit/config/test_session_routing_config.py @@ -0,0 +1,121 @@ +# 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.common.config import EndpointConfig +from aiperf.common.config.endpoint_config import _parse_session_routing_opts + + +def _config(**kwargs) -> EndpointConfig: + return EndpointConfig(model_names=["test-model"], **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_opt=["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_opt=["timeout_seconds=600"], + ) + assert config.session_routing_opts == {"timeout_seconds": 600} + assert isinstance(config.session_routing_opts["timeout_seconds"], int) + + +def test_canonicalization_from_config_file_dict(): + """Opts set directly (config-file form) are canonicalized the same way.""" + config = _config( + session_routing="dynamo_nvext", + session_routing_opts={"timeout_seconds": 600}, + ) + assert config.session_routing_opts == {"timeout_seconds": 600} + + +def test_identity_headers_tiers_canonicalized_to_lists(): + """Comma-separated tier values canonicalize to typed lists, so the + pickled UserConfig that reaches workers re-validates without re-splitting.""" + config = _config( + session_routing="identity_headers", + session_routing_opt=[ + "session=X-Session-ID,X-SMG-Routing-Key", + "root=X-Tree-ID", + ], + ) + assert config.session_routing_opts == { + "session": ["X-Session-ID", "X-SMG-Routing-Key"], + "root": ["X-Tree-ID"], + } + + +def test_identity_headers_no_names_rejected(): + with pytest.raises(ValueError, match="at least one header name"): + _config(session_routing="identity_headers") + + +def test_session_id_header_preset_rejects_opts(): + """The preset lost its header_name option when identity_headers took over + customization; it must reject what it used to accept.""" + with pytest.raises(ValueError): + _config( + session_routing="session_id_header", + session_routing_opt=["header_name=X-Affinity"], + ) + + +def test_opts_without_mode_rejected(): + with pytest.raises(ValueError, match="session-routing-opt"): + _config(session_routing_opt=["session=X-A"]) + + +def test_unknown_opt_key_rejected(): + with pytest.raises(ValueError, match="timeout_secs"): + _config( + session_routing="dynamo_nvext", + session_routing_opt=["timeout_secs=600"], + ) + + +def test_invalid_opt_value_rejected(): + with pytest.raises(ValueError): + _config( + session_routing="dynamo_nvext", + session_routing_opt=["timeout_seconds=0"], + ) + + +def test_parameterless_mode_rejects_any_opt(): + with pytest.raises(ValueError): + _config( + session_routing="smg_routing_key", + session_routing_opt=["anything=x"], + ) + + +def test_parse_routing_opts_duplicate_key_rejected(): + with pytest.raises(ValueError, match="Duplicate"): + _parse_session_routing_opts(["session=X-A", "session=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_session_routing_opts([item]) diff --git a/tests/unit/credit/test_credit_issuer_join_adversarial.py b/tests/unit/credit/test_credit_issuer_join_adversarial.py index f7b5eb2413..2f3ef738a8 100644 --- a/tests/unit/credit/test_credit_issuer_join_adversarial.py +++ b/tests/unit/credit/test_credit_issuer_join_adversarial.py @@ -328,3 +328,33 @@ async def fake_try_issue_credit(turn: TurnToSend): ) await issuer.dispatch_join_turn(pending) assert captured["turn"].has_forks is False + + +@pytest.mark.asyncio +async def test_dispatch_join_turn_stamps_has_branches_from_pending(): + """The join seam must carry the gated turn's any-mode branch flag onto the + resumed TurnToSend, or a SPAWN-declaring join turn could stamp a wrong + is_tree_final=True while its children are pending.""" + issuer = _make_issuer() + issuer.try_issue_credit = AsyncMock(return_value=True) + pending = PendingBranchJoin( + parent_x_correlation_id="corr-parent", + parent_conversation_id="conv-parent", + parent_num_turns=3, + gated_turn_index=2, + parent_has_branches_on_gated_turn=True, + ) + await issuer.dispatch_join_turn(pending) + + turn = issuer.try_issue_credit.await_args.args[0] + assert turn.has_branches is True + + issuer.try_issue_credit.reset_mock() + pending_leaf = PendingBranchJoin( + parent_x_correlation_id="corr-parent", + parent_conversation_id="conv-parent", + parent_num_turns=3, + gated_turn_index=2, + ) + await issuer.dispatch_join_turn(pending_leaf) + assert issuer.try_issue_credit.await_args.args[0].has_branches is False diff --git a/tests/unit/credit/test_issuer_finality.py b/tests/unit/credit/test_issuer_finality.py new file mode 100644 index 0000000000..dedff8e5e2 --- /dev/null +++ b/tests/unit/credit/test_issuer_finality.py @@ -0,0 +1,297 @@ +# 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 (three-touch touch #2). + +Touch #1 (the ``Credit`` struct fields) and touch #3 (worker -> ``RequestInfo``, +``test_create_request_info_plumbs_finality_from_credit``) are already covered. +This file covers touch #2: ``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. + +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 + +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 real ``SessionTreeRegistry`` only needs ``release_session_slot``; 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(_FakeConcurrency()) + + +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-agentic 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_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 + + +def test_build_first_turn_stamps_has_branches_and_gates_finality(): + """End-to-end seam guard: a root whose turn-0 declares a SPAWN branch must + build a ``TurnToSend`` with ``has_branches=True`` / ``has_forks=False`` + (SPAWN does not set the FORK-only flag) and, fed through the real issuer's + ``_finality_for_issue`` against an open tree with nothing outstanding, stamp + ``is_tree_final=False``. Catches a dropped ``has_branches`` stamp at the + ``build_first_turn`` construction seam.""" + from aiperf.common.enums import ConversationBranchMode + from aiperf.common.models import ( + ConversationBranchInfo, + ConversationMetadata, + TurnMetadata, + ) + from aiperf.timing.conversation_source import SampledSession + + meta = ConversationMetadata( + conversation_id="conv-1", + turns=[TurnMetadata(timestamp_ms=0.0, branch_ids=["conv-1:0"])], + branches=[ + ConversationBranchInfo( + branch_id="conv-1:0", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + ) + ], + ) + session = SampledSession( + conversation_id="conv-1", metadata=meta, x_correlation_id="root-1" + ) + turn = session.build_first_turn() + assert turn.has_branches is True + assert turn.has_forks is False + + registry = _make_registry() + registry.open_tree("root-1", CreditPhase.PROFILING, root_pending=True) + issuer, _ = _make_issuer(registry) + + assert issuer._finality_for_issue(turn) == (None, False) + + +# ============================================================================= +# GUARD: the Credit(...) construction site must pass the helper's results through +# ============================================================================= + + +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) + + await issuer.issue_credit(_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 + + +def test_build_turn_at_index_stamps_has_branches(): + """The mid-trace resume seam (agentic warmup k_i > 0) must stamp the + any-mode branch flag from the indexed turn's metadata.""" + from aiperf.common.enums import ConversationBranchMode + from aiperf.common.models import ( + ConversationBranchInfo, + ConversationMetadata, + TurnMetadata, + ) + from aiperf.timing.conversation_source import SampledSession + + meta = ConversationMetadata( + conversation_id="conv-1", + turns=[ + TurnMetadata(timestamp_ms=0.0), + TurnMetadata(timestamp_ms=1.0, branch_ids=["conv-1:1"]), + ], + branches=[ + ConversationBranchInfo( + branch_id="conv-1:1", + child_conversation_ids=["child-conv"], + mode=ConversationBranchMode.SPAWN, + ) + ], + ) + session = SampledSession( + conversation_id="conv-1", metadata=meta, x_correlation_id="root-1" + ) + assert session.build_turn_at_index(1).has_branches is True + assert session.build_turn_at_index(0).has_branches is False diff --git a/tests/unit/credit/test_structs.py b/tests/unit/credit/test_structs.py index 176ead47d2..ffd97514ec 100644 --- a/tests/unit/credit/test_structs.py +++ b/tests/unit/credit/test_structs.py @@ -75,3 +75,24 @@ def test_turn_to_send_does_not_propagate_max_tokens_override(): parent = _make_credit(max_tokens_override=1) next_turn = TurnToSend.from_previous_credit(parent) assert next_turn.max_tokens_override is None + + +def test_turn_to_send_from_previous_credit_stamps_has_branches_from_next_meta(): + """The mainline multi-turn continuation seam must stamp the any-mode + branch flag from the NEW turn's metadata, or a SPAWN-declaring final + turn reads as tree-final while its children are still pending.""" + from aiperf.common.models.dataset_models import TurnMetadata + + parent = _make_credit() + branching = TurnToSend.from_previous_credit( + parent, next_meta=TurnMetadata(timestamp_ms=0.0, branch_ids=["c:0"]) + ) + assert branching.has_branches is True + + leaf = TurnToSend.from_previous_credit( + parent, next_meta=TurnMetadata(timestamp_ms=0.0) + ) + assert leaf.has_branches is False + + no_meta = TurnToSend.from_previous_credit(parent) + assert no_meta.has_branches is False diff --git a/tests/unit/dataset/test_dataset_manager.py b/tests/unit/dataset/test_dataset_manager.py index 3bf7c4e440..632f3549d0 100644 --- a/tests/unit/dataset/test_dataset_manager.py +++ b/tests/unit/dataset/test_dataset_manager.py @@ -11,7 +11,11 @@ from aiperf.common.config.config_defaults import InputDefaults from aiperf.common.config.conversation_config import ConversationConfig, TurnConfig from aiperf.common.config.tokenizer_config import TokenizerConfig -from aiperf.common.enums import ConversationContextMode +from aiperf.common.enums import ( + CacheBustTarget, + ConversationContextMode, + MemoryMapFormat, +) from aiperf.common.exceptions import ServiceError from aiperf.common.messages import ( ConversationRequestMessage, @@ -20,6 +24,7 @@ ) from aiperf.common.messages.command_messages import ProfileConfigureCommand from aiperf.common.models import Conversation, Image, Text, Turn +from aiperf.dataset import mmap_cache from aiperf.dataset.dataset_manager import DatasetManager from aiperf.plugin.enums import ( CustomDatasetType, @@ -1219,6 +1224,57 @@ def test_preformat_skipped_when_cache_bust_enabled( assert conversations[0].turns[0].raw_payload is None assert conversations[1].turns[0].raw_payload is None + def test_preformat_skipped_when_mutating_routing_enabled( + self, initialized_dataset_manager + ): + """A body-mutating session-routing mode must keep structured datasets on + the structured-turns path. Without this bail, the preformatter promotes + the default synthetic/single-turn dataset to PAYLOAD_BYTES and + _select_mmap_format then hard-fails a perfectly valid run.""" + initialized_dataset_manager.user_config.endpoint.session_routing = ( + "dynamo_nvext" + ) + + conversations = [ + Conversation( + session_id="s1", + turns=[Turn(role="user", texts=[Text(contents=["hi"])])], + ), + ] + + with patch( + "aiperf.dataset.dataset_manager.format_conversation_payloads" + ) as mock_fmt: + initialized_dataset_manager._preformat_payloads(conversations) + mock_fmt.assert_not_called() + + assert conversations[0].turns[0].raw_payload is None + # And the resulting structured dataset selects CONVERSATION cleanly. + assert ( + initialized_dataset_manager._select_mmap_format(conversations) + == MemoryMapFormat.CONVERSATION + ) + + def test_preformat_proceeds_with_header_routing(self, initialized_dataset_manager): + """Headers-only routing modes leave the body untouched, so the + PAYLOAD_BYTES fast path stays available (no bail).""" + initialized_dataset_manager.user_config.endpoint.session_routing = ( + "dynamo_headers" + ) + + conversations = [ + Conversation( + session_id="s1", + turns=[Turn(role="user", texts=[Text(contents=["hi"])])], + ), + ] + + with patch( + "aiperf.dataset.dataset_manager.format_conversation_payloads" + ) as mock_fmt: + initialized_dataset_manager._preformat_payloads(conversations) + mock_fmt.assert_called_once() + class TestSelectMmapFormat: """Tests for DatasetManager._select_mmap_format format-selection guard.""" @@ -1301,7 +1357,7 @@ def test_select_format_rejects_payload_bytes_when_cache_bust_enabled( ] with pytest.raises( ValueError, - match=r"--cache-bust is incompatible with the PAYLOAD_BYTES", + match=r"cache-bust must mutate request bodies and is incompatible", ): initialized_dataset_manager._select_mmap_format(conversations) @@ -1325,16 +1381,18 @@ def test_select_format_allows_conversation_when_cache_bust_enabled( == MemoryMapFormat.CONVERSATION ) - def test_select_format_rejects_payload_bytes_when_dynamo_routing_enabled( + def test_select_format_rejects_payload_bytes_when_mutating_routing_enabled( self, initialized_dataset_manager ): - """Dynamo session-control + raw_payload-producing loader must raise. + """Body-mutating session routing + raw_payload-producing loader must raise. nvext.session_control mutates the request body, which the verbatim PAYLOAD_BYTES fast path streams pre-encoded and cannot carry -- the same conflict as cache-bust, refused early with an actionable error. """ - initialized_dataset_manager.user_config.endpoint.use_dynamo_conv_aware_routing = True + initialized_dataset_manager.user_config.endpoint.session_routing = ( + "dynamo_nvext" + ) conversations = [ Conversation( @@ -1342,19 +1400,18 @@ def test_select_format_rejects_payload_bytes_when_dynamo_routing_enabled( turns=[Turn(role="user", raw_payload={"a": 1})], ), ] - with pytest.raises( - ValueError, - match=r"--use-dynamo-conv-aware-routing is incompatible with the PAYLOAD_BYTES", - ): + with pytest.raises(ValueError, match=r"dynamo_nvext"): initialized_dataset_manager._select_mmap_format(conversations) - def test_select_format_allows_conversation_when_dynamo_routing_enabled( + def test_select_format_allows_conversation_when_mutating_routing_enabled( self, initialized_dataset_manager ): - """Dynamo routing with structured turns (no raw_payload) -> CONVERSATION.""" + """Body-mutating routing with structured turns (no raw_payload) -> CONVERSATION.""" from aiperf.common.enums import MemoryMapFormat - initialized_dataset_manager.user_config.endpoint.use_dynamo_conv_aware_routing = True + initialized_dataset_manager.user_config.endpoint.session_routing = ( + "dynamo_nvext" + ) conversations = [ Conversation( session_id="s1", @@ -1365,3 +1422,203 @@ def test_select_format_allows_conversation_when_dynamo_routing_enabled( initialized_dataset_manager._select_mmap_format(conversations) == MemoryMapFormat.CONVERSATION ) + + +# ============================================================================ +# PAYLOAD_BYTES body-mutating feature gates (session-routing + cache-bust) +# ============================================================================ + + +def _raw_payload_conversations() -> list[Conversation]: + """Conversations whose turns carry raw_payload (select PAYLOAD_BYTES).""" + return [ + Conversation(session_id="s1", turns=[Turn(role="user", raw_payload={"a": 1})]) + ] + + +def _payload_bytes_cache_hit( + tmp_path: Path, *, source_loaded: bool = True +) -> mmap_cache.CacheHit: + """Minimal CacheHit whose manifest reports PAYLOAD_BYTES. + + ``source_loaded=True`` marks the entry's payload bytes as shipped by the + dataset itself (hard-fail material); ``False`` marks a preformat-promoted + entry (downgraded to a MISS under a body-mutator). The cache-hit gate is + the first statement of ``_configure_from_cache_hit`` and raises before any + file restore, so the on-disk paths need not exist. + """ + manifest = mmap_cache.CacheManifest( + cache_key="test-key", + created_at=0.0, + num_conversations=1, + total_size_bytes=1, + mmap_format=str(MemoryMapFormat.PAYLOAD_BYTES), + dataset_metadata_json="{}", + all_turns_source_loaded_payloads=source_loaded, + ) + return mmap_cache.CacheHit( + entry_dir=tmp_path, + data_path=tmp_path / "dataset.dat", + index_path=tmp_path / "index.dat", + manifest=manifest, + ) + + +class TestPayloadBytesBodyMutatingGates: + """PAYLOAD_BYTES is refused whenever a body-mutating feature is active. + + Covers both gates that key off ``_body_mutating_feature``: build-path + format selection (``_select_mmap_format``) and cache-hit adoption + (``_configure_from_cache_hit`` / + ``_reject_body_mutators_for_payload_bytes``). + """ + + def test_select_format_rejects_payload_bytes_with_mutating_routing( + self, initialized_dataset_manager + ) -> None: + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + + with pytest.raises(ValueError, match="dynamo_nvext"): + dm._select_mmap_format(_raw_payload_conversations()) + + def test_select_format_allows_payload_bytes_with_header_routing( + self, initialized_dataset_manager + ) -> None: + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_headers" + + assert ( + dm._select_mmap_format(_raw_payload_conversations()) + == MemoryMapFormat.PAYLOAD_BYTES + ) + + @pytest.mark.asyncio + async def test_cache_hit_rejects_payload_bytes_with_mutating_routing( + self, initialized_dataset_manager, tmp_path + ) -> None: + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + + with pytest.raises(ValueError, match="dynamo_nvext"): + await dm._configure_from_cache_hit(_payload_bytes_cache_hit(tmp_path)) + + @pytest.mark.asyncio + async def test_cache_hit_rejects_payload_bytes_with_cache_bust( + self, initialized_dataset_manager, tmp_path + ) -> None: + dm = initialized_dataset_manager + dm.user_config.input.prompt.cache_bust.target = CacheBustTarget.SYSTEM_PREFIX + + with pytest.raises(ValueError, match="cache-bust"): + await dm._configure_from_cache_hit(_payload_bytes_cache_hit(tmp_path)) + + def test_cache_hit_allows_payload_bytes_when_clean( + self, initialized_dataset_manager + ) -> None: + dm = initialized_dataset_manager + # No routing, no cache-bust: the pre-check must pass (no raise). + assert dm.user_config.endpoint.session_routing is None + assert dm.user_config.input.prompt.cache_bust.target == CacheBustTarget.NONE + + dm._reject_body_mutators_for_payload_bytes(MemoryMapFormat.PAYLOAD_BYTES) + + def test_promoted_cache_hit_downgraded_to_miss_under_mutating_routing( + self, initialized_dataset_manager, tmp_path + ) -> None: + """A preformat-promoted PAYLOAD_BYTES entry is rebuildable: with the + body-mutator active the preformatter bails, so the rebuild stays + structured. The lookup must therefore treat the hit as a MISS instead + of hard-failing the run.""" + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + hit = _payload_bytes_cache_hit(tmp_path, source_loaded=False) + + assert dm._downgrade_body_mutator_cache_hit(hit) is None + + def test_source_loaded_cache_hit_not_downgraded_and_hard_fails( + self, initialized_dataset_manager, tmp_path + ) -> None: + """Payload bytes shipped by the dataset itself cannot be rebuilt as + structured turns: the hit survives lookup and the restore-path gate + raises the actionable hard-fail.""" + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + hit = _payload_bytes_cache_hit(tmp_path, source_loaded=True) + + assert dm._downgrade_body_mutator_cache_hit(hit) is hit + + def test_promoted_cache_hit_kept_without_body_mutator( + self, initialized_dataset_manager, tmp_path + ) -> None: + """No body-mutating feature active: promoted entries stay valid hits.""" + dm = initialized_dataset_manager + hit = _payload_bytes_cache_hit(tmp_path, source_loaded=False) + + assert dm._downgrade_body_mutator_cache_hit(hit) is hit + + def test_conversation_cache_hit_never_downgraded( + self, initialized_dataset_manager, tmp_path + ) -> None: + """CONVERSATION-format entries carry structured turns and are always + compatible with body mutators.""" + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + hit = _payload_bytes_cache_hit(tmp_path, source_loaded=False) + hit.manifest.mmap_format = str(MemoryMapFormat.CONVERSATION) + + assert dm._downgrade_body_mutator_cache_hit(hit) is hit + + def test_try_cache_lookup_wires_the_downgrade( + self, initialized_dataset_manager, tmp_path + ) -> None: + """The downgrade must be applied AT the lookup site, not only exist as + a helper: a promoted PAYLOAD_BYTES hit under a mutating mode comes back + as a MISS from _try_cache_lookup itself.""" + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + hit = _payload_bytes_cache_hit(tmp_path, source_loaded=False) + + with ( + patch( + "aiperf.dataset.dataset_manager.mmap_cache.cache_enabled", + return_value=True, + ), + patch( + "aiperf.dataset.dataset_manager.mmap_cache.compute_cache_key_from_user_config", + return_value="key-1", + ), + patch("aiperf.dataset.dataset_manager.mmap_cache.lookup", return_value=hit), + ): + assert dm._try_cache_lookup() is None + + def test_lookup_under_lock_wires_the_downgrade( + self, initialized_dataset_manager, tmp_path + ) -> None: + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + dm._cache_key_for_run = "key-1" + hit = _payload_bytes_cache_hit(tmp_path, source_loaded=False) + + with patch( + "aiperf.dataset.dataset_manager.mmap_cache.lookup", return_value=hit + ): + assert dm._lookup_under_lock() is None + + def test_populate_after_run_skipped_under_body_mutator( + self, initialized_dataset_manager, tmp_path + ) -> None: + """A body-mutator run must not write its CONVERSATION build under the + shared cache key (which excludes routing settings) -- that would demote + every later feature-free run of the same dataset off the PAYLOAD_BYTES + fast path.""" + dm = initialized_dataset_manager + dm.user_config.endpoint.session_routing = "dynamo_nvext" + dm._cache_hit_used = False + dm._cache_key_for_run = "key-1" + dm._backing_store = MagicMock() + dm.dataset_metadata = MagicMock() + + with patch.object(dm, "_run_mmap_paths") as run_paths: + dm._populate_cache_after_run() + run_paths.assert_not_called() diff --git a/tests/unit/dataset/test_dataset_manager_cache.py b/tests/unit/dataset/test_dataset_manager_cache.py index 95f77cb772..9d734274fc 100644 --- a/tests/unit/dataset/test_dataset_manager_cache.py +++ b/tests/unit/dataset/test_dataset_manager_cache.py @@ -189,7 +189,9 @@ async def test_hf_weka_dataset_change_invalidates_cache( cfg_b.input.file = None cfg_b.input.custom_dataset_type = None cfg_b.input.public_dataset = PublicDatasetType.WEKA_HF - cfg_b.input.hf_weka_dataset = "semianalysisai/cc-traces-weka-with-subagents-051826" + cfg_b.input.hf_weka_dataset = ( + "semianalysisai/cc-traces-weka-with-subagents-051826" + ) key_a = mmap_cache.compute_cache_key_from_user_config(cfg_a) key_b = mmap_cache.compute_cache_key_from_user_config(cfg_b) diff --git a/tests/unit/dataset/test_dataset_manager_inputs_json_adversarial.py b/tests/unit/dataset/test_dataset_manager_inputs_json_adversarial.py index 11aa4d64f0..e16dcdfb1b 100644 --- a/tests/unit/dataset/test_dataset_manager_inputs_json_adversarial.py +++ b/tests/unit/dataset/test_dataset_manager_inputs_json_adversarial.py @@ -230,9 +230,11 @@ def _make_mgr(self, convs: list[Conversation]) -> DatasetManager: mgr = object.__new__(DatasetManager) # Minimal user_config: any non-None object is enough to pass the guard. mgr.user_config = Mock() - # Disable cache-bust so the preformat path runs (the cache-bust early - # return bails preformatting whenever target != NONE). + # Disable both body-mutators so the preformat path runs (the + # _body_mutating_feature early return bails preformatting whenever + # cache-bust or a body-mutating session-routing mode is active). mgr.user_config.input.prompt.cache_bust.target = CacheBustTarget.NONE + mgr.user_config.endpoint.session_routing = None # Stub the logger mixin attrs that _preformat_payloads uses. mgr.info = Mock() return mgr diff --git a/tests/unit/exporters/test_console_api_error_exporter.py b/tests/unit/exporters/test_console_api_error_exporter.py index 28d1591a4f..9ec262173c 100644 --- a/tests/unit/exporters/test_console_api_error_exporter.py +++ b/tests/unit/exporters/test_console_api_error_exporter.py @@ -131,8 +131,11 @@ def test_detects_raw_serde_unknown_variant_error(self): insight = DynamoSessionControlDetector.detect(make_summary(err)) assert insight is not None - assert "bind" in insight.problem - assert any("--use-legacy-dynamo-session-control" in f for f in insight.fixes) + assert "bind" in insight.title + assert "session_control" in insight.problem + assert any("--session-routing dynamo_nvext" in c for c in insight.causes) + assert any("--session-routing" in f for f in insight.fixes) + assert not any("use-legacy" in f for f in insight.fixes) def test_detects_json_wrapped_unknown_variant_error(self): """The serde message is often wrapped in a JSON error envelope.""" @@ -165,6 +168,35 @@ def test_returns_none_when_no_errors(self): assert DynamoSessionControlDetector.detect(None) is None assert DynamoSessionControlDetector.detect([]) is None + def test_detect_item_without_error_details_skipped_returns_none(self): + summary = [MockErrorDetailsCount(None, 1)] + + assert DynamoSessionControlDetector.detect(summary) is None + + def test_detect_none_message_returns_none(self): + summary = make_summary(MockErrorDetails(message=None)) + + assert DynamoSessionControlDetector.detect(summary) is None + + def test_detects_unknown_field_session_control_error(self): + """Current Dynamo main (#[serde(deny_unknown_fields)] on NvExt) rejects + nvext.session_control as an unknown FIELD; recommend dynamo_headers.""" + err = MockErrorDetails( + message="Failed to deserialize the JSON body: nvext: unknown field `session_control`" + ) + insight = DynamoSessionControlDetector.detect(make_summary(err)) + assert insight is not None + assert any("--session-routing dynamo_headers" in f for f in insight.fixes) + + def test_bind_rejection_fixes_reference_new_flag(self): + err = MockErrorDetails( + message="unknown variant `bind`, expected `open` or `close`" + ) + insight = DynamoSessionControlDetector.detect(make_summary(err)) + assert insight is not None + assert not any("use-legacy" in f for f in insight.fixes) + assert any("--session-routing" in f for f in insight.fixes) + @pytest.mark.asyncio async def test_exporter_prints_panel_for_bind_rejection(self): mock_console = MagicMock(spec=Console) @@ -182,5 +214,7 @@ async def test_exporter_prints_panel_for_bind_rejection(self): assert mock_console.print.call_count >= 2 _, args, _ = mock_console.print.mock_calls[1] panel = args[0] - assert "session_control action: bind" in str(panel.title) - assert "--use-legacy-dynamo-session-control" in str(panel.renderable) + assert "Unsupported Dynamo session_control action: bind" in str(panel.title) + panel_text = str(panel.renderable) + assert "--session-routing" in panel_text + assert "use-legacy" not in panel_text 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/test_session_tree_finality.py b/tests/unit/timing/test_session_tree_finality.py new file mode 100644 index 0000000000..ed4649a095 --- /dev/null +++ b/tests/unit/timing/test_session_tree_finality.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from aiperf.common.enums import CreditPhase +from aiperf.timing.session_tree import SessionTreeRegistry + + +class _FakeConcurrency: + """Minimal stand-in for the concurrency manager the registry releases to. + + ``SessionTreeRegistry`` requires a concurrency manager exposing + ``release_session_slot(phase)``; these finality-query tests only exercise + the state map, so the release is a no-op. + """ + + def release_session_slot(self, phase: CreditPhase) -> None: + pass + + +def _make_registry() -> SessionTreeRegistry: + return SessionTreeRegistry(_FakeConcurrency()) + + +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 rootless drain would pop the tree + # (release path) 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_spawn_declaring_final_turn_gated_by_has_branches(): + """Same registry state that yields True with ``has_branches=False`` must + yield False when the turn declares ANY branch: a SPAWN-declaring final turn + (``has_forks`` stays False) registers its children only at return-intercept, + AFTER this stamp, so it can never be provably-last.""" + registry = _registry_with_tree(root="root-3") + assert registry.is_last_tree_request( + "root-3", is_final_turn=True, is_root_credit=True, has_branches=False + ) + assert not registry.is_last_tree_request( + "root-3", is_final_turn=True, is_root_credit=True, has_branches=True + ) + + +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_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"} + assert credit_fields <= set(Credit.__struct_fields__) + assert credit_fields <= set(RequestInfo.model_fields) diff --git a/tests/unit/transports/test_base_transport.py b/tests/unit/transports/test_base_transport.py index 1f0aa690ca..acd4884a76 100644 --- a/tests/unit/transports/test_base_transport.py +++ b/tests/unit/transports/test_base_transport.py @@ -7,7 +7,6 @@ import pytest from aiperf.common.enums import CreditPhase, ModelSelectionStrategy -from aiperf.common.environment import Environment from aiperf.common.models.model_endpoint_info import ( EndpointInfo, ModelEndpointInfo, @@ -154,28 +153,6 @@ def test_build_headers_merges_endpoint_headers(self, transport, request_info): assert headers["Custom-Header"] == "custom-value" assert headers["User-Agent"] == AIPERF_USER_AGENT - def test_build_headers_can_alias_correlation_id_as_session_id( - self, transport, request_info, monkeypatch - ): - """Test opt-in session affinity header for external routers.""" - monkeypatch.setattr(Environment.HTTP, "X_SESSION_ID_FROM_CORRELATION_ID", True) - - headers = transport.build_headers(request_info) - - assert headers["X-Session-ID"] == "test-correlation-id" - - def test_build_headers_can_alias_correlation_id_as_smg_routing_key( - self, transport, request_info, monkeypatch - ): - """Test opt-in SGLang Model Gateway affinity header.""" - monkeypatch.setattr( - Environment.HTTP, "X_SMG_ROUTING_KEY_FROM_CORRELATION_ID", True - ) - - headers = transport.build_headers(request_info) - - assert headers["X-SMG-Routing-Key"] == "test-correlation-id" - def test_build_headers_transport_headers_override(self, request_info): """Test that transport headers can override endpoint headers.""" diff --git a/tests/unit/workers/test_dynamo_session_control.py b/tests/unit/workers/test_dynamo_session_control.py deleted file mode 100644 index b72d80b4ce..0000000000 --- a/tests/unit/workers/test_dynamo_session_control.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Unit tests for the Dynamo nvext.session_control helpers.""" - -from aiperf.workers.dynamo_session_control import ( - build_session_control, - merge_session_control, -) - - -class TestBuildSessionControl: - """Policy: which lifecycle action each turn emits (modern / bind contract).""" - - def test_non_final_turn_binds_with_timeout(self): - """Non-final turns re-bind (idempotent, refreshes the router TTL).""" - sc = build_session_control( - session_id="conv-1", is_final_turn=False, timeout_seconds=300 - ) - assert sc == {"session_id": "conv-1", "action": "bind", "timeout": 300} - - def test_modern_ignores_already_opened(self): - """Modern mode always re-binds; already_opened only affects legacy.""" - sc = build_session_control( - session_id="conv-1", - is_final_turn=False, - timeout_seconds=300, - already_opened=True, - ) - assert sc == {"session_id": "conv-1", "action": "bind", "timeout": 300} - - def test_final_turn_closes_without_timeout(self): - """Final turn closes the session; timeout is irrelevant to close.""" - sc = build_session_control( - session_id="conv-1", is_final_turn=True, timeout_seconds=300 - ) - assert sc == {"session_id": "conv-1", "action": "close"} - - def test_session_id_and_timeout_are_passed_through(self): - """session_id is the routing key; timeout flows through on bind.""" - sc = build_session_control( - session_id="x-corr-abc", - is_final_turn=False, - timeout_seconds=42, - ) - assert sc["session_id"] == "x-corr-abc" - assert sc["timeout"] == 42 - - -class TestBuildSessionControlLegacy: - """Legacy (v1.2.x) contract: open on first request, close on last, else bare. - - The 'first request' is signalled by ``already_opened=False`` (the caller - tracks it per worker), NOT by turn_index -- agentic replay's first request - is the warmup turn k_i, never turn 0. - """ - - def test_first_request_opens_with_timeout(self): - """The first request for a session (not yet opened) emits open.""" - sc = build_session_control( - session_id="conv-1", - is_final_turn=False, - timeout_seconds=300, - legacy=True, - already_opened=False, - ) - assert sc == {"session_id": "conv-1", "action": "open", "timeout": 300} - - def test_already_opened_sends_session_id_only(self): - """Once opened, subsequent turns carry only session_id (sticky routing).""" - sc = build_session_control( - session_id="conv-1", - is_final_turn=False, - timeout_seconds=300, - legacy=True, - already_opened=True, - ) - assert sc == {"session_id": "conv-1"} - - def test_final_turn_closes(self): - """Final turn closes, same as modern mode (even if already opened).""" - sc = build_session_control( - session_id="conv-1", - is_final_turn=True, - timeout_seconds=300, - legacy=True, - already_opened=True, - ) - assert sc == {"session_id": "conv-1", "action": "close"} - - def test_single_turn_closes_rather_than_opens(self): - """is_final_turn wins: a single-turn (never-opened) session emits close.""" - sc = build_session_control( - session_id="conv-1", - is_final_turn=True, - timeout_seconds=300, - legacy=True, - already_opened=False, - ) - assert sc == {"session_id": "conv-1", "action": "close"} - - def test_never_emits_bind(self): - """Legacy mode must never emit the 'bind' action (rejected by v1.2.x).""" - actions = { - build_session_control( - session_id="c", - is_final_turn=False, - timeout_seconds=300, - legacy=True, - already_opened=opened, - ).get("action") - for opened in (False, True) - } - assert "bind" not in actions - - -class TestMergeSessionControl: - """Mechanism: overlay session_control under nvext without mutating input.""" - - def test_adds_nvext_session_control_to_bare_payload(self): - payload = {"messages": [], "model": "m"} - sc = {"session_id": "c1", "action": "bind", "timeout": 300} - - merged = merge_session_control(payload, sc) - - assert merged["nvext"]["session_control"] == sc - assert merged["messages"] == [] - assert merged["model"] == "m" - - def test_preserves_existing_nvext_and_session_control_fields(self): - payload = {"nvext": {"trace": "keep", "session_control": {"existing": "keep"}}} - sc = {"session_id": "c1", "action": "bind", "timeout": 300} - - merged = merge_session_control(payload, sc) - - assert merged["nvext"] == { - "trace": "keep", - "session_control": { - "existing": "keep", - "session_id": "c1", - "action": "bind", - "timeout": 300, - }, - } - - def test_does_not_mutate_input_payload_or_nested_dicts(self): - """Safe to call on a cached Turn.raw_payload / shared extra_body.""" - nested_sc = {"existing": "keep"} - nvext = {"session_control": nested_sc} - payload = {"nvext": nvext} - sc = {"session_id": "c1", "action": "close"} - - merged = merge_session_control(payload, sc) - - # Inputs are left pristine at every level. - assert payload == {"nvext": {"session_control": {"existing": "keep"}}} - assert nvext == {"session_control": {"existing": "keep"}} - assert nested_sc == {"existing": "keep"} - assert merged is not payload - # The returned copy carries the overlay. - assert merged["nvext"]["session_control"] == { - "existing": "keep", - "session_id": "c1", - "action": "close", - } diff --git a/tests/unit/workers/test_inference_client.py b/tests/unit/workers/test_inference_client.py index 139291481d..9476cda617 100644 --- a/tests/unit/workers/test_inference_client.py +++ b/tests/unit/workers/test_inference_client.py @@ -20,6 +20,22 @@ 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, + IdentityHeadersRouting, + SessionIdHeaderRouting, + SessionRoutingBase, + SmgRoutingKeyRouting, +) + +_ROUTING_CLASSES = { + "dynamo_headers": DynamoHeadersRouting, + "dynamo_nvext": DynamoNvextRouting, + "smg_routing_key": SmgRoutingKeyRouting, + "session_id_header": SessionIdHeaderRouting, + "identity_headers": IdentityHeadersRouting, +} @pytest.fixture @@ -494,16 +510,22 @@ def test_enrich_keeps_payload_bytes_by_default( assert enriched.request_info.payload_bytes == request_info.payload_bytes -class TestInferenceClientDynamoSessionControl: - """Chokepoint injection of nvext.session_control for Dynamo routing. +class TestInferenceClientSessionRouting: + """Session-routing plugins wired through the InferenceClient chokepoint. - The verbatim PAYLOAD_BYTES path is refused against this feature at dataset - load, so injection only ever runs on the structured (format_payload) body. + 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. """ - @pytest.fixture - def model_endpoint(self): - return ModelEndpointInfo( + 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, @@ -511,20 +533,17 @@ def model_endpoint(self): endpoint=EndpointInfo( type=EndpointType.CHAT, base_url="http://localhost:8000/v1/test", - use_dynamo_conv_aware_routing=True, - dynamo_session_timeout_seconds=123, + session_routing=session_routing, + session_routing_opts=session_routing_opts or {}, ), ) - - @pytest.fixture - def inference_client(self, model_endpoint, mock_http_transport_entry): 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 = { - "messages": [{"role": "user", "content": "hi"}], "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], } def mock_get_class(protocol, name): @@ -532,6 +551,8 @@ def mock_get_class(protocol, name): 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 ( @@ -544,204 +565,375 @@ def mock_get_class(protocol, name): return_value=[mock_http_transport_entry], ), ): - return InferenceClient( + 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, inference_client, *, is_final_turn, x_correlation_id="corr-1" - ): + self, + client: InferenceClient, + *, + x_correlation_id: str = "corr-1", + parent_correlation_id: str | None = None, + is_final_turn: bool = False, + ) -> RequestInfo: return RequestInfo( - model_endpoint=inference_client.model_endpoint, - turns=[Turn(role="user", texts=[Text(contents=["hi"])])], + model_endpoint=client.model_endpoint, + turns=[Turn(role="user", texts=[Text(contents=["hello"])])], turn_index=0, - credit_num=1, + credit_num=0, credit_phase=CreditPhase.PROFILING, - x_request_id="rid", + x_request_id="req-1", x_correlation_id=x_correlation_id, - conversation_id="conv", + parent_correlation_id=parent_correlation_id, + # Mirrors the worker, which always passes + # credit.effective_root_correlation_id (own ID for roots). + root_correlation_id=x_correlation_id, + conversation_id="conv-template", is_final_turn=is_final_turn, ) - async def _sent_payload(self, inference_client, request_info): - inference_client.transport.send_request = AsyncMock( - return_value=RequestRecord(request_info=request_info) + def _sent_payload(self, client: InferenceClient): + payload = client.transport.send_request.call_args.kwargs["payload"] + if isinstance(payload, bytes): + return orjson.loads(payload) + return 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" ) - await inference_client.send_request(request_info) - return orjson.loads( - inference_client.transport.send_request.call_args.kwargs["payload"] + 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_non_final_turn_binds_with_x_correlation_id_and_timeout( - self, inference_client - ): - payload = await self._sent_payload( - inference_client, - self._request_info(inference_client, is_final_turn=False), + 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"}, ) - assert payload["nvext"]["session_control"] == { + + 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, } - # Endpoint-built fields are preserved. - assert payload["messages"] == [{"role": "user", "content": "hi"}] - @pytest.mark.asyncio - async def test_final_turn_closes_session(self, inference_client): - payload = await self._sent_payload( - inference_client, - self._request_info(inference_client, is_final_turn=True), - ) - assert payload["nvext"]["session_control"] == { + 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_disabled_leaves_payload_untouched(self, inference_client): - inference_client.model_endpoint.endpoint.use_dynamo_conv_aware_routing = False - payload = await self._sent_payload( - inference_client, - self._request_info(inference_client, is_final_turn=False), + async def test_session_id_header_preset(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, session_routing="session_id_header" ) - assert "nvext" not in payload + request_info = self._request_info(client) + await client._send_request_to_transport(request_info) -class TestInferenceClientLegacySessionControl: - """Legacy (v1.2.x) open/close lifecycle, with the agentic-replay edge case. - - The critical property: 'open' fires on the FIRST request the worker sends - for a session, tracked per-worker -- NOT on turn_index 0. Agentic replay - warms at k_i and profiles from k_i+1, so the first request a worker sees for - a session carries a NON-ZERO turn_index; a turn_index==0 gate would never - emit 'open' for those sessions. - """ + assert request_info.endpoint_headers["X-Session-ID"] == "corr-1" + assert "nvext" not in self._sent_payload(client) - @pytest.fixture - def model_endpoint(self): - return 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", - use_dynamo_conv_aware_routing=True, - use_legacy_dynamo_session_control=True, - dynamo_session_timeout_seconds=123, - ), + async def test_identity_headers_custom_tiers(self, mock_http_transport_entry): + client = self._build_client( + mock_http_transport_entry, + session_routing="identity_headers", + session_routing_opts={"session": "X-Affinity", "root": "X-Tree-ID"}, ) + request_info = self._request_info(client) - @pytest.fixture - def inference_client(self, model_endpoint, mock_http_transport_entry): - 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 = { - "messages": [{"role": "user", "content": "hi"}], - "model": "test-model", - } + await client._send_request_to_transport(request_info) - def mock_get_class(protocol, name): - if protocol == "endpoint": - return lambda **kwargs: mock_endpoint - if protocol == "transport": - return lambda **kwargs: mock_transport - raise ValueError(f"Unknown protocol: {protocol}") + assert request_info.endpoint_headers["X-Affinity"] == "corr-1" + assert request_info.endpoint_headers["X-Tree-ID"] == "corr-1" + assert "nvext" not in self._sent_payload(client) - 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], - ), - ): - return InferenceClient( - model_endpoint=model_endpoint, service_id="test-service-id" - ) + @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") - async def _sent_sc( - self, inference_client, *, turn_index, is_final_turn, x_correlation_id="corr-1" + 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_payload_bytes_with_mutating_plugin_yields_error_record( + self, mock_http_transport_entry ): - request_info = RequestInfo( - model_endpoint=inference_client.model_endpoint, - turns=[Turn(role="user", texts=[Text(contents=["hi"])])], - turn_index=turn_index, - credit_num=1, - credit_phase=CreditPhase.PROFILING, - x_request_id="rid", - x_correlation_id=x_correlation_id, - conversation_id="conv", - is_final_turn=is_final_turn, + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" ) - inference_client.transport.send_request = AsyncMock( - return_value=RequestRecord(request_info=request_info) + request_info = self._request_info(client) + request_info.payload_bytes = b'{"a":1}' + + record = await client.send_request(request_info) + + assert record.error is not None + assert "PAYLOAD_BYTES" in record.error.message + + @pytest.mark.asyncio + async def test_payload_bytes_with_header_plugin_gets_headers( + self, mock_http_transport_entry + ): + client = self._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" ) - await inference_client.send_request(request_info) - payload = orjson.loads( - inference_client.transport.send_request.call_args.kwargs["payload"] + request_info = self._request_info(client, parent_correlation_id="parent-corr") + request_info.payload_bytes = b'{"a":1}' + + 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" ) - return payload["nvext"]["session_control"] + # The verbatim bytes are forwarded to the transport untouched. + assert client.transport.send_request.call_args.kwargs["payload"] == b'{"a":1}' @pytest.mark.asyncio - async def test_open_fires_on_first_request_with_nonzero_turn_index( - self, inference_client + 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 ): - """Agentx fix: the worker's first request for a session is the warmup - turn at k_i (non-zero turn_index), and it must still emit 'open'.""" - sc = await self._sent_sc(inference_client, turn_index=5, is_final_turn=False) - assert sc == {"session_id": "corr-1", "action": "open", "timeout": 123} + 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") - @pytest.mark.asyncio - async def test_open_once_then_session_id_only_then_close(self, inference_client): - """Full lifecycle across the warmup->profiling boundary on one worker.""" - # warmup turn k_i: first request -> open - warm = await self._sent_sc(inference_client, turn_index=5, is_final_turn=False) - assert warm["action"] == "open" - # profiling resume k_i+1: already opened -> session_id only, NO action - mid = await self._sent_sc(inference_client, turn_index=6, is_final_turn=False) - assert mid == {"session_id": "corr-1"} - # another profiling turn: still session_id only - mid2 = await self._sent_sc(inference_client, turn_index=7, is_final_turn=False) - assert mid2 == {"session_id": "corr-1"} - # final turn -> close - final = await self._sent_sc(inference_client, turn_index=8, is_final_turn=True) - assert final == {"session_id": "corr-1", "action": "close"} - # 'open' emitted exactly once for the session. + 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_close_clears_tracking_state(self, inference_client): - """The opened-sessions set is bounded: close drops the entry.""" - await self._sent_sc(inference_client, turn_index=5, is_final_turn=False) - assert "corr-1" in inference_client._dynamo_opened_sessions - await self._sent_sc(inference_client, turn_index=6, is_final_turn=True) - assert "corr-1" not in inference_client._dynamo_opened_sessions + 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_each_session_opens_independently(self, inference_client): - """Distinct sessions each get their own 'open'.""" - a = await self._sent_sc( - inference_client, turn_index=2, is_final_turn=False, x_correlation_id="a" + 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" ) - b = await self._sent_sc( - inference_client, turn_index=9, is_final_turn=False, x_correlation_id="b" + client._routing.transform_body = MagicMock( + side_effect=RuntimeError("bad body transform") ) - assert a["action"] == "open" - assert b["action"] == "open" - assert {"a", "b"} <= inference_client._dynamo_opened_sessions + 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 + + +class TestRoutingCapabilityFlags: + """Init-time capability detection and the mutates_body contract check.""" + + @pytest.mark.parametrize( + "mode, opts, stamps_headers, transforms_body", + [ + param("dynamo_headers", None, True, False, id="dynamo_headers"), + param("dynamo_nvext", None, False, True, id="dynamo_nvext"), + param("smg_routing_key", None, True, False, id="smg_routing_key"), + param("session_id_header", None, True, False, id="session_id_header"), + param("identity_headers", {"session": "X-A"}, True, False, id="identity_headers"), + ], + ) # fmt: skip + def test_flags_per_builtin( + self, mock_http_transport_entry, mode, opts, stamps_headers, transforms_body + ): + client = TestInferenceClientSessionRouting()._build_client( + mock_http_transport_entry, session_routing=mode, session_routing_opts=opts + ) + assert client._routing_stamps_headers is stamps_headers + assert client._routing_transforms_body is transforms_body @pytest.mark.asyncio - async def test_never_emits_bind(self, inference_client): - """Legacy mode must never put 'bind' on the wire (v1.2.x rejects it).""" - for ti in range(4): - sc = await self._sent_sc( - inference_client, turn_index=ti, is_final_turn=False - ) - assert sc.get("action") != "bind" + async def test_body_only_plugin_never_calls_headers( + self, mock_http_transport_entry + ): + """The capability flag must skip the no-op base headers() entirely for + a body-only plugin (spy patched AFTER init so it does not flip the + init-time flag).""" + client = TestInferenceClientSessionRouting()._build_client( + mock_http_transport_entry, session_routing="dynamo_nvext" + ) + spy = MagicMock(return_value={}) + client._routing.headers = spy + request_info = TestInferenceClientSessionRouting()._request_info(client) + + await client._send_request_to_transport(request_info) + + spy.assert_not_called() + + @pytest.mark.asyncio + async def test_header_only_plugin_never_calls_transform_body( + self, mock_http_transport_entry + ): + client = TestInferenceClientSessionRouting()._build_client( + mock_http_transport_entry, session_routing="dynamo_headers" + ) + spy = MagicMock(side_effect=lambda payload, ctx: payload) + client._routing.transform_body = spy + request_info = TestInferenceClientSessionRouting()._request_info(client) + + await client._send_request_to_transport(request_info) + + spy.assert_not_called() + + def test_transform_override_without_mutates_body_rejected_at_init( + self, mock_http_transport_entry + ): + """Split-brain plugin: transform_body overridden but mutates_body left + False would pass every PAYLOAD_BYTES gate and silently never fire. + Worker init must refuse it.""" + + class ForgotDeclaration(SessionRoutingBase): + def transform_body(self, payload, ctx): + return {**payload, "stamped": True} + + with pytest.raises(ValueError, match="mutates_body"): + _build_custom_routing_client(mock_http_transport_entry, ForgotDeclaration) + + def test_mutates_body_without_transform_rejected_at_init( + self, mock_http_transport_entry + ): + class DeclaredButInert(SessionRoutingBase): + mutates_body = True + + with pytest.raises(ValueError, match="mutates_body"): + _build_custom_routing_client(mock_http_transport_entry, DeclaredButInert) + + def test_instance_bound_methods_detected_as_capabilities( + self, mock_http_transport_entry + ): + """A plugin selecting strategies by binding instance attributes in + __init__ is a legal shape; detection must see it.""" + + class StrategyBound(SessionRoutingBase): + mutates_body = True + + def __init__(self, options): + super().__init__(options) + self.headers = lambda ctx: {"X-S": ctx.x_correlation_id} + self.transform_body = lambda payload, ctx: {**payload, "s": 1} + + client = _build_custom_routing_client(mock_http_transport_entry, StrategyBound) + assert client._routing_stamps_headers is True + assert client._routing_transforms_body is True + + +def _build_custom_routing_client(mock_http_transport_entry, routing_cls): + """Build an InferenceClient whose session_routing resolves to routing_cls.""" + 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="custom_mode", + ), + ) + mock_transport = MagicMock() + mock_endpoint = MagicMock() + mock_endpoint.get_endpoint_headers.return_value = {} + mock_endpoint.get_endpoint_params.return_value = {} + + 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_cls + 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], + ), + ): + return InferenceClient( + model_endpoint=model_endpoint, service_id="test-service-id" + ) diff --git a/tests/unit/workers/test_session_routing.py b/tests/unit/workers/test_session_routing.py new file mode 100644 index 0000000000..b4706ba981 --- /dev/null +++ b/tests/unit/workers/test_session_routing.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pydantic import ValidationError +from pytest import param + +from aiperf.workers.session_routing import ( + DynamoHeadersRouting, + DynamoNvextOptions, + DynamoNvextRouting, + IdentityHeadersOptions, + IdentityHeadersRouting, + RoutingContext, + 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 + + def test_plugin_session_control_wins_over_dataset_shipped_keys(self): + """Merge precedence: the plugin's live session identity must override + any session_control keys the dataset shipped, or a recorded + session_id would leak into live routing.""" + payload = { + "nvext": { + "session_control": { + "session_id": "recorded-stale", + "action": "open", + "keep": "me", + } + } + } + plugin = DynamoNvextRouting(DynamoNvextOptions(timeout_seconds=42)) + merged = plugin.transform_body(payload, _ctx()) + sc = merged["nvext"]["session_control"] + assert sc["session_id"] == "corr-1" + assert sc["action"] == "bind" + assert sc["timeout"] == 42 + assert sc["keep"] == "me" + + def test_nvext_present_without_session_control_preserved(self): + plugin = DynamoNvextRouting(DynamoNvextOptions()) + merged = plugin.transform_body({"nvext": {"trace": "keep"}}, _ctx()) + assert merged["nvext"]["trace"] == "keep" + assert merged["nvext"]["session_control"]["action"] == "bind" + + def test_non_dict_nvext_replaced(self): + """A malformed (non-dict) nvext value is replaced rather than crashed on.""" + plugin = DynamoNvextRouting(DynamoNvextOptions()) + merged = plugin.transform_body({"nvext": "bogus"}, _ctx()) + assert merged["nvext"]["session_control"]["session_id"] == "corr-1" + + +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_emits_x_session_id(self): + plugin = SessionIdHeaderRouting(SessionIdHeaderRouting.Options()) + assert plugin.headers(_ctx()) == {"X-Session-ID": "corr-1"} + assert plugin.mutates_body is False + + def test_rejects_any_opt(self): + with pytest.raises(ValidationError): + SessionIdHeaderRouting.Options(header_name="X-Affinity") + + +class TestIdentityHeadersRouting: + def test_session_tier_single_name(self): + plugin = IdentityHeadersRouting(IdentityHeadersOptions(session="X-Affinity")) + assert plugin.headers(_ctx()) == {"X-Affinity": "corr-1"} + assert plugin.mutates_body is False + + def test_session_tier_multiple_names_comma_separated(self): + """One mode, N additive headers, same correlation-ID value on each -- + the layered-router topology (e.g. ingress LB + SMG).""" + plugin = IdentityHeadersRouting( + IdentityHeadersOptions(session="X-Session-ID, X-SMG-Routing-Key") + ) + assert plugin.headers(_ctx()) == { + "X-Session-ID": "corr-1", + "X-SMG-Routing-Key": "corr-1", + } + + def test_list_form_not_resplit(self): + """Canonicalized opts re-enter as lists; must not be re-split.""" + plugin = IdentityHeadersRouting(IdentityHeadersOptions(session=["X-A", "X-B"])) + assert plugin.headers(_ctx()) == {"X-A": "corr-1", "X-B": "corr-1"} + + def test_parent_tier_omitted_for_roots(self): + plugin = IdentityHeadersRouting( + IdentityHeadersOptions(session="X-S", parent="X-P") + ) + assert plugin.headers(_ctx()) == {"X-S": "corr-1"} + + def test_parent_tier_emitted_for_children(self): + plugin = IdentityHeadersRouting( + IdentityHeadersOptions(session="X-S", parent="X-P") + ) + headers = plugin.headers(_ctx(parent_correlation_id="parent-1")) + assert headers == {"X-S": "corr-1", "X-P": "parent-1"} + + def test_root_tier_whole_tree_affinity(self): + plugin = IdentityHeadersRouting(IdentityHeadersOptions(root="X-Tree")) + headers = plugin.headers( + _ctx(parent_correlation_id="parent-1", root_correlation_id="root-1") + ) + assert headers == {"X-Tree": "root-1"} + + def test_root_tier_equals_session_for_roots(self): + """effective_root == x_corr for root sessions, so root-tier affinity + degrades gracefully on flat (non-tree) workloads.""" + plugin = IdentityHeadersRouting( + IdentityHeadersOptions(session="X-S", root="X-Tree") + ) + assert plugin.headers(_ctx()) == {"X-S": "corr-1", "X-Tree": "corr-1"} + + def test_dynamo_headers_preset_expressible(self): + """The dynamo_headers preset semantics, spelled via the generic.""" + plugin = IdentityHeadersRouting( + IdentityHeadersOptions( + session="X-Dynamo-Session-ID", parent="X-Dynamo-Parent-Session-ID" + ) + ) + preset = DynamoHeadersRouting(DynamoHeadersRouting.Options()) + for ctx in (_ctx(), _ctx(parent_correlation_id="parent-1")): + assert plugin.headers(ctx) == preset.headers(ctx) + + def test_no_names_anywhere_rejected(self): + with pytest.raises(ValidationError, match="at least one header name"): + IdentityHeadersOptions() + + def test_duplicate_name_across_tiers_rejected_case_insensitive(self): + with pytest.raises(ValidationError, match="duplicate header name"): + IdentityHeadersOptions(session="X-Affinity", root="x-affinity") + + def test_duplicate_name_within_tier_rejected(self): + with pytest.raises(ValidationError, match="duplicate header name"): + IdentityHeadersOptions(session="X-Affinity,x-affinity") + + def test_empty_name_entry_rejected(self): + with pytest.raises(ValidationError, match="non-empty"): + IdentityHeadersOptions(session="X-Affinity,") + + def test_unknown_opt_rejected(self): + with pytest.raises(ValidationError): + IdentityHeadersOptions(header_name="X-Affinity") + + @pytest.mark.parametrize( + "bad_name", + [ + param("X Foo", id="interior-space"), + param("X:Foo", id="colon"), + param("X-Foo\r\nX-Evil: 1", id="crlf-injection"), + param("X-Fo\u00e9", id="non-ascii"), + ], + ) # fmt: skip + def test_non_token_header_names_rejected(self, bad_name): + """Names must be RFC 9110 tokens; anything else fails at config time + instead of corrupting (or injecting into) the wire request.""" + with pytest.raises(ValidationError, match="RFC 9110"): + IdentityHeadersOptions(session=[bad_name]) + + def test_list_form_whitespace_only_name_rejected(self): + """YAML list inputs are stripped like CLI comma-splits; a + whitespace-only entry must hit the non-empty check, not emit a + malformed header.""" + with pytest.raises(ValidationError, match="non-empty"): + IdentityHeadersOptions(session=[" "]) + + def test_list_form_stripped_before_duplicate_check(self): + """' X-A ' (YAML list) and 'x-a' (another tier) are the same header + on the wire; stripping must happen before the cross-tier dup check.""" + with pytest.raises(ValidationError, match="duplicate header name"): + IdentityHeadersOptions(session=[" X-A "], root=["x-a"]) + + def test_list_form_names_stripped_in_output(self): + plugin = IdentityHeadersRouting(IdentityHeadersOptions(session=[" X-A "])) + assert plugin.headers(_ctx()) == {"X-A": "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 5327045df8..fd90a06468 100644 --- a/tests/unit/workers/test_worker.py +++ b/tests/unit/workers/test_worker.py @@ -1,15 +1,17 @@ # 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 +from pytest import param from aiperf.common.config.service_config import ServiceConfig from aiperf.common.config.user_config import UserConfig from aiperf.common.enums import CreditPhase from aiperf.common.models import ( Conversation, + ErrorDetails, ParsedResponse, ReasoningResponseData, RequestRecord, @@ -18,7 +20,7 @@ Turn, ) from aiperf.credit.structs import Credit, CreditContext -from aiperf.workers.worker import Worker +from aiperf.workers.worker import Worker, _is_terminal_context_overflow from tests.harness.fake_communication import FakeCommunication as FakeCommunication from tests.harness.fake_service_manager import FakeServiceManager as FakeServiceManager from tests.harness.fake_tokenizer import FakeTokenizer @@ -72,6 +74,33 @@ async def test_create_request_info_overrides_only_outgoing_turn(self, mock_worke assert request_info.turns[-1].max_tokens == 1 assert original.max_tokens == 4096 + async def test_create_request_info_plumbs_finality_from_credit(self, mock_worker): + # Real Credit struct (not a MagicMock, which would auto-create the + # attributes and mask a missed plumb) carrying both finality facts. + credit_context = CreditContext( + credit=Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="test-conv", + x_correlation_id="test-correlation", + turn_index=0, + num_turns=1, + issued_at_ns=0, + is_parent_final=True, + is_tree_final=True, + ), + drop_perf_ns=0, + ) + + request_info = mock_worker._create_request_info( + x_request_id="request-id", + credit_context=credit_context, + turns=[Turn()], + ) + + assert request_info.is_parent_final is True + assert request_info.is_tree_final is True + async def test_process_response( self, monkeypatch, mock_worker, sample_request_record ): @@ -557,3 +586,199 @@ async def emit_boom(*args, **kwargs): assert sample_credit_context.returned is True assert sample_credit_context.cancelled is False credit_send.assert_awaited_once() + + +# --- Terminal Eviction / Session-Routing Hook Tests --- + + +_OVERFLOW_BODY = "This model's maximum context length is 8192 tokens" + + +class TestTerminalContextOverflowClassifier: + """``_is_terminal_context_overflow``: a context-overflow error on a + non-final, non-cancelled turn is terminal (agentic_replay recycles the lane + and sends no final/cancel credit). Final-turn / cancelled returns go through + the normal eviction path and must NOT be classified as overflow-terminal.""" + + @pytest.mark.parametrize( + "is_final, cancelled, error, expected", + [ + param(False, False, _OVERFLOW_BODY, True, id="nonfinal-overflow"), + param(False, False, "connection reset by peer", False, id="nonfinal-plain-error"), + param(False, False, None, False, id="nonfinal-no-error"), + param(True, False, _OVERFLOW_BODY, False, id="final-turn"), + param(False, True, _OVERFLOW_BODY, False, id="cancelled"), + ], + ) # fmt: skip + def test_classifier(self, is_final, cancelled, error, expected) -> None: + credit = MagicMock(is_final_turn=is_final) + ctx = MagicMock(cancelled=cancelled, error=error) + assert _is_terminal_context_overflow(credit, ctx) is expected + + +@pytest.mark.asyncio +class TestSessionRoutingTerminalHooks: + """Every terminal disposition reaches ``InferenceClient.notify_session_end`` + so routing plugins get their idempotent post-session hook on ALL four + terminal paths: final turn and cancellation (the finally block of + ``Worker._process_credit``, where notify precedes session eviction), + terminal context overflow (same finally block, elif branch), and the + cancel-before-start branch of ``_on_credit_drop_message_task_done`` (which + that finally block never sees). + """ + + def _make_context(self, *, num_turns: int) -> CreditContext: + return CreditContext( + credit=Credit( + id=7, + phase=CreditPhase.PROFILING, + conversation_id="conv-template", + x_correlation_id="conv-X", + turn_index=0, + num_turns=num_turns, + issued_at_ns=0, + ), + drop_perf_ns=0, + ) + + async def _drive_process_credit( + self, mock_worker, monkeypatch, ctx: CreditContext, *, outcome: str + ) -> Mock: + """Drive the REAL ``_on_credit_drop_message_task`` -> ``_process_credit`` + path (finally block included) with the session step stubbed to produce + the given outcome; return the notify_session_end mock.""" + mock_worker._is_payload_bytes = False + + async def behavior(*args, **kwargs): + if outcome == "cancel": + raise asyncio.CancelledError() + if outcome == "overflow": + # Mirror _execute_request propagating an error record's body + # onto the credit context (worker returns the error upstream). + ctx.error = ErrorDetails(message=_OVERFLOW_BODY) + elif outcome == "plain-error": + raise ValueError("connection reset by peer") + + monkeypatch.setattr(mock_worker, "_process_credit_with_session", behavior) + monkeypatch.setattr(mock_worker.credit_dealer_client, "send", AsyncMock()) + monkeypatch.setattr(mock_worker, "_send_inference_result_message", AsyncMock()) + notify = Mock() + monkeypatch.setattr(mock_worker.inference_client, "notify_session_end", notify) + + await mock_worker._on_credit_drop_message_task(ctx) + return notify + + async def test_final_turn_notifies_session_end(self, mock_worker, monkeypatch): + ctx = self._make_context(num_turns=1) + notify = await self._drive_process_credit( + mock_worker, monkeypatch, ctx, outcome="ok" + ) + notify.assert_called_once_with("conv-X") + + async def test_cancelled_notifies_session_end(self, mock_worker, monkeypatch): + ctx = self._make_context(num_turns=3) + notify = await self._drive_process_credit( + mock_worker, monkeypatch, ctx, outcome="cancel" + ) + assert ctx.cancelled is True + notify.assert_called_once_with("conv-X") + + async def test_terminal_context_overflow_notifies_session_end( + self, mock_worker, monkeypatch + ): + """Under agentic_replay, a non-final turn whose error is a context + overflow ends the session without any final/cancel credit (the lane is + recycled), so the finally block's elif branch must fire the hook.""" + mock_worker._overflow_ends_session = True + ctx = self._make_context(num_turns=3) + notify = await self._drive_process_credit( + mock_worker, monkeypatch, ctx, outcome="overflow" + ) + notify.assert_called_once_with("conv-X") + + async def test_context_overflow_does_not_notify_under_non_agentic_strategies( + self, mock_worker, monkeypatch + ): + """request_rate / user_centric strategies keep issuing a session's + remaining turns after an overflow error, so the overflow branch must + NOT fire the post-session hook mid-session there. The mock_worker + fixture's default timing mode is non-agentic.""" + assert mock_worker._overflow_ends_session is False + ctx = self._make_context(num_turns=3) + notify = await self._drive_process_credit( + mock_worker, monkeypatch, ctx, outcome="overflow" + ) + notify.assert_not_called() + + async def test_nonfinal_plain_error_does_not_notify(self, mock_worker, monkeypatch): + """A non-final, non-overflow error is not terminal: the session may + still receive its remaining turns, so the hook must NOT fire.""" + ctx = self._make_context(num_turns=3) + notify = await self._drive_process_credit( + mock_worker, monkeypatch, ctx, outcome="plain-error" + ) + notify.assert_not_called() + + def _done_callback( + self, *, returned: bool, task_cancelled: bool = True + ) -> 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.cancelled = False + ctx.credit = credit + task = MagicMock() + task.cancelled.return_value = task_cancelled + Worker._on_credit_drop_message_task_done.__get__(worker)(task, ctx) + return worker + + async 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") + + async 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() + + async def test_done_callback_uncancelled_send_failure_does_not_notify( + self, + ) -> None: + """A task that ran to completion but whose CreditReturn send raised + reaches this branch with returned=False and cancelled=False. That is + NOT a terminal disposition for a non-final turn, so the hook must not + fire (for final/cancelled turns the finally block already notified).""" + worker = self._done_callback(returned=False, task_cancelled=False) + worker.inference_client.notify_session_end.assert_not_called() + + async def test_worker_stop_sweeps_remaining_sessions( + self, mock_worker, monkeypatch + ) -> None: + """Sessions still resident at shutdown (mid-conversation at phase end, + no in-flight credit to cancel) get their post-session hook fired by + the stop sweep.""" + notify = Mock() + monkeypatch.setattr(mock_worker.inference_client, "notify_session_end", notify) + mock_worker.session_manager._cache["conv-A"] = MagicMock() + mock_worker.session_manager._cache["conv-B"] = MagicMock() + + await mock_worker._worker_stop() + + assert notify.call_count == 2 + notify.assert_any_call("conv-A") + notify.assert_any_call("conv-B") From 074a67d871eb3d7cd004ccde844cc2b34044dee8 Mon Sep 17 00:00:00 2001 From: Anthony Casagrande Date: Wed, 8 Jul 2026 23:25:40 -0700 Subject: [PATCH 2/2] feat(routing): lineage-scoped affinity for dynamo_nvext Adds scope={conversation,lineage} to the dynamo_nvext mode (--session-routing-opt scope=lineage; default conversation, unchanged behavior). Under lineage, every session in an agent tree binds nvext.session_control with the tree ROOT's correlation ID instead of its own, so fork/spawn children co-locate on the worker (and data-parallel rank) already holding the shared parent prefix -- an explicit client-side affinity for Dynamo deployments without KV-event prefix indexing, where an unbound child request is otherwise placed arbitrarily and re-prefills the shared parent prefix once per child. Close discipline: a shared affinity key must never be torn down while sibling sessions may still run (a straggler's next bind would re-place it arbitrarily and lose co-location), so under lineage the close fires only on a request the credit issuer stamped is_tree_final -- provably the last request of the whole tree, available under agentic replay, where the has_branches gate guarantees a branch-declaring final turn never reads as tree-final while children are pending. Under indeterminate modes (reactive dag_jsonl) nothing closes and the session_control TTL reclaims the key. Lineage-scope design and motivation from SemiAnalysisAI/aiperf#14 (YAMY1234), rebased onto the session-routing plugin layer: the flag surface there (--dynamo-session-affinity-scope + EndpointDefaults + EndpointInfo plumbing) collapses into this one Options field, and the root-final-turn close is replaced by the is_tree_final discipline. Verified with back-to-back --export-level raw mock-server runs of --scenario inferencex-agentx-mvp against the PR #14 implementation: child binds are wire-identical (keyed on the tree-root ID in both), and the lineage close fired exactly once -- on the provably-last wire request of a fully-drained tree. Co-Authored-By: YAMY1234 Co-Authored-By: Claude Fable 5 Signed-off-by: Anthony Casagrande --- docs/plugins/plugin-system.md | 2 +- src/aiperf/plugin/plugins.yaml | 7 +++- src/aiperf/workers/session_routing.py | 34 ++++++++++++++-- tests/unit/workers/test_session_routing.py | 46 ++++++++++++++++++++++ 4 files changed, 83 insertions(+), 6 deletions(-) diff --git a/docs/plugins/plugin-system.md b/docs/plugins/plugin-system.md index b5af198608..82869b2abd 100644 --- a/docs/plugins/plugin-system.md +++ b/docs/plugins/plugin-system.md @@ -146,7 +146,7 @@ the request-serialization chokepoint. The base class is `SessionRoutingBase` | 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). | +| `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`. Options: `timeout_seconds` (default 300); `scope` (default `conversation`; `lineage` binds every session in an agent tree with the tree root's correlation ID for whole-lineage co-location on deployments without KV-event prefix indexing, closing the shared key only on a request stamped provably-last for the whole tree). | | `smg_routing_key` | `SmgRoutingKeyRouting` | SGLang Model Gateway `manual`-policy stickiness via `X-SMG-Routing-Key`. No options. | | `session_id_header` | `SessionIdHeaderRouting` | Preset: additive `X-Session-ID` header carrying the session's correlation ID. No options. | | `identity_headers` | `IdentityHeadersRouting` | Fully generic tiered identity headers: emits exactly what you configure, nothing by default. Options (each a comma-separable name list): `session` (this session's ID), `parent` (immediate parent's ID; omitted for roots), `root` (session-tree root's ID, for whole-tree affinity). Requires at least one name across tiers; names must be unique (case-insensitive) across all tiers combined, including within a single tier. | diff --git a/src/aiperf/plugin/plugins.yaml b/src/aiperf/plugin/plugins.yaml index 6c3c55474e..230aa767a5 100644 --- a/src/aiperf/plugin/plugins.yaml +++ b/src/aiperf/plugin/plugins.yaml @@ -449,7 +449,12 @@ session_routing: Only for Dynamo builds that implement session_control (current upstream main does not; use dynamo_headers there). Mutates request bodies, so it is incompatible with the PAYLOAD_BYTES mmap fast path. - Option: timeout_seconds (default 300). + Options: timeout_seconds (default 300); scope (default + conversation) -- scope=lineage binds every session in an agent + tree with the tree ROOT's correlation ID so the whole lineage + co-locates with the shared parent prefix (for deployments without + KV-event prefix indexing), closing the shared key only on a + request stamped provably-last for the whole tree. smg_routing_key: class: aiperf.workers.session_routing:SmgRoutingKeyRouting diff --git a/src/aiperf/workers/session_routing.py b/src/aiperf/workers/session_routing.py index c20748ecb0..a2e9a53382 100644 --- a/src/aiperf/workers/session_routing.py +++ b/src/aiperf/workers/session_routing.py @@ -27,7 +27,7 @@ import re from abc import ABC -from typing import Annotated, Any, ClassVar, Generic, TypeVar +from typing import Annotated, Any, ClassVar, Generic, Literal, TypeVar from msgspec import Struct from pydantic import BeforeValidator, ConfigDict, Field, model_validator @@ -141,6 +141,17 @@ class DynamoNvextOptions(AIPerfBaseModel): ge=1, description="Dynamo session_control inactivity timeout carried on every bind.", ) + scope: Literal["conversation", "lineage"] = Field( + default="conversation", + description="Affinity-key scope. 'conversation' (default) binds each " + "session with its own correlation ID and closes on its final turn. " + "'lineage' binds every session in an agent tree with the tree ROOT's " + "correlation ID so the whole lineage co-locates on the worker holding " + "the shared parent prefix (for Dynamo deployments without KV-event " + "prefix indexing); the shared key is closed only on a request stamped " + "provably-last for the whole tree (is_tree_final, agentic replay) -- " + "otherwise the session_control TTL reclaims it.", + ) class DynamoNvextRouting(SessionRoutingBase[DynamoNvextOptions]): @@ -150,6 +161,14 @@ class DynamoNvextRouting(SessionRoutingBase[DynamoNvextOptions]): 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). + + Under ``scope=lineage`` the affinity key is the tree root's correlation ID + and the close discipline changes: a shared key must never be torn down + while sibling sessions may still run (a later bind would re-place the + straggler arbitrarily and lose co-location), so close fires only on a + request the issuer stamped ``is_tree_final`` -- conservative-exact under + agentic replay, never under indeterminate modes, where the TTL reaper + reclaims the key instead. """ mutates_body: ClassVar[bool] = True @@ -160,18 +179,25 @@ def __init__(self, options: DynamoNvextOptions) -> None: # Hot path: one attribute hop per request instead of two through the # pydantic options model. self._timeout = options.timeout_seconds + self._lineage = options.scope == "lineage" def transform_body( self, payload: dict[str, Any], ctx: RoutingContext ) -> dict[str, Any]: - if ctx.is_final_turn: + if self._lineage: + session_id = ctx.root_correlation_id or ctx.x_correlation_id + is_close = ctx.is_tree_final + else: + session_id = ctx.x_correlation_id + is_close = ctx.is_final_turn + if is_close: session_control: dict[str, Any] = { - "session_id": ctx.x_correlation_id, + "session_id": session_id, "action": "close", } else: session_control = { - "session_id": ctx.x_correlation_id, + "session_id": session_id, "action": "bind", "timeout": self._timeout, } diff --git a/tests/unit/workers/test_session_routing.py b/tests/unit/workers/test_session_routing.py index b4706ba981..d11c74451e 100644 --- a/tests/unit/workers/test_session_routing.py +++ b/tests/unit/workers/test_session_routing.py @@ -129,6 +129,52 @@ def test_non_dict_nvext_replaced(self): merged = plugin.transform_body({"nvext": "bogus"}, _ctx()) assert merged["nvext"]["session_control"]["session_id"] == "corr-1" + def test_lineage_scope_child_binds_with_root_id(self): + """Under scope=lineage, every session in the tree binds the ROOT's + correlation ID so the whole lineage co-locates.""" + plugin = DynamoNvextRouting( + DynamoNvextOptions(scope="lineage", timeout_seconds=77) + ) + ctx = _ctx( + parent_correlation_id="root-1", + root_correlation_id="root-1", + is_final_turn=True, # child's own final turn must NOT close + ) + sc = plugin.transform_body({}, ctx)["nvext"]["session_control"] + assert sc == {"session_id": "root-1", "action": "bind", "timeout": 77} + + def test_lineage_scope_root_final_turn_binds_not_closes(self): + """A shared key must never be torn down while siblings may still run: + even the root's own final turn only binds unless the issuer stamped + the request provably-last for the whole tree.""" + plugin = DynamoNvextRouting(DynamoNvextOptions(scope="lineage")) + sc = plugin.transform_body({}, _ctx(is_final_turn=True))["nvext"][ + "session_control" + ] + assert sc["action"] == "bind" + assert sc["session_id"] == "corr-1" + + def test_lineage_scope_closes_only_on_tree_final(self): + plugin = DynamoNvextRouting(DynamoNvextOptions(scope="lineage")) + ctx = _ctx( + root_correlation_id="root-1", + is_final_turn=True, + is_tree_final=True, + ) + sc = plugin.transform_body({}, ctx)["nvext"]["session_control"] + assert sc == {"session_id": "root-1", "action": "close"} + + def test_conversation_scope_default_unchanged(self): + """Default scope keeps per-conversation identity and final-turn close.""" + plugin = DynamoNvextRouting(DynamoNvextOptions()) + assert plugin.transform_body({}, _ctx(is_final_turn=True))["nvext"][ + "session_control" + ] == {"session_id": "corr-1", "action": "close"} + + def test_invalid_scope_rejected(self): + with pytest.raises(ValidationError): + DynamoNvextOptions(scope="tree") + class TestSmgRoutingKeyRouting: def test_emits_routing_key(self):