perf: encode json with orjson in prime-runs and prime-traces

packages/prime-runs/pyproject.toml

@@ -19,6 +19,7 @@ authors = [
 # `tenacity`. Nothing is added to what verifiers and prime-rl already carry.
 dependencies = [
     "httpx>=0.25.0",
+    "orjson>=3.10.0",
     "prime-traces>=0.0.3",
 ]
 keywords = ["evals", "evaluations", "training", "runs", "traces", "observability"]

packages/prime-runs/src/prime_runs/_http.py

@@ -15,13 +15,13 @@
 Error mapping and backoff come from ``prime_traces.core.client``.
 """
 
-import json
 import sys
 import time
 from typing import Any, Dict, Mapping, Optional, Tuple, Union
 from urllib.parse import urlsplit, urlunsplit
 
 import httpx
+import orjson
 from prime_traces.core.client import AMBIGUOUS_TRANSPORT_ERRORS, raise_for_response, retry_delay
 
 from . import _fork
@@ -72,10 +72,8 @@ def _user_agent() -> str:
 
 
 def encode_json(value: Any) -> bytes:
-    """Compact UTF-8 JSON; bare ``NaN`` is refused here rather than as a 400."""
-    return json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode(
-        "utf-8"
-    )
+    """Compact UTF-8 JSON; NaN and infinity go out as ``null``, never as a bare ``NaN``."""
+    return orjson.dumps(value, option=orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS)
 
 
 class PlatformClient:

packages/prime-runs/src/prime_runs/projection.py

@@ -3,11 +3,12 @@
 training viewer's Parquet table (moved here from prime-rl). Duck-typed: verifiers
 ``Trace``/``Episode`` satisfy it structurally and are not imported."""
 
-import json
 import logging
 import math
 from typing import Any, Dict, Iterable, List, Optional, Sequence
 
+import orjson
+
 from ._http import encode_json
 
 logger = logging.getLogger(__name__)
@@ -38,7 +39,7 @@ def _nullify_non_finite(value: Any) -> Any:
 
 
 def _json_column(value: Any) -> str:
-    return json.dumps(_nullify_non_finite(value), allow_nan=False)
+    return orjson.dumps(value, option=orjson.OPT_NON_STR_KEYS).decode()
 
 
 def summary_trace_index(episode: Any) -> int:

packages/prime-runs/tests/conftest.py

@@ -58,9 +58,9 @@ def paths(self) -> List[str]:
         return [f"{r.method} {r.url.path}" for r in self.requests]
 
     def bodies_for(self, path: str) -> List[Any]:
-        import json
+        import orjson
 
-        return [json.loads(r.content) for r in self.requests if r.url.path == path and r.content]
+        return [orjson.loads(r.content) for r in self.requests if r.url.path == path and r.content]
 
 
 @pytest.fixture
@@ -149,9 +149,9 @@ def rft_routes() -> Dict[str, Any]:
     """The happy path for a training run: register, metrics, samples, finalize."""
 
     def presign(request: httpx.Request) -> httpx.Response:
-        import json
+        import orjson
 
-        step = json.loads(request.content)["step"]
+        step = orjson.loads(request.content)["step"]
         return httpx.Response(
             200,
             json={

packages/prime-runs/tests/test_config_source.py

@@ -1,7 +1,6 @@
 """The config a run is actually configured with: the file someone wrote, kept"""
 
-import json
-
+import orjson
 import pytest
 from conftest import RecordingHandler
 
@@ -139,7 +138,7 @@ def test_an_online_run_sends_the_source_in_create_metadata(tmp_path, online):
     run.finish()
 
     create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/")
-    metadata = json.loads(create.content)["metadata"]
+    metadata = orjson.loads(create.content)["metadata"]
     assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML
     assert metadata[CONFIG_SOURCE_KEY]["format"] == "toml"
     assert metadata[CONFIG_SOURCE_KEY]["filename"] == "eval.toml"
@@ -157,7 +156,7 @@ def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path, online):
     run.finish()
 
     create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/")
-    metadata = json.loads(create.content)["metadata"]
+    metadata = orjson.loads(create.content)["metadata"]
     assert metadata["model"] == "deepseek/deepseek-v4-flash"
     assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML
     assert run.config_source.filename == "eval.toml"

packages/prime-runs/tests/test_init.py

@@ -315,14 +315,16 @@ def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run(
     assert handler.paths().count("POST /api/v1/evaluations/eval-abc/finalize") == 1
 
 
-def test_records_reject_nonfinite_json_instead_of_sending_an_opaque_400(online):
-    run, _ = online(on_error="raise")
-    run.log_traces([{"sample_id": "bad", "reward": float("nan")}])
-
-    with pytest.raises(ValueError, match="Out of range float values"):
-        run.flush()
+def test_records_send_nonfinite_floats_as_null_instead_of_an_opaque_400(online):
+    run, handler = online(on_error="raise")
+    run.log_traces([{"sample_id": "nan", "reward": float("nan")}])
+    run.flush()
     run.finish()
 
+    posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")
+    (sample,) = [s for body in posted for s in body["samples"]]
+    assert sample["sample_id"] == "nan" and sample["reward"] is None
+
 
 # ----------------------------------------------------------------- training
 

packages/prime-runs/tests/test_metrics_sink.py

@@ -50,10 +50,10 @@ def flaky(request: httpx.Request) -> httpx.Response:
 def test_a_persistent_transient_failure_loses_only_its_step(
     make_platform_client, rft_routes, no_sleep
 ):
-    import json
+    import orjson
 
     def down_for_step_one(request: httpx.Request) -> httpx.Response:
-        if json.loads(request.content)["metrics"]["step"] == 1:
+        if orjson.loads(request.content)["metrics"]["step"] == 1:
             return httpx.Response(502)
         return httpx.Response(200, json={"data": {"status": "success"}})
 

packages/prime-runs/tests/test_train_samples_sink.py

@@ -111,10 +111,10 @@ def test_nothing_to_encode_makes_no_request(make_platform_client, rft_routes):
 
 
 def test_a_failed_presign_loses_that_step_only(make_platform_client, rft_routes):
-    import json
+    import orjson
 
     def flaky(request: httpx.Request) -> httpx.Response:
-        if json.loads(request.content)["step"] == 10:
+        if orjson.loads(request.content)["step"] == 10:
             return httpx.Response(502)
         return rft_routes["POST /api/v1/rft/samples/presign"](request)
 
@@ -319,7 +319,8 @@ def test_the_default_encoder_writes_the_viewer_table():
     pa = pytest.importorskip("pyarrow")
     pq = pytest.importorskip("pyarrow.parquet")
     import io
-    import json
+
+    import orjson
 
     from prime_runs.projection import episodes_to_parquet_bytes, train_sample_schema
 
@@ -341,20 +342,20 @@ def test_the_default_encoder_writes_the_viewer_table():
     assert [row["reward"] for row in rows] == [1.0, 0.0]
     assert [row["advantage"] for row in rows] == [0.25, -0.25]
     assert [row["env_name"] for row in rows] == ["gsm8k", "gsm8k"]
-    trajectory = json.loads(rows[0]["trajectory"])
+    trajectory = orjson.loads(rows[0]["trajectory"])
     assert trajectory[-1]["advantage"] == 0.25
     assert rows[0]["num_output_tokens"] == 5
-    info = json.loads(rows[0]["info"])
+    info = orjson.loads(rows[0]["info"])
     assert info["native_wrapper"]["id"] == "e1"
     assert isinstance(pa.Table, type)
 
 
 def test_the_default_encoder_writes_non_finite_numbers_as_null():
     pytest.importorskip("pyarrow")
     import io
-    import json
     import math
 
+    import orjson
     import pyarrow.parquet as pq
 
     from prime_runs.projection import episodes_to_parquet_bytes
@@ -365,9 +366,9 @@ def test_the_default_encoder_writes_non_finite_numbers_as_null():
     payload = episodes_to_parquet_bytes([episode], "run-abc", 10)
 
     (row,) = pq.read_table(io.BytesIO(payload)).to_pylist()
-    strict = dict(parse_constant=lambda constant: pytest.fail(f"bare {constant} in a column"))
-    assert json.loads(row["metrics"], **strict) == {"kl": None}
-    assert json.loads(row["trajectory"], **strict)[-1]["advantage"] is None
+    # orjson.loads rejects a bare NaN or Infinity
+    assert orjson.loads(row["metrics"]) == {"kl": None}
+    assert orjson.loads(row["trajectory"])[-1]["advantage"] is None
     assert math.isinf(row["advantage"])  # the float column keeps the value
 
 

packages/prime-runs/tests/test_upload_routing.py

@@ -1,8 +1,7 @@
 """Exercise exclusive routing through init, the worker, and real Traces HTTP calls."""
 
-import json
-
 import httpx
+import orjson
 import pytest
 from _fakes import make_episode, make_trace, make_train_episode
 from conftest import FakeSink, RecordingHandler
@@ -123,7 +122,7 @@ def _opt_out_via(source, monkeypatch, home):
         monkeypatch.setenv("PRIME_TRACES_OPT_OUT", "true")
     else:
         (home / ".prime").mkdir()
-        (home / ".prime" / "config.json").write_text(json.dumps({"traces_opt_out": True}))
+        (home / ".prime" / "config.json").write_bytes(orjson.dumps({"traces_opt_out": True}))
     return {}
 
 

packages/prime-traces/pyproject.toml

@@ -11,6 +11,7 @@ authors = [
 ]
 dependencies = [
     "httpx>=0.25.0",
+    "orjson>=3.10.0",
     "pydantic>=2.0.0",
 ]
 keywords = ["traces", "observability", "evals", "rollouts"]

packages/prime-traces/src/prime_traces/async_traces.py

@@ -86,7 +86,7 @@ def _record_lines_for(
 
     ``aiter_batches`` consumes a synchronous source entirely in a worker
     thread, so handing it the sync generator — rather than adapting it to an
-    async one here — is what keeps ``json.dumps`` over large records off the
+    async one here — is what keeps encoding large records off the
     event loop.
     """
     if isinstance(records, AsyncIterable):

packages/prime-traces/src/prime_traces/core/client.py

@@ -16,7 +16,6 @@
 """
 
 import gzip as gzip_module
-import json as json_module
 import random
 import sys
 import time
@@ -25,6 +24,7 @@
 from typing import Any, Dict, Iterator, Optional
 
 import httpx
+import orjson
 
 from ..exceptions import (
     AmbiguousDeleteError,
@@ -317,7 +317,7 @@ def _multipart_files(
             traces_part = ("traces.jsonl", body, "application/x-ndjson")
 
         return {
-            "metadata": (None, json_module.dumps(metadata).encode(), "application/json"),
+            "metadata": (None, orjson.dumps(metadata), "application/json"),
             "traces": traces_part,
         }
 

packages/prime-traces/src/prime_traces/core/config.py

@@ -4,11 +4,12 @@
 environment variables, env taking precedence.
 """
 
-import json
 import os
 from pathlib import Path
 from typing import Optional
 
+import orjson
+
 
 class Config:
     """Minimal configuration class for SDK packages.
@@ -29,8 +30,8 @@ def _load_config(self) -> None:
         config_data: object = {}
         if self.config_file.exists():
             try:
-                config_data = json.loads(self.config_file.read_text())
-            except (json.JSONDecodeError, OSError, UnicodeDecodeError):
+                config_data = orjson.loads(self.config_file.read_bytes())
+            except (orjson.JSONDecodeError, OSError):
                 config_data = {}
         # Valid JSON that is not an object (a list, a bare string) must degrade
         # the same way invalid JSON does: every accessor assumes a dict, and

packages/prime-traces/src/prime_traces/traces.py

@@ -18,7 +18,6 @@
 field registry).
 """
 
-import json
 import tempfile
 import time
 from pathlib import Path
@@ -37,6 +36,8 @@
 )
 from urllib.parse import quote
 
+import orjson
+
 from .batching import (
     DEFAULT_TARGET_BATCH_BYTES,
     Batch,
@@ -77,6 +78,8 @@ def to_record(self) -> Mapping[str, Any]: ...
 
 TraceRecord = Union[Mapping[str, Any], SupportsToRecord]
 
+RECORD_OPTIONS = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS | orjson.OPT_APPEND_NEWLINE
+
 
 def _encode_record(record: TraceRecord, record_number: int) -> bytes:
     """Serialize one in-memory record as a compact UTF-8 JSONL line.
@@ -95,18 +98,9 @@ def _encode_record(record: TraceRecord, record_number: int) -> bytes:
         if not isinstance(value, Mapping):
             raise TypeError(f"Record {record_number} to_record() must return a mapping")
 
-    # Compact separators match common JSONL writers, ensure_ascii=False
-    # keeps the wire representation UTF-8, and allow_nan=False prevents
-    # emitting JavaScript-only NaN/Infinity values that strict JSON rejects.
-    return (
-        json.dumps(
-            dict(value),
-            ensure_ascii=False,
-            allow_nan=False,
-            separators=(",", ":"),
-        ).encode("utf-8")
-        + b"\n"
-    )
+    # orjson writes compact UTF-8 and encodes NaN/Infinity as null, which
+    # strict JSON parsers accept. Numpy arrays let callers skip a list copy.
+    return orjson.dumps(dict(value), option=RECORD_OPTIONS)
 
 
 def _record_lines(records: Iterable[TraceRecord]) -> Iterator[bytes]:

packages/prime-traces/tests/test_async_upload.py

@@ -10,10 +10,10 @@
 import asyncio
 import gzip
 import hashlib
-import json
 import threading
 
 import httpx
+import orjson
 import pytest
 from test_upload_wire import parse_multipart
 
@@ -52,7 +52,7 @@ def handler(request: httpx.Request) -> httpx.Response:
 
         parts = parse_multipart(captured["content"], request.headers["content-type"])
         assert parts["traces"][1] == RAW
-        assert json.loads(parts["metadata"][1]) == {
+        assert orjson.loads(parts["metadata"][1]) == {
             "schema_version": 1,
             "context": {"source": "hosted_eval"},
         }