feat: use msgspec for faster JSON encoding and decoding

CLAUDE.md

@@ -78,7 +78,7 @@ except Exception as e:
 
 **Logging**: Lambda for expensive logs: `self.debug(lambda: f"{len(self._x())}")`. Direct string: `self.info("Starting")`.
 
-**JSON**: Always orjson: `orjson.loads(s)`, `orjson.dumps(d)`
+**JSON**: Always msgspec: `msgspec.json.decode(s)`, `msgspec.json.encode(d)`
 
 ## Mixins & Base Classes
 

pyproject.toml

@@ -30,8 +30,8 @@ dependencies = [
   "ffmpeg-python~=0.2.0",
   "jinja2~=3.1.5",  # NOTE: Versions prior to 3.1.5 have vuln exploits
   "jmespath~=1.0.1",
+  "msgspec~=0.20.0",
   "numpy~=1.26.4",
-  "orjson~=3.10.18",
   "pillow~=11.1.0",
   "prometheus_client~=0.23.1",
   "psutil~=7.0.0",

src/aiperf/common/config/config_validators.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 from typing import Any
 
-import orjson
+import msgspec
 
 from aiperf.common.enums.service_enums import ServiceType
 from aiperf.common.utils import load_json_str
@@ -151,7 +151,7 @@ def parse_str_or_dict_as_tuple_list(input: Any | None) -> list[tuple[str, Any]]
         if input.startswith("{"):
             try:
                 return [(key, value) for key, value in load_json_str(input).items()]
-            except orjson.JSONDecodeError as e:
+            except msgspec.DecodeError as e:
                 raise ValueError(
                     f"User Config: {input} - must be a valid JSON string"
                 ) from e

src/aiperf/common/config/user_config.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 from typing import Annotated, Any
 
-from orjson import JSONDecodeError
+import msgspec
 from pydantic import BeforeValidator, Field, model_validator
 from typing_extensions import Self
 
@@ -141,7 +141,7 @@ def _should_use_fixed_schedule_for_mooncake_trace(self) -> bool:
                     try:
                         data = load_json_str(line)
                         return "timestamp" in data and data["timestamp"] is not None
-                    except (JSONDecodeError, KeyError):
+                    except (msgspec.DecodeError, KeyError):
                         continue
         except (OSError, FileNotFoundError):
             _logger.warning(

src/aiperf/common/messages/base_messages.py

@@ -3,7 +3,7 @@
 import time
 from typing import ClassVar
 
-import orjson
+import msgspec
 from pydantic import Field
 
 from aiperf.common.enums.message_enums import MessageType
@@ -44,10 +44,10 @@ def __str__(self) -> str:
         return self.model_dump_json(exclude_none=True)
 
     def to_json_bytes(self) -> bytes:
-        """Serialize message to JSON bytes using orjson for optimal performance.
+        """Serialize message to JSON bytes using msgspec for optimal performance.
 
-        This method uses orjson for high-performance serialization (6x faster for
-        large records >20KB). It automatically excludes None fields to minimize
+        This method uses msgspec for high-performance serialization (ultra-fast
+        serialization for large records). It automatically excludes None fields to minimize
         message size.
 
         Returns:
@@ -57,7 +57,7 @@ def to_json_bytes(self) -> bytes:
             Prefer this method over model_dump_json() for ZMQ message passing
             and other high-throughput scenarios.
         """
-        return orjson.dumps(self.model_dump(exclude_none=True, mode="json"))
+        return msgspec.json.encode(self.model_dump(exclude_none=True, mode="json"))
 
 
 class RequiresRequestNSMixin(Message):

src/aiperf/common/mixins/buffered_jsonl_writer_mixin.py

@@ -7,7 +7,7 @@
 from typing import Generic
 
 import aiofiles
-import orjson
+import msgspec
 
 from aiperf.common.environment import Environment
 from aiperf.common.hooks import on_init, on_stop
@@ -57,28 +57,30 @@ def __init__(
     async def _open_file(self) -> None:
         """Open the file handle for writing in binary mode (called automatically on initialization)."""
         async with self._file_lock:
-            # Binary mode for optimal performance with orjson
+            # Binary mode for optimal performance with msgspec
             self._file_handle = await aiofiles.open(self.output_file, mode="wb")
 
     async def buffered_write(self, record: BaseModelT) -> None:
         """Write a Pydantic model to the buffer with automatic flushing.
 
-        This method serializes the provided Pydantic model to JSON bytes using orjson
+        This method serializes the provided Pydantic model to JSON bytes using msgspec
         and adds it to the internal buffer. If the buffer reaches the configured batch
         size, it automatically flushes the buffer to disk.
 
-        Uses binary mode with orjson for optimal performance:
-        - 6x faster for large records (>20KB)
+        Uses binary mode with msgspec for optimal performance:
+        - Ultra-fast serialization for large records
         - No encode/decode overhead
         - Efficient for all record sizes
 
         Args:
             record: A Pydantic BaseModel instance to write
         """
         try:
-            # Serialize to bytes using orjson (faster for large records)
+            # Serialize to bytes using msgspec (faster for large records)
             # Use exclude_none=True to omit None fields (smaller output)
-            json_bytes = orjson.dumps(record.model_dump(exclude_none=True, mode="json"))
+            json_bytes = msgspec.json.encode(
+                record.model_dump(exclude_none=True, mode="json")
+            )
 
             buffer_to_flush = None
             async with self._buffer_lock:

src/aiperf/common/models/record_models.py

@@ -7,7 +7,7 @@
 from functools import cached_property
 from typing import Any, AnyStr
 
-import orjson
+import msgspec
 from pydantic import (
     Field,
     RootModel,
@@ -250,7 +250,7 @@ def get_json(self) -> JsonObject | None:
             if not self.text:
                 return None
             return load_json_str(self.text)
-        except orjson.JSONDecodeError:
+        except msgspec.DecodeError:
             return None
 
 
@@ -347,7 +347,7 @@ def get_json(self) -> JsonObject | None:
             if data_content in ("", None, "[DONE]"):
                 return None
             return load_json_str(data_content)
-        except orjson.JSONDecodeError:
+        except msgspec.DecodeError:
             return None
 
 

src/aiperf/common/models/sequence_distribution.py

@@ -35,8 +35,8 @@
 import re
 from dataclasses import dataclass
 
+import msgspec
 import numpy as np
-import orjson
 
 from aiperf.common import random_generator as rng
 from aiperf.common.aiperf_logger import AIPerfLogger
@@ -362,7 +362,7 @@ def _parse_pairs_from_json(cls, json_str: str) -> list[SequenceLengthPair]:
         """Parse JSON format and extract pairs: {"pairs": [{"isl": 256, "isl_stddev": 10, "osl": 128, "osl_stddev": 5, "prob": 40}, ...]}"""
         try:
             data = load_json_str(json_str)
-        except orjson.JSONDecodeError as e:
+        except msgspec.DecodeError as e:
             raise ValueError(f"Invalid JSON format: {e}") from e
 
         if "pairs" not in data:

src/aiperf/common/utils.py

@@ -7,7 +7,7 @@
 from collections.abc import Callable
 from typing import Any
 
-import orjson
+import msgspec
 
 from aiperf.common import aiperf_logger
 from aiperf.common.aiperf_logger import AIPerfLogger
@@ -90,11 +90,9 @@ def load_json_str(
           run validation checks on the object. Defaults to identity function.
     """
     try:
-        # Note: orjson may not parse JSON the same way as Python's standard json library,
-        # notably being stricter on UTF-8 conformance.
-        # Refer to https://github.com/ijl/orjson?tab=readme-ov-file#str for details.
-        return func(orjson.loads(json_str))
-    except orjson.JSONDecodeError as e:
+        # Note: msgspec is strict on UTF-8 conformance and provides fast JSON parsing
+        return func(msgspec.json.decode(json_str))
+    except msgspec.DecodeError as e:
         snippet = json_str[:200] + ("..." if len(json_str) > 200 else "")
         _logger.exception(f"Failed to parse JSON string: '{snippet}' - {e!r}")
         raise

src/aiperf/endpoints/template_endpoint.py

@@ -8,7 +8,7 @@
 
 import jinja2
 import jmespath
-import orjson
+import msgspec
 
 from aiperf.common.decorators import implements_protocol
 from aiperf.common.enums import EndpointType
@@ -145,8 +145,8 @@ def format_payload(self, request_info: RequestInfo) -> dict[str, Any]:
         rendered = self._template.render(**template_vars)
 
         try:
-            payload = orjson.loads(rendered)
-        except orjson.JSONDecodeError as e:
+            payload = msgspec.json.decode(rendered)
+        except msgspec.DecodeError as e:
             self.error(f"Template did not render valid JSON: {rendered} - {e!r}")
             raise ValueError(
                 f"Template did not render valid JSON {e!r}: {rendered[:100]}"

src/aiperf/transports/aiohttp_transport.py

@@ -7,7 +7,7 @@
 from collections.abc import Mapping
 from typing import Any
 
-import orjson
+import msgspec
 
 from aiperf.common.enums import TransportType
 from aiperf.common.exceptions import NotInitializedError
@@ -148,8 +148,8 @@ async def send_request(
             url = self.build_url(request_info)
             headers = self.build_headers(request_info)
 
-            # Serialize with orjson for performance
-            json_str = orjson.dumps(payload).decode("utf-8")
+            # Serialize with msgspec for performance
+            json_str = msgspec.json.encode(payload).decode("utf-8")
             record = await self.aiohttp_client.post_request(url, json_str, headers)
             record.request_headers = headers
         except Exception as e:

tests/aiperf_mock_server/config.py

@@ -2,11 +2,11 @@
 # SPDX-License-Identifier: Apache-2.0
 """Mock server configuration."""
 
-import json
 import logging
 import os
 from typing import Annotated, Any, Literal
 
+import msgspec
 from cyclopts import Parameter
 from pydantic import Field, model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -153,5 +153,5 @@ def _get_env_key(config_key: str) -> str:
 def _serialize_env_value(value: Any) -> str:
     """Serialize value for environment variable storage."""
     if isinstance(value, list | dict):
-        return json.dumps(value)
+        return msgspec.json.encode(value).decode("utf-8")
     return str(value)

tests/integration/test_gpu_telemetry.py

@@ -4,7 +4,7 @@
 
 import platform
 
-import orjson
+import msgspec
 import pytest
 
 from aiperf.common.models.telemetry_models import TelemetryRecord
@@ -98,7 +98,7 @@ async def test_gpu_telemetry_export(
 
         # Validate each line is valid JSON and can be parsed as TelemetryRecord
         for line in lines:
-            record_dict = orjson.loads(line)
+            record_dict = msgspec.json.decode(line)
             record = TelemetryRecord.model_validate(record_dict)
 
             # Verify required fields are present

tests/integration/test_random_generator_canary.py

@@ -8,9 +8,9 @@
 that any changes to the codebase don't silently break determinism.
 """
 
-import json
 from pathlib import Path
 
+import msgspec
 import pytest
 
 from aiperf.common.utils import load_json_str
@@ -115,5 +115,5 @@ def _save_reference(self, data: dict) -> None:
             data: Inputs data to save as reference
         """
         self.REFERENCE_FILE.parent.mkdir(parents=True, exist_ok=True)
-        with open(self.REFERENCE_FILE, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2)
+        with open(self.REFERENCE_FILE, "wb") as f:
+            f.write(msgspec.json.format(msgspec.json.encode(data), indent=2))

tests/integration/utils.py

@@ -3,11 +3,10 @@
 """Utility functions for integration tests."""
 
 import base64
-import json
 import subprocess
 from pathlib import Path
 
-import orjson
+import msgspec
 
 from aiperf.common.aiperf_logger import AIPerfLogger
 from tests.integration.models import VideoDetails
@@ -34,7 +33,7 @@ def create_rankings_dataset(tmp_path: Path, num_entries: int) -> Path:
                     {"name": "passages", "contents": [f"AI passage {i}"]},
                 ]
             }
-            f.write(orjson.dumps(entry).decode("utf-8") + "\n")
+            f.write(msgspec.json.encode(entry).decode("utf-8") + "\n")
     return dataset_path
 
 
@@ -62,7 +61,7 @@ def extract_base64_video_details(base64_data: str) -> VideoDetails:
     ]
     result = subprocess.run(cmd, input=video_bytes, capture_output=True, check=True)
 
-    probe_data = json.loads(result.stdout)
+    probe_data = msgspec.json.decode(result.stdout)
     format_info = probe_data["format"]
     video_stream = next(s for s in probe_data["streams"] if s["codec_type"] == "video")