fix(dataset): apply --synthesis-max-osl cap in loader overrides

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
 
     # ------------------------------------------------------------------

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 = "<unknown>"
+) -> 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,12 +306,20 @@ 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)
 
         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

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

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

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

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
     # ============================================================================