diff --git a/src/aiperf/common/models/record_models.py b/src/aiperf/common/models/record_models.py index 29877a5395..19d296a79a 100644 --- a/src/aiperf/common/models/record_models.py +++ b/src/aiperf/common/models/record_models.py @@ -203,6 +203,13 @@ class MetricRecordMetadata(AIPerfBaseModel): description="The x_correlation_id of the parent session that spawned this record's session via a " "DAG subagent fork. None for root sessions. Use to group sibling branches of the same DAG.", ) + root_correlation_id: str | None = Field( + default=None, + description="The x_correlation_id of the depth-0 root of this record's session TREE. Stable " + "across the whole tree (root + every descendant subagent at any depth); equals x_correlation_id " + "for a root session. Groups every record of one agentic session (root + subagents) under a single " + "lane and lets analysis reconstruct exactly-N session-tree concurrency.", + ) class TimesliceResult(AIPerfBaseModel): @@ -661,6 +668,14 @@ class RecordContext(AIPerfBaseModel): description="The x_correlation_id of the parent session that spawned this session via a DAG " "subagent fork. None for root sessions. Sourced from the originating Credit.", ) + root_correlation_id: str | None = Field( + default=None, + description="The x_correlation_id of the depth-0 root of this record's session TREE. " + "Stable across the whole tree (root + every descendant subagent at any depth); equals " + "x_correlation_id for a root session. Sourced from the originating Credit. Use to group " + "every record of one agentic session (root + subagents) and to reconstruct exactly-N " + "session-tree concurrency in analysis.", + ) # --- Hoisted metric inputs (avoid shipping full Turn structs) ------------- @@ -683,24 +698,6 @@ class RecordContext(AIPerfBaseModel): "the full ``turns`` list on the wire. None for non-ASR requests.", ) - # --- Records-pipeline reads (read by inference_result_parser, raw_record_writer) ---- - - turns: list[Turn] = Field( - default_factory=list, - description="The actual turns of the request. This will include assistant turns as well as user turns in multi-turn conversations. " - "Read by the records pipeline (``inference_result_parser``, ``raw_record_writer_processor``) for response parsing and raw export.", - ) - system_message: str | None = Field( - default=None, - description="Optional shared system message to prepend to the first turn. " - "Extracted from conversation.system_message at request time. Read by the records pipeline.", - ) - user_context_message: str | None = Field( - default=None, - description="Optional per-conversation user context message to prepend to the first turn. " - "Extracted from conversation.user_context_message at request time. Read by the records pipeline.", - ) - class RequestInfo(RecordContext): """Full request info used Worker-side for transport dispatch. @@ -718,6 +715,29 @@ class RequestInfo(RecordContext): ..., description="The model endpoint that the request was sent to.", ) + turns: list[Turn] = Field( + default_factory=list, + description="The actual turns of the request, consumed by " + "``format_payload`` to build the wire body. Lives on ``RequestInfo`` " + "(not ``RecordContext``) so the full Turn list never crosses the " + "ZMQ hop to the record processor — only the canonical " + "``payload_bytes`` travel.", + ) + system_message: str | None = Field( + default=None, + description="Optional shared system message extracted from " + "``Conversation.system_message`` at request time. Consumed by the " + "endpoint's ``format_payload`` (or top-level ``instructions`` on the " + "Responses API) and inlined into ``payload_bytes`` before transport; " + "lives on ``RequestInfo`` because the record processor reads only " + "``payload_bytes`` downstream.", + ) + user_context_message: str | None = Field( + default=None, + description="Optional per-conversation user context message extracted " + "from ``Conversation.user_context_message`` at request time. Same " + "inlining contract as ``system_message``.", + ) endpoint_headers: dict[str, str] = Field( default_factory=dict, description="Endpoint-specific headers (auth, API keys, custom headers).", @@ -820,11 +840,6 @@ class RequestRecord(AIPerfBaseModel): "Includes detailed timing for connection establishment, DNS resolution, request/response events, etc. " "The type of the trace data is determined by the transport and library used.", ) - turns: list[Turn] = Field( - default_factory=list, - description="Deep copy of the request turns. This is a copy of the turns from request_info, " - "made to avoid mutating the original session data when stripping multimodal content.", - ) @field_validator("trace_data", mode="before") @classmethod diff --git a/src/aiperf/records/record_processor_service.py b/src/aiperf/records/record_processor_service.py index 305a42f7d1..b841e80c3b 100644 --- a/src/aiperf/records/record_processor_service.py +++ b/src/aiperf/records/record_processor_service.py @@ -206,6 +206,7 @@ def _create_metric_record_metadata( cancellation_time_ns=cancellation_time_ns, agent_depth=record.request_info.agent_depth, parent_correlation_id=record.request_info.parent_correlation_id, + root_correlation_id=record.request_info.root_correlation_id, ) @on_pull_message(MessageType.INFERENCE_RESULTS) @@ -349,7 +350,7 @@ def _free_record_data( The only data sent downstream in MetricRecordsMessage is metadata, results, trace_data, and error -- so everything else can be released here. - We assign None to fields typed as non-optional lists (turns, responses) to let + We assign None to fields typed as non-optional lists (responses) to let the GC reclaim the underlying objects. Using .clear() would keep the empty list alive, and reassigning [] would allocate a new object for no reason. """ @@ -357,13 +358,8 @@ def _free_record_data( error = record.error if self.run.cfg.artifacts.export_level != ExportLevel.RAW: record.responses = None - record.turns = None record.trace_data = None record.request_headers = None - if record.request_info: - record.request_info.turns = None - record.request_info.system_message = None - record.request_info.user_context_message = None parsed_record.responses = None return trace_data, error diff --git a/src/aiperf/workers/inference_client.py b/src/aiperf/workers/inference_client.py index 27d2098edb..027154e83a 100644 --- a/src/aiperf/workers/inference_client.py +++ b/src/aiperf/workers/inference_client.py @@ -206,10 +206,12 @@ def _enrich_request_record( the record before the ZMQ hop to the record processor. The full ``RequestInfo`` carries transport-only extras - (``model_endpoint``, ``turns``, ``endpoint_headers``, - ``endpoint_params``, ``drop_perf_ns``, ``cancel_after_ns``, ...) that - the record-processor pipeline never reads; downcasting saves - ~500-900 bytes per record at high throughput. + (``model_endpoint``, ``turns``, ``system_message``, + ``user_context_message``, ``endpoint_headers``, ``endpoint_params``, + ``drop_perf_ns``, ``cancel_after_ns``, ...) that the record-processor + pipeline never reads; downcasting saves ~500-900 bytes per record at + high throughput. The full ``turns`` list never travels — live records + drive off the canonical ``payload_bytes``. """ ctx_field_names = set(RecordContext.model_fields.keys()) ri_dump = request_info.model_dump(include=ctx_field_names) @@ -244,10 +246,6 @@ def _finalize_request_record( if self.strip_record_payload_bytes and record.request_info is not None: record.request_info.payload_bytes = None - # Copy turns with stripped multimodal data to avoid mutating original session - # and reduce memory usage (placeholders instead of large image/audio/video data) - record.turns = [turn.copy_with_stripped_media() for turn in request_info.turns] - # If this is the first turn, calculate the credit drop latency if request_info.turn_index == 0 and request_info.drop_perf_ns is not None: record.credit_drop_latency = ( diff --git a/tests/unit/common/models/test_record_context.py b/tests/unit/common/models/test_record_context.py index 61f2e5c19d..bf27248462 100644 --- a/tests/unit/common/models/test_record_context.py +++ b/tests/unit/common/models/test_record_context.py @@ -48,14 +48,18 @@ def test_request_info_has_transport_extras(self): ri_fields = set(RequestInfo.model_fields.keys()) ctx_fields = set(RecordContext.model_fields.keys()) extras = ri_fields - ctx_fields - # ``turns``, ``system_message``, ``user_context_message`` were hoisted - # onto RecordContext because the records pipeline reads them - # post-transport. The remaining transport-only extras are the - # endpoint/URL/timing fields. + # ``turns``, ``system_message``, ``user_context_message`` live ONLY on + # RequestInfo (worker-side) so the full Turn list never crosses the ZMQ + # hop to the record processor — only the canonical ``payload_bytes`` (a + # RecordContext field) travels. They are transport-only extras here. assert {"model_endpoint", "endpoint_headers", "drop_perf_ns"}.issubset(extras) - assert "turns" not in extras - assert "system_message" not in extras - assert "user_context_message" not in extras + assert "turns" in extras + assert "system_message" in extras + assert "user_context_message" in extras + # The hoisted scalars stay on RecordContext (they cross the wire). + assert "max_tokens" not in extras + assert "audio_duration_seconds" not in extras + assert "payload_bytes" not in extras class TestRequestRecordHoldsRecordContext: diff --git a/tests/unit/workers/test_inference_client_enrich_record.py b/tests/unit/workers/test_inference_client_enrich_record.py index fe36b7a340..31d6370145 100644 --- a/tests/unit/workers/test_inference_client_enrich_record.py +++ b/tests/unit/workers/test_inference_client_enrich_record.py @@ -48,12 +48,15 @@ def test_transport_extras_dropped(self): ri = _make_request_info() record = RequestRecord() enriched = InferenceClient._enrich_request_record(record, ri) - # downcast strips model_endpoint / endpoint_headers / drop_perf_ns, - # but ``turns``, ``system_message``, ``user_context_message`` were - # hoisted onto RecordContext (records pipeline reads them) and - # therefore survive the downcast. + # Downcast strips RequestInfo-only transport/payload-builder fields while + # keeping the hoisted scalar inputs the records pipeline reads. dump = enriched.request_info.model_dump() assert "model_endpoint" not in dump assert "endpoint_headers" not in dump assert "drop_perf_ns" not in dump - assert "turns" in dump + assert "turns" not in dump + assert "system_message" not in dump + assert "user_context_message" not in dump + assert "payload_bytes" in dump + assert "max_tokens" in dump + assert "audio_duration_seconds" in dump