feat: use msgspec for faster JSON encoding and decoding #493

ajcasagrande · 2025-11-25T15:05:50Z

Important

Converting to draft as still running benchmarks to compare, and the jury is still out

Replaces all usage of json and orjson with msgspec.json for better performance. See: https://github.com/jcrist/msgspec

Summary by CodeRabbit

Chores
- Updated JSON serialization dependency from orjson to msgspec across the application and test suite.
- Migrated all JSON encoding and decoding operations to use the new library.
- Updated error handling to work with msgspec's exception types.
- No changes to functionality or end-user APIs.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

github-actions · 2025-11-25T15:06:00Z

Try out this PR

Quick install:

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

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@deac901f5795c530d0e56e11f8218f1487b0713e

Last updated for commit: deac901 • Browse code

codecov · 2025-11-25T15:08:45Z

Codecov Report

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/aiperf/common/models/record_models.py	33.33%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

coderabbitai · 2025-11-25T15:14:59Z

Walkthrough

The codebase migrates from orjson to msgspec for JSON serialization and deserialization. Changes include updating the dependency in pyproject.toml, replacing all orjson.loads()/dumps() with msgspec.json.decode()/encode(), updating exception handling from orjson.JSONDecodeError to msgspec.DecodeError, and adjusting all affected tests and test assertions accordingly.

Changes

Cohort / File(s)	Summary
Dependency Configuration `pyproject.toml`	Removed `orjson~=3.10.18` dependency; added `msgspec~=0.20.0` dependency.
Core JSON Serialization `src/aiperf/common/messages/base_messages.py`, `src/aiperf/common/mixins/buffered_jsonl_writer_mixin.py`, `src/aiperf/transports/aiohttp_transport.py`	Updated JSON encoding to use `msgspec.json.encode()` replacing `orjson.dumps()` for serializing model data to bytes.
Core JSON Deserialization & Error Handling `src/aiperf/common/config/config_validators.py`, `src/aiperf/common/config/user_config.py`, `src/aiperf/common/models/record_models.py`, `src/aiperf/common/models/sequence_distribution.py`, `src/aiperf/common/utils.py`, `src/aiperf/endpoints/template_endpoint.py`	Updated JSON decoding to use `msgspec.json.decode()` replacing `orjson.loads()`, and updated exception handling to catch `msgspec.DecodeError` instead of `orjson.JSONDecodeError`.
Test Infrastructure & Config `tests/aiperf_mock_server/config.py`, `tests/unit/dataset/loader/conftest.py`	Updated JSON serialization to use `msgspec.json.encode()` in config serialization; changed file write mode from text ("w") to binary ("wb") in test fixtures.
Integration Tests `tests/integration/test_gpu_telemetry.py`, `tests/integration/test_random_generator_canary.py`, `tests/integration/utils.py`	Replaced JSON operations with `msgspec.json.decode()` and `msgspec.json.encode()` for test data serialization/deserialization and file I/O operations.
Unit Tests — Messages & Models `tests/unit/common/messages/test_messages.py`, `tests/unit/common/mixins/test_buffered_jsonl_writer_mixin.py`, `tests/unit/common/models/test_auto_routed_messages.py`, `tests/unit/common/models/test_sequence_distribution.py`	Replaced JSON decoding with `msgspec.json.decode()` in test assertions; renamed test methods and updated docstrings to reference `msgspec` instead of `orjson`.
Unit Tests — Dataset Loaders `tests/unit/dataset/loader/test_multi_turn.py`, `tests/unit/dataset/loader/test_single_turn.py`, `tests/unit/dataset/test_dataset_manager_inputs_json.py`	Replaced `json.dumps()` and `json.loads()` with `msgspec.json.encode()` and `msgspec.json.decode()` for test data generation and parsing.
Unit Tests — Exporters `tests/unit/exporters/test_metrics_json_exporter.py`, `tests/unit/exporters/test_timeslice_metrics_json_exporter.py`	Updated JSON parsing in test assertions from `json.loads()/json.load()` to `msgspec.json.decode()`.
Unit Tests — Post Processors `tests/unit/post_processors/test_raw_record_writer_processor.py`, `tests/unit/post_processors/test_record_export_results_processor.py`, `tests/unit/post_processors/test_telemetry_export_results_processor.py`	Replaced `orjson.loads()` with `msgspec.json.decode()` for parsing JSONL output lines in test validations.
Unit Tests — Transports & Server `tests/unit/transports/test_aiohttp_client.py`, `tests/unit/transports/test_aiohttp_transport.py`, `tests/unit/server/test_config.py`	Updated JSON serialization to use `msgspec.json.encode().decode("utf-8")` in request/response tests; renamed test method to reflect `msgspec`; adjusted test expectation for serialized list formatting.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Exception handling behavior: Verify msgspec.DecodeError is functionally equivalent to orjson.JSONDecodeError in all error paths, particularly around error message content and propagation.
Binary vs. text mode file operations: Check that conftest.py fixture changes (from "w" to "wb") correctly handle mixed string/bytes inputs and don't break dependent tests.
Test expectation changes: Review updated test assertions in test_config.py regarding list serialization formatting (spacing changes between elements).
Serialization output compatibility: Confirm that msgspec.json.encode() produces identical or compatible JSON output to orjson.dumps() for downstream consumers.

