From 63ee21b6ecf485df6de23b788c9aeca68c6b4b97 Mon Sep 17 00:00:00 2001 From: Ishan Dhanani Date: Thu, 9 Jul 2026 12:16:46 +0000 Subject: [PATCH] feat: support session IDs in request bodies Signed-off-by: Ishan Dhanani --- docs/benchmark-datasets.md | 2 ++ docs/cli-options.md | 8 +++++++ .../common/models/model_endpoint_info.py | 6 +++++ src/aiperf/config/endpoint.py | 11 ++++++++++ .../config/flags/_converter_endpoint.py | 1 + src/aiperf/config/flags/_section_fields.py | 1 + src/aiperf/config/flags/cli_config.py | 16 ++++++++++++++ .../config/schema/aiperf-config.schema.json | 13 +++++++++++ src/aiperf/transports/aiohttp_transport.py | 7 ++++++ .../config/test_converter_endpoint_session.py | 17 ++++++++++++++ .../unit/transports/test_aiohttp_transport.py | 22 +++++++++++++++++++ 11 files changed, 104 insertions(+) create mode 100644 tests/unit/config/test_converter_endpoint_session.py diff --git a/docs/benchmark-datasets.md b/docs/benchmark-datasets.md index f0d2f3f4e4..94b35aaf62 100644 --- a/docs/benchmark-datasets.md +++ b/docs/benchmark-datasets.md @@ -86,6 +86,8 @@ aiperf profile \ --fixed-schedule ``` +For an endpoint such as SGLang that expects the session identifier in the request body, add `--session-header x-dynamo-session-id --session-body-field session_id`. This copies each recorded Exgentic session ID into the top-level `session_id` field without changing the dataset. + `source_model` selects the model that produced the trace; `--model` selects the target model receiving the replay. `benchmark` selects an Exgentic v2 workload. Invalid filters report the available harness/model combinations. The v1 dataset contains 22 combinations across five harnesses and six canonical source models. Fixed-schedule mode emits each recorded call as an independently scheduled one-turn request using its start offset from the source session. Calls that overlapped in the trace therefore overlap during replay. Selected source sessions start together at offset zero. Without `--fixed-schedule`, each source session remains a closed-loop multi-turn conversation: AIPerf waits for one live response before applying the recorded residual delay and sending the next request. diff --git a/docs/cli-options.md b/docs/cli-options.md index ca7c620a3c..9ce404457a 100644 --- a/docs/cli-options.md +++ b/docs/cli-options.md @@ -379,6 +379,10 @@ Content type for request body serialization. By default, requests are sent as 'a HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`). +#### `--session-body-field` `` + +Optional request body field used to carry the same per-session affinity identifier as `--session-header` (or `X-Correlation-ID` by default). For example, use `--session-body-field session_id` for SGLang session-aware requests. + ### Tokenizer #### `--tokenizer` `` @@ -1786,6 +1790,10 @@ Content type for request body serialization. By default, requests are sent as 'a HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`). +#### `--session-body-field` `` + +Optional request body field used to carry the same per-session affinity identifier as `--session-header` (or `X-Correlation-ID` by default). For example, use `--session-body-field session_id` for SGLang session-aware requests. + ### Tokenizer #### `--tokenizer` `` diff --git a/src/aiperf/common/models/model_endpoint_info.py b/src/aiperf/common/models/model_endpoint_info.py index e5e98bb62b..793b7b28f4 100644 --- a/src/aiperf/common/models/model_endpoint_info.py +++ b/src/aiperf/common/models/model_endpoint_info.py @@ -139,6 +139,11 @@ def _redact_headers(self, value: list[tuple[str, str]]) -> list[tuple[str, str]] description="HTTP header name to use for the per-session affinity identifier. " "When set, replaces the default `X-Correlation-ID` header name with this value.", ) + session_body_field: str | None = Field( + default=None, + description="Optional request body field used to carry the same per-session " + "affinity identifier as the configured session header.", + ) collect_trace_chunks: bool = Field( default=False, description="Collect per-chunk trace data (timestamps and sizes) for HTTP trace export. " @@ -213,6 +218,7 @@ def from_run(cls, run: BenchmarkRun) -> ModelEndpointInfo: collect_trace_chunks=False, template=getattr(ep, "template", None), session_header=getattr(ep, "session_header", None), + session_body_field=getattr(ep, "session_body_field", None), ), transport=ep.transport, ) diff --git a/src/aiperf/config/endpoint.py b/src/aiperf/config/endpoint.py index a5a8fee1e0..a1a3919c38 100644 --- a/src/aiperf/config/endpoint.py +++ b/src/aiperf/config/endpoint.py @@ -316,6 +316,17 @@ def _redact_headers(self, value: dict[str, str]) -> dict[str, str]: ), ] + session_body_field: Annotated[ + str | None, + Field( + default=None, + description=( + "Optional request body field used to carry the same per-session " + "affinity identifier as the configured session header." + ), + ), + ] + wait_for_model_timeout: Annotated[ float, Field( diff --git a/src/aiperf/config/flags/_converter_endpoint.py b/src/aiperf/config/flags/_converter_endpoint.py index afe4773c61..4cb50ee786 100644 --- a/src/aiperf/config/flags/_converter_endpoint.py +++ b/src/aiperf/config/flags/_converter_endpoint.py @@ -72,6 +72,7 @@ def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None: "download_video_content": "download_video_content", "request_content_type": "request_content_type", "session_header": "session_header", + "session_body_field": "session_body_field", } diff --git a/src/aiperf/config/flags/_section_fields.py b/src/aiperf/config/flags/_section_fields.py index ec2cefd831..d331c8be03 100644 --- a/src/aiperf/config/flags/_section_fields.py +++ b/src/aiperf/config/flags/_section_fields.py @@ -23,6 +23,7 @@ "model_names", "model_selection_strategy", "request_content_type", + "session_body_field", "session_header", "streaming", "timeout_seconds", diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py index cd790bc571..a79cb21f5a 100644 --- a/src/aiperf/config/flags/cli_config.py +++ b/src/aiperf/config/flags/cli_config.py @@ -389,6 +389,22 @@ class CLIConfig(BaseConfig): ), ] = None + session_body_field: Annotated[ + str | None, + Field( + description=( + "Optional request body field used to carry the same per-session " + "affinity identifier as `--session-header` (or `X-Correlation-ID` " + "by default). For example, use `--session-body-field session_id` " + "for SGLang session-aware requests." + ), + ), + CLIParameter( + name=("--session-body-field",), + group=Groups.ENDPOINT, + ), + ] = None + @property def url(self) -> str: """Return the first URL for backward compatibility.""" diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json index 98f43fdd30..8b54c834ce 100644 --- a/src/aiperf/config/schema/aiperf-config.schema.json +++ b/src/aiperf/config/schema/aiperf-config.schema.json @@ -7214,6 +7214,19 @@ "description": "HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header. Useful when the inference server expects a custom session-affinity header (e.g. `--session-header X-Session-ID`).", "title": "Sessionheader" }, + "sessionBodyField": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional request body field used to carry the same per-session affinity identifier as the configured session header.", + "title": "Sessionbodyfield" + }, "waitForModelTimeout": { "default": 0.0, "description": "Enable a pre-flight readiness probe by setting this to a positive value (seconds). aiperf applies this timeout to each URL/model probe before starting the benchmark, aborting with a non-zero exit if any probe times out. For multiple URLs or models, worst-case wall-clock time can be roughly this timeout multiplied by the number of URL/model probes. The probe strategy is controlled by `--wait-for-model-mode`, which defaults to sending a 1-token inference request. 0 (default) disables the probe. Eliminates the need for external shell-based readiness loops in containers and Kubernetes recipes.", diff --git a/src/aiperf/transports/aiohttp_transport.py b/src/aiperf/transports/aiohttp_transport.py index bfebab23c7..7128b3eb0f 100644 --- a/src/aiperf/transports/aiohttp_transport.py +++ b/src/aiperf/transports/aiohttp_transport.py @@ -319,6 +319,13 @@ async def send_request( try: url = self.build_url(request_info) headers = self.build_headers(request_info) + session_body_field = request_info.model_endpoint.endpoint.session_body_field + if session_body_field: + session_header = ( + request_info.model_endpoint.endpoint.session_header + or "X-Correlation-ID" + ) + payload[session_body_field] = headers[session_header] use_form_data = ( request_info.model_endpoint.endpoint.request_content_type == RequestContentType.MULTIPART_FORM_DATA diff --git a/tests/unit/config/test_converter_endpoint_session.py b/tests/unit/config/test_converter_endpoint_session.py new file mode 100644 index 0000000000..e5bfe51676 --- /dev/null +++ b/tests/unit/config/test_converter_endpoint_session.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from aiperf.config.flags._converter_endpoint import build_endpoint +from aiperf.config.flags.cli_config import CLIConfig + + +def test_session_affinity_fields_are_added_to_endpoint() -> None: + endpoint = build_endpoint( + CLIConfig( + session_header="x-dynamo-session-id", + session_body_field="session_id", + ) + ) + + assert endpoint["session_header"] == "x-dynamo-session-id" + assert endpoint["session_body_field"] == "session_id" diff --git a/tests/unit/transports/test_aiohttp_transport.py b/tests/unit/transports/test_aiohttp_transport.py index 8102edc917..8ad6be2182 100644 --- a/tests/unit/transports/test_aiohttp_transport.py +++ b/tests/unit/transports/test_aiohttp_transport.py @@ -5,9 +5,11 @@ from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock, patch +import orjson import pytest from aiperf.common.enums import ConnectionReuseStrategy, CreditPhase +from aiperf.common.models import Turn from aiperf.common.models.record_models import RequestInfo, RequestRecord from aiperf.plugin import plugins from aiperf.plugin.enums import TransportType @@ -297,6 +299,26 @@ async def test_send_request_serializes_payload_with_orjson( assert b"messages" in json_bytes assert b"gpt-4" in json_bytes + @pytest.mark.asyncio + async def test_send_request_copies_resolved_session_header_to_body(self): + model_endpoint = create_model_endpoint_info() + model_endpoint.endpoint.session_header = "x-dynamo-session-id" + model_endpoint.endpoint.session_body_field = "session_id" + transport = AioHttpTransport(model_endpoint=model_endpoint) + await self._setup_initialized_transport_with_mock(transport) + + request_info = create_request_info(model_endpoint) + request_info.turns = [ + Turn(extra_headers={"x-dynamo-session-id": "source-session-1"}) + ] + payload = {"model": "test-model", "messages": []} + + await transport.send_request(request_info, payload) + + args = self._extract_call_args(transport.aiohttp_client.post_request.call_args) + assert args["headers"]["x-dynamo-session-id"] == "source-session-1" + assert orjson.loads(args["json_bytes"])["session_id"] == "source-session-1" + @pytest.mark.asyncio async def test_send_request_handles_exception( self, transport, model_endpoint_non_streaming