diff --git a/src/aiperf/dataset/loader/burst_gpt.py b/src/aiperf/dataset/loader/burst_gpt.py index 570d84ef71..a72f835419 100644 --- a/src/aiperf/dataset/loader/burst_gpt.py +++ b/src/aiperf/dataset/loader/burst_gpt.py @@ -102,8 +102,6 @@ def load_dataset(self) -> dict[str, list[BurstGPTTrace]]: items.append(trace) - self._log_filtering_summary() - data = self._group_traces(items) self.debug( lambda: ( @@ -115,6 +113,15 @@ def load_dataset(self) -> dict[str, list[BurstGPTTrace]]: if _has_meaningful_synthesis(self._synthesis): data = self._apply_synthesis(data) + # Apply --synthesis-max-osl AFTER synthesis (matches the base + # load_dataset). This override otherwise returned uncapped output + # lengths -- the base-loader refactor moved the cap out of + # _filter_and_cap_trace into this grouped pass, which overriding + # loaders must call explicitly. Log AFTER the cap so the capped count + # is reflected in the summary (matches the base order). + data = self._cap_grouped_traces_max_osl(data) + self._log_filtering_summary() + return data # ------------------------------------------------------------------ diff --git a/src/aiperf/dataset/loader/sagemaker_data_capture.py b/src/aiperf/dataset/loader/sagemaker_data_capture.py index 5843fde65d..e65b4991e6 100644 --- a/src/aiperf/dataset/loader/sagemaker_data_capture.py +++ b/src/aiperf/dataset/loader/sagemaker_data_capture.py @@ -15,6 +15,7 @@ from __future__ import annotations import base64 +import binascii from datetime import datetime from pathlib import Path from typing import Any @@ -36,20 +37,29 @@ def _parse_iso8601_to_ms(iso_str: str) -> float: return dt.timestamp() * 1000.0 -def _decode_payload(capture_entry: dict[str, Any]) -> dict[str, Any] | None: +def _decode_payload( + capture_entry: dict[str, Any], event_id: str = "" +) -> dict[str, Any] | None: """Decode a captureData entry (endpointInput or endpointOutput). Handles JSON (raw) and BASE64 encoded payloads. Returns None for - CSV or unknown encodings. + CSV or unknown encodings. A malformed base64 or inner-JSON payload is + re-raised as DatasetLoaderError so one bad record does not abort the + whole load (matching the per-record attribution in _parse_trace). """ data = capture_entry.get("data") if data is None: return None encoding = capture_entry.get("encoding", "BASE64") - if encoding == "JSON": - return orjson.loads(data) - if encoding == "BASE64": - return orjson.loads(base64.b64decode(data)) + try: + if encoding == "JSON": + return orjson.loads(data) + if encoding == "BASE64": + return orjson.loads(base64.b64decode(data)) + except (binascii.Error, ValueError, orjson.JSONDecodeError) as e: + raise DatasetLoaderError( + f"Capture record {event_id} payload decode failed: {e}" + ) from e return None @@ -120,7 +130,9 @@ def _parse_trace(self, record: dict) -> SageMakerDataCaptureTrace: ) from e try: - input_data = _decode_payload(record["captureData"]["endpointInput"]) + input_data = _decode_payload( + record["captureData"]["endpointInput"], event_id + ) except KeyError as e: raise DatasetLoaderError( f"Capture record {event_id} missing captureData.endpointInput: {e}" @@ -132,7 +144,9 @@ def _parse_trace(self, record: dict) -> SageMakerDataCaptureTrace: "Only OpenAI-compatible chat endpoints are supported." ) - output_data = _decode_payload(record["captureData"].get("endpointOutput", {})) + output_data = _decode_payload( + record["captureData"].get("endpointOutput", {}), event_id + ) usage = output_data.get("usage", {}) if isinstance(output_data, dict) else {} max_tokens = input_data.get("max_tokens") @@ -292,7 +306,6 @@ def load_dataset( items = [t for t in raw_items if self._filter_and_cap_trace(t)] items.sort(key=lambda t: t.timestamp) - self._log_filtering_summary() self.info(f"Loaded {len(items):,} traces from {source_desc}") data = self._group_traces(items) @@ -300,4 +313,13 @@ def load_dataset( if _has_meaningful_synthesis(self._synthesis): data = self._apply_synthesis(data) + # Apply --synthesis-max-osl AFTER synthesis (matches the base + # load_dataset). This override otherwise returned uncapped output + # lengths -- the base-loader refactor moved the cap out of + # _filter_and_cap_trace into this grouped pass, which overriding + # loaders must call explicitly. Log AFTER the cap so the capped count + # is reflected in the summary (matches the base order). + data = self._cap_grouped_traces_max_osl(data) + self._log_filtering_summary() + return data diff --git a/src/aiperf/dataset/synthesis/synthesizer.py b/src/aiperf/dataset/synthesis/synthesizer.py index cb0a5d0d3a..be2abe76da 100644 --- a/src/aiperf/dataset/synthesis/synthesizer.py +++ b/src/aiperf/dataset/synthesis/synthesizer.py @@ -99,8 +99,12 @@ def synthesize_traces(self, traces: list[dict]) -> list[dict]: isl = self.params.max_isl # Only set input_length if the original trace used input_length - # (not text_input or messages) to avoid validation errors - if trace.get("text_input") is None and trace.get("messages") is None: + # (not text_input, messages, or payload) to avoid validation errors + if ( + trace.get("text_input") is None + and trace.get("messages") is None + and trace.get("payload") is None + ): synthetic_trace["input_length"] = isl # Apply timestamp scaling if present diff --git a/tests/unit/dataset/loader/test_burst_gpt_trace.py b/tests/unit/dataset/loader/test_burst_gpt_trace.py index 97add653c0..dbfd7bf2fb 100644 --- a/tests/unit/dataset/loader/test_burst_gpt_trace.py +++ b/tests/unit/dataset/loader/test_burst_gpt_trace.py @@ -225,3 +225,29 @@ def test_reconstruct_raises_on_length_mismatch(self, tmp_path: Path) -> None: ] with pytest.raises(ValueError, match="synth_dicts length"): loader._reconstruct_traces(originals, synth_dicts) + + +class TestMaxOslCapApplied: + """--synthesis-max-osl must cap output_length for BurstGPT too. The + base-loader refactor moved the cap out of _filter_and_cap_trace into a + grouped pass run by the base load_dataset; this loader overrides + load_dataset and must call _cap_grouped_traces_max_osl explicitly (it once + returned uncapped output lengths).""" + + def test_output_length_capped_after_load(self, tmp_path: Path) -> None: + path = _make_csv_file([_make_row(response_tokens="200")], tmp_path) + loader = _make_loader(path) + loader._max_osl = 50 + data = loader.load_dataset() + osls = [t.output_length for traces in data.values() for t in traces] + assert osls and all(o == 50 for o in osls) + assert loader._capped_max_osl == 1 + + def test_under_cap_unchanged(self, tmp_path: Path) -> None: + path = _make_csv_file([_make_row(response_tokens="30")], tmp_path) + loader = _make_loader(path) + loader._max_osl = 50 + data = loader.load_dataset() + osls = [t.output_length for traces in data.values() for t in traces] + assert osls == [30] + assert loader._capped_max_osl == 0 diff --git a/tests/unit/dataset/loader/test_sagemaker_data_capture.py b/tests/unit/dataset/loader/test_sagemaker_data_capture.py index 8093c9798c..1b330c75bc 100644 --- a/tests/unit/dataset/loader/test_sagemaker_data_capture.py +++ b/tests/unit/dataset/loader/test_sagemaker_data_capture.py @@ -773,6 +773,46 @@ def test_parse_trace_missing_capture_data_raises_dataset_loader_error( with pytest.raises(DatasetLoaderError, match="captureData.endpointInput"): loader._parse_trace(record) + def test_parse_trace_invalid_base64_input_raises_dataset_loader_error( + self, tmp_path: Path + ) -> None: + """Bad base64 in the input payload is attributed per-record, not raw binascii.Error.""" + loader = self._make_loader(tmp_path) + record = { + "captureData": { + "endpointInput": { + "data": "not-valid-base64!!!", + "encoding": "BASE64", + }, + }, + "eventMetadata": { + "eventId": "bad-b64", + "inferenceTime": "2026-04-29T00:00:00Z", + }, + } + with pytest.raises(DatasetLoaderError, match="bad-b64 payload decode failed"): + loader._parse_trace(record) + + def test_parse_trace_invalid_inner_json_raises_dataset_loader_error( + self, tmp_path: Path + ) -> None: + """Bad inner JSON in the input payload is attributed per-record, not raw JSONDecodeError.""" + loader = self._make_loader(tmp_path) + record = { + "captureData": { + "endpointInput": { + "data": "not valid json {{{", + "encoding": "JSON", + }, + }, + "eventMetadata": { + "eventId": "bad-json", + "inferenceTime": "2026-04-29T00:00:00Z", + }, + } + with pytest.raises(DatasetLoaderError, match="bad-json payload decode failed"): + loader._parse_trace(record) + def test_load_dataset_empty_file_returns_empty(self, tmp_path: Path) -> None: f = tmp_path / "empty.jsonl" f.write_text("") @@ -789,3 +829,37 @@ def test_load_dataset_empty_file_returns_empty(self, tmp_path: Path) -> None: # a directory) and is invoked by AIPerfConfig validation, not as a public # CLIConfig method. There's no v2 path that mirrors the original "count # JSONL entries across a SageMaker date-partitioned directory tree" call. + + +class TestMaxOslCapApplied: + """--synthesis-max-osl must cap output_length for SageMaker too. This loader + overrides load_dataset and must call _cap_grouped_traces_max_osl explicitly + (the base-loader refactor moved the cap into that grouped pass; the override + once returned uncapped output lengths).""" + + def _loader_for(self, tmp_path: Path, line: str) -> SageMakerDataCaptureLoader: + path = tmp_path / "capture.jsonl" + path.write_text(line + "\n") + return SageMakerDataCaptureLoader( + filename=str(path), + run=make_run_from_cli(CLIConfig(model_names=["test-model"])), + prompt_generator=MagicMock(), + ) + + def test_output_length_capped_after_load(self, tmp_path: Path) -> None: + line = _make_capture_record(max_tokens=200, completion_tokens=200) + loader = self._loader_for(tmp_path, line) + loader._max_osl = 50 + data = loader.load_dataset() + osls = [t.output_length for traces in data.values() for t in traces] + assert osls and all(o == 50 for o in osls) + assert loader._capped_max_osl == 1 + + def test_under_cap_unchanged(self, tmp_path: Path) -> None: + line = _make_capture_record(max_tokens=30, completion_tokens=30) + loader = self._loader_for(tmp_path, line) + loader._max_osl = 50 + data = loader.load_dataset() + osls = [t.output_length for traces in data.values() for t in traces] + assert osls == [30] + assert loader._capped_max_osl == 0 diff --git a/tests/unit/dataset/synthesis/test_synthesizer.py b/tests/unit/dataset/synthesis/test_synthesizer.py index b9eb118963..57d556d3b0 100644 --- a/tests/unit/dataset/synthesis/test_synthesizer.py +++ b/tests/unit/dataset/synthesis/test_synthesizer.py @@ -429,6 +429,26 @@ def test_text_input_not_modified(self) -> None: assert "input_length" not in synthetic[0] assert "input_length" not in synthetic[1] + def test_payload_trace_not_given_input_length(self) -> None: + """Traces with a pre-built payload must not get input_length added. + + MooncakeTrace requires exactly one of input_length / text_input / + messages / payload, so injecting input_length into a payload trace + makes the synthesized dataset fail validation on reconstruction. + """ + traces = [ + {"payload": {"prompt": "Hello", "max_tokens": 50}, "output_length": 50}, + {"payload": {"prompt": "World"}, "output_length": 100, "timestamp": 1000}, + ] + synthesizer = Synthesizer() + synthetic = synthesizer.synthesize_traces(traces) + + # payload should be preserved verbatim, input_length should NOT be added + assert synthetic[0]["payload"] == {"prompt": "Hello", "max_tokens": 50} + assert synthetic[1]["payload"] == {"prompt": "World"} + assert "input_length" not in synthetic[0] + assert "input_length" not in synthetic[1] + # ============================================================================ # Input Length Preservation Tests # ============================================================================