Poem

🐰 Hops with glee through code so clean,
Orjson out, msgspec in the scene!
Swift and speedy, bytes aligned,
One JSON lib, no more entwined—
Tests all pass, the refactor's done,
A rabbit's hop to faster fun! 🎉

Pre-merge checks

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: replacing orjson with msgspec for JSON operations across the codebase for performance improvements.
Docstring Coverage	✅ Passed	Docstring coverage is 94.62% which is sufficient. The required threshold is 80.00%.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

coderabbitai

Actionable comments posted: 2

🧹 Nitpick comments (11)

CLAUDE.md (1)

81-81: Update retrieved learnings to reflect msgspec migration.

The coding guideline has been updated to use msgspec, but the retrieved learnings still reference orjson. Based on learnings, the previous standard was to use orjson. Since this PR intentionally migrates to msgspec, the learnings should be updated to reflect this new standard.
tests/unit/dataset/loader/conftest.py (1)
27-36: Consider clarifying the fallback case.

The file mode change to binary ("wb") and the type handling logic are correct for msgspec compatibility. However, the else branch on line 35 is unclear—what type would reach this case that isn't str or bytes? Consider either removing this fallback or adding a comment explaining when it would be triggered. If no other types are expected, raising a TypeError would be more explicit.

Apply this diff to make the handling more explicit:
-                if isinstance(line, str):
-                    f.write(line.encode("utf-8") + b"\n")
-                elif isinstance(line, bytes):
-                    f.write(line + b"\n")
-                else:
-                    f.write(line + b"\n")
+                if isinstance(line, str):
+                    f.write(line.encode("utf-8") + b"\n")
+                elif isinstance(line, bytes):
+                    f.write(line + b"\n")
+                else:
+                    raise TypeError(f"Expected str or bytes, got {type(line)}")
tests/unit/common/models/test_auto_routed_messages.py (1)

5-5: Routing tests correctly exercise msgspec-encoded inputs; consider aligning docstring with actual input type

Using msgspec.json.encode for the from_json inputs (plus the derived str/bytes/bytearray variants) is consistent with the new JSON backend and keeps these tests close to real usage. The only mismatch is that test_json_string_routing now passes bytes while its docstring advertises a JSON string; either decoding to UTF‑8 before calling from_json, or updating the docstring to mention “JSON bytes” would remove that confusion.

Please double‑check that Message.from_json and load_json_str are intended to support all four input types (dict, str, bytes, bytearray) with msgspec as the backend, and update the docstring accordingly if bytes are now the primary path.

Also applies to: 143-152, 178-185

