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

[ajcasagrande]

The bug: --synthesis-max-osl silently doesn't cap for loaders that override load_dataset

The base-loader refactor moved the max-OSL cap out of _filter_and_cap_trace into a grouped pass (_cap_grouped_traces_max_osl) run by BaseTraceDatasetLoader.load_dataset (src/aiperf/dataset/loader/base_trace_loader.py:256), applied after synthesis. Two loaders override load_dataset and never call it, so --synthesis-max-osl returned uncapped output lengths with no warning:

• BurstGPT (src/aiperf/dataset/loader/burst_gpt.py): CSV override now calls _cap_grouped_traces_max_osl(data) after synthesis and logs the filtering summary after the cap, so the capped count shows up in the summary (matching the base order).
• SageMaker Data Capture (src/aiperf/dataset/loader/sagemaker_data_capture.py): same fix, same ordering. Also hardens _decode_payload to re-raise malformed base64 / inner-JSON payloads as DatasetLoaderError with per-record eventId attribution instead of leaking raw binascii.Error / JSONDecodeError that aborts the whole load without saying which record was bad.

Synthesizer payload guard

Synthesizer.synthesize_traces (src/aiperf/dataset/synthesis/synthesizer.py) injected input_length into any trace lacking text_input/messages — including traces that use a pre-built payload. MooncakeTrace requires exactly one of input_length / text_input / messages / payload, so synthesis over payload traces produced datasets that fail validation. The guard now also skips traces with payload.

Tests (all pinned via revert-verification)

Each src change was reverted with git apply -R, its tests confirmed FAILING, then re-applied and confirmed green:

• Revert sagemaker fix → 3 failures (TestMaxOslCapApplied::test_output_length_capped_after_load, test_parse_trace_invalid_base64_input_raises_dataset_loader_error, test_parse_trace_invalid_inner_json_raises_dataset_loader_error)
• Revert burst_gpt fix → 1 failure (TestMaxOslCapApplied::test_output_length_capped_after_load)
• Revert synthesizer fix → 1 failure (TestSynthesizer::test_payload_trace_not_given_input_length, new test written for this PR — the fix previously shipped test-less)

Evidence

• uv run pytest -n auto tests/unit/dataset/loader/test_sagemaker_data_capture.py tests/unit/dataset/loader/test_burst_gpt_trace.py tests/unit/dataset/synthesis/test_synthesizer.py → 120 passed
• pytest --collect-only clean (120 collected)
• ruff format --check / ruff check clean on all touched files
• make check-ruff-baselined → OK (118 total, 118 baselined, 0 new)

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Improved trace loading so output-length limits are applied at the right stage, with summaries now reflecting the final capped values.
  • Made SageMaker capture import more resilient by providing clearer record-level errors when payload decoding or parsing fails.
  • Preserved existing payload-based traces during synthesis without incorrectly adding input-length values.

• Tests

  • Added coverage for output-length capping and for clearer decode-error handling in dataset loaders.
  • Added a test to confirm payload traces remain unchanged during synthesis.

[github-actions]

Try out this PR

Quick install:

pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@760ee619314c72131f540c8040fc60c2b5bc5555

Recommended with virtual environment (using uv):

uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@760ee619314c72131f540c8040fc60c2b5bc5555

---

Last updated for commit: 760ee61 • Browse code

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: db83bfa3-d5db-428b-bdcc-631f948bbaa2

📥 Commits

Reviewing files that changed from the base of the PR and between 0314074 and 760ee61.

📒 Files selected for processing (6)

• src/aiperf/dataset/loader/burst_gpt.py
• src/aiperf/dataset/loader/sagemaker_data_capture.py
• src/aiperf/dataset/synthesis/synthesizer.py
• tests/unit/dataset/loader/test_burst_gpt_trace.py
• tests/unit/dataset/loader/test_sagemaker_data_capture.py
• tests/unit/dataset/synthesis/test_synthesizer.py

---

Walkthrough

This PR reorders max output-length capping relative to filtering summary logging in the BurstGPT and SageMaker dataset loaders, adds event-attributed error handling for SageMaker payload decoding, and gates Synthesizer's input_length assignment on payload absence. Unit tests cover these behaviors.

Changes

Dataset loader capping, decoding, and synthesis behavior

Layer / File(s)	Summary
BurstGPT loader capping and logging order src/aiperf/dataset/loader/burst_gpt.py, tests/unit/dataset/loader/test_burst_gpt_trace.py	Removes an early filtering summary log call and instead applies grouped max-OSL capping after optional synthesis, followed by the filtering summary log; new tests verify capped and under-cap output lengths.
SageMaker payload decode error handling and capping order src/aiperf/dataset/loader/sagemaker_data_capture.py, tests/unit/dataset/loader/test_sagemaker_data_capture.py	_decode_payload now accepts an event_id, handles JSON/BASE64 explicitly, and wraps decode failures into DatasetLoaderError; _parse_trace passes event_id at both call sites; load_dataset moves the trace-count log earlier and applies capping before the filtering summary log; new tests cover decode error paths and cap behavior.
Synthesizer input_length gating on payload presence src/aiperf/dataset/synthesis/synthesizer.py, tests/unit/dataset/synthesis/test_synthesizer.py	input_length is now only set when text_input, messages, and payload are all None; a new test confirms payload-based traces retain their payload without an injected input_length.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Poem

A rabbit hopped through logs and caps,
Capping OSL before mishaps,
Base64 crumbs and JSON bits,
Wrapped in errors that finally fit,
Payloads spared their input_length spree —
Hop hop hurray, the tests agree! 🐰✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main fix: applying the synthesis max-OSL cap in loader overrides.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[datadog-official]

⚠️ Warnings

🚦 1 Pipeline job failed

Run Integration Tests | integration-tests (ubuntu-latest, 3.12)     

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 760ee61 | Docs | Give us feedback!}