src/aiperf/common/config/config_validators.py (1)

8-8: Exception handling for msgspec-backed JSON parsing is correct; reconcile with global JSON-backend guideline

Catching msgspec.DecodeError around load_json_str(input) and re‑raising a ValueError keeps the function’s external behavior intact while aligning with the new msgspec JSON backend. The only concern is process-level: repo guidelines currently say to “Always use orjson for JSON operations”, but this module (and others in the PR) now depend on msgspec instead. It would be good to either (a) confirm that msgspec is the new standard and update the guidelines/CL docs, or (b) explicitly document this module as an intentional exception to avoid future confusion.

Please verify the current contract of load_json_str (i.e., that it raises msgspec.DecodeError for malformed JSON) and update the coding guidelines or an ADR to reflect msgspec as the canonical JSON backend if that’s now the intended direction. As per coding guidelines, this change currently diverges from the “Always use orjson for JSON operations” rule.

Also applies to: 150-157

src/aiperf/common/models/sequence_distribution.py (1)

38-38: JSON distribution parsing correctly wraps msgspec.DecodeError; ensure backend contract stays stable

Catching msgspec.DecodeError in _parse_pairs_from_json and re‑raising a ValueError("Invalid JSON format: ...") is a sensible way to surface malformed distribution specs while delegating actual parsing to load_json_str. This matches the broader move to msgspec, but also hard‑codes knowledge of the JSON backend into this model module, similar to config_validators.

Please verify that load_json_str guarantees raising msgspec.DecodeError (and not some other exception type) for malformed JSON so this handler remains effective, and consider documenting msgspec as the canonical JSON backend in the project guidelines to avoid future confusion about the “Always use orjson” rule that’s still present in the shared coding guidelines.

Also applies to: 361-367

src/aiperf/transports/aiohttp_transport.py (1)

10-10: Payload serialization via msgspec is correct; consider bytes path later

Switching to msgspec.json.encode(payload).decode("utf-8") cleanly preserves the existing str payload contract to AioHttpClient.post_request while using the new backend.

If you later relax AioHttpClient to accept bytes, you could skip the .decode("utf-8") here and pass the encoded bytes directly for a tiny efficiency win, but that’s purely optional.

Also applies to: 151-153

tests/unit/transports/test_aiohttp_client.py (1)

9-9: End‑to‑end transport tests correctly switched to msgspec serialization

Using msgspec.json.encode(...).decode("utf-8") for both mocked responses and request bodies brings the tests in line with the transport’s new serialization behavior while preserving the str API expected by AioHttpClient.post_request.

If you find yourself adding more such tests, consider a small local helper like encode_json(obj) -> str to DRY up the encode/decode pattern, but it’s not required here.

Also applies to: 326-334, 355-358

src/aiperf/common/messages/base_messages.py (1)

6-7: Message.to_json_bytes correctly migrated to msgspec; update global JSON guideline

Using msgspec.json.encode(self.model_dump(exclude_none=True, mode="json")) here is a clean swap from the previous orjson implementation: callers still get compact JSON bytes with None fields stripped, and the docstring accurately reflects the new backend and performance intent.

Given this central message type now standardizes on msgspec, it would be good to update CLAUDE/coding guidelines that currently say “Always use orjson for JSON operations” so they instead point to msgspec as the preferred JSON backend going forward. Based on learnings, that will keep future changes consistent with this PR’s direction.

Also applies to: 46-61

tests/unit/dataset/loader/test_single_turn.py (1)

4-4: Verify create_jsonl_file handles msgspec-encoded bytes

msgspec.json.encode(...) returns bytes, whereas most other tests here pass JSON str literals into create_jsonl_file. This is fine as long as the fixture writes bytes (or normalizes both bytes/str internally); otherwise you may hit type errors when constructing the JSONL file.

If create_jsonl_file still assumes str, either update it to accept both bytes and str, or decode here with .decode("utf-8") before adding to content.

Also applies to: 257-273

tests/unit/dataset/loader/test_multi_turn.py (1)

4-4: Ensure dataset JSONL helper supports msgspec-encoded bytes (and mixed bytes/str)

All these tests now use msgspec.json.encode(...), which returns bytes, and in some cases (test_load_dataset_skips_empty_lines) the content list mixes bytes and str (the empty-line sentinel).

This is only safe if create_jsonl_file has been updated to:

accept both bytes and str, and

write them in a compatible way (e.g., binary mode for bytes, or decoding bytes to UTF‑8 before writing as text).

If the fixture still assumes str and does operations like line + "\n" in text mode, these tests will fail. Consider either:

normalizing each encoded payload with .decode("utf-8") before adding it to content, or

explicitly updating create_jsonl_file to handle bytes | str robustly.

Also applies to: 126-136, 158-176, 195-203, 223-239, 261-271, 286-301, 322-343, 373-386, 403-414

src/aiperf/endpoints/template_endpoint.py (1)

11-12: msgspec-based decoding looks good; consider asserting object type and reconciling JSON-lib guidelines

Switching from orjson.loads(rendered) to msgspec.json.decode(rendered) and catching msgspec.DecodeError is a valid replacement here, and the error propagation remains clear.

Two follow-ups you might consider:

After decoding, defensively assert the template produced a JSON object before calling payload.update(...), e.g. if not isinstance(payload, dict): raise ValueError("Template must render a JSON object, got ..."), so non-object templates fail fast with a clearer message.

The current coding guidelines still state “Always use orjson for JSON operations”; if msgspec is now the project-wide standard (as this PR suggests), those guidelines should be updated to avoid contradicting this implementation.

As per coding guidelines, the JSON backend choice should be documented consistently.

Also applies to: 145-153

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 563d9d8 and deac901.

📒 Files selected for processing (31)

CLAUDE.md (1 hunks)
pyproject.toml (1 hunks)
src/aiperf/common/config/config_validators.py (2 hunks)
src/aiperf/common/config/user_config.py (2 hunks)
src/aiperf/common/messages/base_messages.py (3 hunks)
src/aiperf/common/mixins/buffered_jsonl_writer_mixin.py (2 hunks)
src/aiperf/common/models/record_models.py (3 hunks)
src/aiperf/common/models/sequence_distribution.py (2 hunks)
src/aiperf/common/utils.py (2 hunks)
src/aiperf/endpoints/template_endpoint.py (2 hunks)
src/aiperf/transports/aiohttp_transport.py (2 hunks)
tests/aiperf_mock_server/config.py (2 hunks)
tests/integration/test_gpu_telemetry.py (2 hunks)
tests/integration/test_random_generator_canary.py (2 hunks)
tests/integration/utils.py (3 hunks)
tests/unit/common/messages/test_messages.py (11 hunks)
tests/unit/common/mixins/test_buffered_jsonl_writer_mixin.py (2 hunks)
tests/unit/common/models/test_auto_routed_messages.py (3 hunks)
tests/unit/common/models/test_sequence_distribution.py (1 hunks)
tests/unit/dataset/loader/conftest.py (1 hunks)
tests/unit/dataset/loader/test_multi_turn.py (10 hunks)
tests/unit/dataset/loader/test_single_turn.py (2 hunks)
tests/unit/dataset/test_dataset_manager_inputs_json.py (7 hunks)
tests/unit/exporters/test_metrics_json_exporter.py (12 hunks)
tests/unit/exporters/test_timeslice_metrics_json_exporter.py (8 hunks)
tests/unit/post_processors/test_raw_record_writer_processor.py (5 hunks)
tests/unit/post_processors/test_record_export_results_processor.py (5 hunks)
tests/unit/post_processors/test_telemetry_export_results_processor.py (14 hunks)
tests/unit/server/test_config.py (1 hunks)
tests/unit/transports/test_aiohttp_client.py (3 hunks)
tests/unit/transports/test_aiohttp_transport.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (6)

**/*.py