perf: encode json with orjson

src/prime_rl/entrypoints/dashboard.py

@@ -17,7 +17,6 @@
 import fcntl
 import functools
 import importlib.util
-import json
 import os
 import shutil
 import subprocess
@@ -28,6 +27,8 @@
 from pathlib import Path
 from typing import Callable
 
+import orjson
+
 from prime_rl.utils.links import PRIME_EVAL_LINKS, PRIME_LINKS, WANDB_LINKS
 from prime_rl.utils.logger import get_logger
 from prime_rl.utils.pathing import CACHE_DIR
@@ -53,22 +54,22 @@ def register_output_dir(output_dir: Path) -> None:
     """Add the dir to the daemon's registry (idempotent, atomic)."""
     with registry_lock():
         try:
-            dirs = json.loads(DIRS_FILE.read_text())
+            dirs = orjson.loads(DIRS_FILE.read_bytes())
         except (OSError, ValueError):
             dirs = []
         entry = str(output_dir.resolve())
         if entry in dirs:
             return
         dirs.append(entry)
         tmp = DIRS_FILE.with_suffix(".tmp")
-        tmp.write_text(json.dumps(dirs))
+        tmp.write_bytes(orjson.dumps(dirs))
         tmp.replace(DIRS_FILE)
 
 
 def find_daemon(timeout: float = 1.0) -> dict | None:
     """The live daemon's record ({pid, url, ...}), or None."""
     try:
-        info = json.loads(DAEMON_FILE.read_text())
+        info = orjson.loads(DAEMON_FILE.read_bytes())
         os.kill(info["pid"], 0)
     except (OSError, ValueError, KeyError, TypeError):
         return None

src/prime_rl/entrypoints/eval.py

@@ -5,7 +5,6 @@
 """
 
 import asyncio
-import json
 import os
 import re
 import signal
@@ -16,6 +15,8 @@
 from subprocess import Popen
 from typing import Any
 
+import orjson
+
 from prime_rl.configs.eval import EvalConfig
 from prime_rl.utils.config import cli, dump_resolved_config
 from prime_rl.utils.process import DEFAULT_COMMON_ENV_VARS, cleanup_processes, set_proc_title
@@ -43,7 +44,7 @@ def set_nested(target: dict[str, Any], keys: list[str], value: Any) -> None:
 
 def parse_value(raw: str) -> Any:
     """A JSON list/object stays structured; every other value is left to pydantic."""
-    return json.loads(raw) if raw.startswith(("[", "{")) else raw
+    return orjson.loads(raw) if raw.startswith(("[", "{")) else raw
 
 
 def expand_shorthands(argv: list[str]) -> list[str]:
@@ -87,7 +88,7 @@ def expand_shorthands(argv: list[str]) -> list[str]:
                 "The <taskset-id> shorthand describes a single source and cannot be combined "
                 "with a config file that defines [[source]] blocks - use one or the other"
             )
-        out += ["--source", json.dumps([source])]
+        out += ["--source", orjson.dumps([source]).decode()]
     return out
 
 
@@ -142,7 +143,7 @@ def main():
     logger.info("Starting eval")
 
     write_launch_artifacts(config_dir, "eval")
-    (config_dir / "eval.json").write_text(json.dumps(dump_resolved_config(config), indent=2))
+    (config_dir / "eval.json").write_bytes(orjson.dumps(dump_resolved_config(config), option=orjson.OPT_INDENT_2))
     components: list[tuple[str, Path | str]] = [("Eval", config_dir / "eval.json")]
     # One env server per source without an explicit `serve.address`, like the rl launcher's:
     # `env-server @ <path>` binds an OS-assigned port and publishes it to the source's

src/prime_rl/entrypoints/inference.py

@@ -1,4 +1,3 @@
-import json
 import os
 import signal
 import subprocess
@@ -7,6 +6,8 @@
 from threading import Event, Thread
 from typing import Any
 
+import orjson
+
 from prime_rl.configs.inference import InferenceConfig, MooncakeKVCacheOffloadConfig
 from prime_rl.utils.config import cli, dump_resolved_config
 from prime_rl.utils.logger import setup_logger
@@ -36,7 +37,7 @@ def vllm_overrides_fragment(overrides: dict[str, Any]) -> str:
     """
     if not overrides:
         return ""
-    return ", " + json.dumps(overrides)[1:-1].replace('"', '\\"')
+    return ", " + orjson.dumps(overrides).decode()[1:-1].replace('"', '\\"')
 
 
 def write_config(
@@ -52,8 +53,7 @@ def write_config(
     config_dict = dump_resolved_config(config, exclude=exclude)
     if engine_only:
         config_dict["router"] = None
-    with open(config_path, "w") as f:
-        json.dump(config_dict, f, indent=2)
+    config_path.write_bytes(orjson.dumps(config_dict, option=orjson.OPT_INDENT_2))
     return config_path
 
 
@@ -66,7 +66,7 @@ def write_pd_configs(config: InferenceConfig, output_dir: Path) -> None:
             config.for_pd_role(role), exclude={"deployment", "slurm", "output_dir", "dry_run"}
         )
         role_config["router"] = None
-        (output_dir / f"inference-{role}.json").write_text(json.dumps(role_config, indent=2))
+        (output_dir / f"inference-{role}.json").write_bytes(orjson.dumps(role_config, option=orjson.OPT_INDENT_2))
 
 
 def mooncake_offloads(config: InferenceConfig | None) -> dict[str, MooncakeKVCacheOffloadConfig]:

src/prime_rl/entrypoints/rl.py

@@ -1,10 +1,11 @@
-import json
 import os
 import subprocess
 import sys
 from pathlib import Path
 from urllib.parse import urlparse
 
+import orjson
+
 from prime_rl.configs.algorithm import FrozenModelConfig
 from prime_rl.configs.inference import VllmRouterConfig
 from prime_rl.configs.orchestrator import EnvConfig
@@ -65,8 +66,9 @@ def env_server_names(config: RLConfig, split: str) -> list[str]:
 def write_config(config: RLConfig, output_dir: Path, exclude: set[str] | None = None) -> None:
     """Write resolved config to disk, excluding launcher-only fields."""
     output_dir.mkdir(parents=True, exist_ok=True)
-    with open(output_dir / RL_CONFIG, "w") as f:
-        json.dump(dump_resolved_config(config, exclude=exclude), f, indent=2)
+    (output_dir / RL_CONFIG).write_bytes(
+        orjson.dumps(dump_resolved_config(config, exclude=exclude), option=orjson.OPT_INDENT_2)
+    )
 
 
 def rl_config_components(config: RLConfig, config_dir: Path) -> list[tuple[str, Path | str]]:
@@ -87,8 +89,7 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None:
     output_dir.mkdir(parents=True, exist_ok=True)
     filenames = {"trainer": TRAINER_CONFIG, "orchestrator": ORCHESTRATOR_CONFIG, "inference": INFERENCE_CONFIG}
     for name, component in dump_rl_components(config).items():
-        with open(output_dir / filenames[name], "w") as f:
-            json.dump(component, f, indent=2)
+        (output_dir / filenames[name]).write_bytes(orjson.dumps(component, option=orjson.OPT_INDENT_2))
     if config.inference is not None:
         write_pd_configs(config.inference, output_dir)
 
@@ -216,7 +217,7 @@ def rl_local(config: RLConfig):
                 **DEFAULT_COMMON_ENV_VARS,
                 "LOGURU_FORCE_COLORS": "1",
                 "WANDB_PROGRAM": "uv run rl",
-                "WANDB_ARGS": json.dumps(start_command),
+                "WANDB_ARGS": orjson.dumps(start_command).decode(),
                 **config.env_vars,
                 **config.orchestrator.env_vars,
                 **wandb_shared_env,
@@ -242,7 +243,7 @@ def rl_local(config: RLConfig):
                 **DEFAULT_TRAINER_ENV_VARS,
                 "LOGURU_FORCE_COLORS": "1",
                 "WANDB_PROGRAM": "uv run rl",
-                "WANDB_ARGS": json.dumps(start_command),
+                "WANDB_ARGS": orjson.dumps(start_command).decode(),
                 **config.env_vars,
                 **config.trainer.env_vars,
                 **wandb_shared_env,

src/prime_rl/entrypoints/sft.py

@@ -1,9 +1,10 @@
-import json
 import os
 import subprocess
 import sys
 from pathlib import Path
 
+import orjson
+
 from prime_rl.configs.eval import SFTOnlineEvalConfig
 from prime_rl.configs.monitors import EvalMonitorsConfig, PrimeEvalMonitorConfig, TrainMonitorsConfig
 from prime_rl.configs.orchestrator import OnlineEvalSourceConfig
@@ -121,8 +122,7 @@ def sft_config_components(config: SFTConfig, config_dir: Path) -> list[tuple[str
 def write_config(config: SFTConfig, config_path: Path, exclude: set[str] | None = None) -> None:
     """Write resolved config to disk, excluding launcher-only fields."""
     config_path.parent.mkdir(parents=True, exist_ok=True)
-    with open(config_path, "w") as f:
-        json.dump(dump_resolved_config(config, exclude=exclude), f, indent=2)
+    config_path.write_bytes(orjson.dumps(dump_resolved_config(config, exclude=exclude), option=orjson.OPT_INDENT_2))
 
 
 def write_eval_subconfigs(config: SFTConfig, config_dir: Path, strip_router: bool = False) -> None:
@@ -136,11 +136,10 @@ def write_eval_subconfigs(config: SFTConfig, config_dir: Path, strip_router: boo
         if strip_router:
             # Per-rank processes run bare engines; the sbatch starts the single global router.
             inference_dict["router"] = None
-        with open(config_dir / INFERENCE_CONFIG, "w") as f:
-            json.dump(inference_dict, f, indent=2)
+        (config_dir / INFERENCE_CONFIG).write_bytes(orjson.dumps(inference_dict, option=orjson.OPT_INDENT_2))
 
-    with open(config_dir / ONLINE_EVAL_CONFIG, "w") as f:
-        json.dump(dump_resolved_config(build_online_eval_config(config)), f, indent=2)
+    online_eval = dump_resolved_config(build_online_eval_config(config))
+    (config_dir / ONLINE_EVAL_CONFIG).write_bytes(orjson.dumps(online_eval, option=orjson.OPT_INDENT_2))
 
     # One EnvServerConfig per launcher-managed eval source: `env-server @ <path>` binds an
     # OS-assigned port and publishes it to the source's address file, where the
@@ -351,7 +350,7 @@ def sft_local(config: SFTConfig):
             "WANDB_SHARED_PRIMARY": "trainer",
             "WANDB_SHARED_FINISHER": "online-eval",
             "WANDB_PROGRAM": "uv run sft",
-            "WANDB_ARGS": json.dumps(sys.argv),
+            "WANDB_ARGS": orjson.dumps(sys.argv).decode(),
         }
 
     with ProcessGroup() as processes:

src/prime_rl/eval/online.py

@@ -13,9 +13,10 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import os
 
+import orjson
+
 from prime_rl import monitors
 from prime_rl.configs.eval import SFTOnlineEvalConfig
 from prime_rl.configs.shared import FileSystemWeightBroadcastConfig
@@ -215,7 +216,7 @@ def main():
     config_dir, log_dir = prepare_attempt_dirs(config.output_dir)
     os.environ["PRL_ATTEMPT_CONFIG_DIR"] = str(config_dir)
     os.environ["PRL_ATTEMPT_LOG_DIR"] = str(log_dir)
-    (config_dir / "eval.json").write_text(json.dumps(dump_resolved_config(config), indent=2))
+    (config_dir / "eval.json").write_bytes(orjson.dumps(dump_resolved_config(config), option=orjson.OPT_INDENT_2))
     setup_logger(config.log.level, json_logging=config.log.json_logging)
     asyncio.run(run_online_eval(config))
 

src/prime_rl/inference/json_logging.py

@@ -10,13 +10,14 @@
 """
 
 import datetime
-import json
 import logging
 import tempfile
 import traceback
 from pathlib import Path
 from typing import Any
 
+import orjson
+
 
 class JsonFormatter(logging.Formatter):
     """Emit one JSON object per record matching `build_log_entry` in
@@ -34,7 +35,7 @@ def format(self, record: logging.LogRecord) -> str:
         }
         if record.exc_info:
             entry["exception"] = "".join(traceback.format_exception(*record.exc_info))
-        return json.dumps(entry, default=str)
+        return orjson.dumps(entry, default=str).decode()
 
 
 # Loguru-only level names that stdlib `logging` doesn't recognize.
@@ -83,6 +84,6 @@ def write_logging_config(level: str = "info") -> Path:
     config = build_dict_config(level)
     fd, path_str = tempfile.mkstemp(prefix="prime_rl_logging_", suffix=".json")
     path = Path(path_str)
-    with open(fd, "w", encoding="utf-8") as f:
-        json.dump(config, f)
+    with open(fd, "wb") as f:
+        f.write(orjson.dumps(config))
     return path

src/prime_rl/monitors/file/monitor.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import shutil
 import time
 from pathlib import Path
@@ -73,7 +72,7 @@ async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None:
         row = {"step": step, "time": time.time(), **sanitized}
         if self.producer is not None:
             row["producer"] = self.producer
-        self.file.write(json.dumps(row) + "\n")
+        self.file.write(orjson.dumps(row).decode() + "\n")
 
     async def log_eval_plan(self, env_name: str, step: int, expected: int) -> None:
         """Merge the epoch's expected count into ``plan.json`` (atomic replace)."""

src/prime_rl/monitors/file/traces/live.py

@@ -16,7 +16,6 @@
 
 import argparse
 import hashlib
-import json
 import sys
 import threading
 import time
@@ -273,7 +272,7 @@ def main() -> None:
         if folded is None:
             sys.exit(f"no live trace {args.trace_id} under {get_live_dir(args.run_dir)}")
         dispatch, trace = folded
-        print(json.dumps({"dispatch": dispatch, "trace": trace}, indent=2, default=str))
+        print(orjson.dumps({"dispatch": dispatch, "trace": trace}, default=str, option=orjson.OPT_INDENT_2).decode())
         return
     rows = live_rows(args.run_dir)
     if not rows:

src/prime_rl/monitors/wandb/monitor.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import os
 import sys
 import time
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
+import orjson
 import wandb
 from wandb.errors import CommError, UsageError
 from wandb.sdk.mailbox.mailbox_handle import ServerResponseError
@@ -41,7 +41,7 @@ async def init(
         wandb_args = os.environ.get("WANDB_ARGS")
         if wandb_args:
             self.logger.debug(f"Found WANDB_ARGS in environment variables {wandb_args}")
-            sys.argv = json.loads(wandb_args)
+            sys.argv = orjson.loads(wandb_args)
 
         # W&B advertises Weave on every init when it sees openai/verifiers imported.
         os.environ.setdefault("WANDB_DISABLE_WEAVE", "1")

src/prime_rl/trainer/lora.py

@@ -201,9 +201,10 @@ def save_lora_config(model: nn.Module, save_path, rank: int, alpha: float, dropo
         alpha: LoRA alpha scaling parameter
         dropout: LoRA dropout rate
     """
-    import json
     from pathlib import Path
 
+    import orjson
+
     save_path = Path(save_path)
 
     target_modules = {name.split(".")[-1] for name, module in model.named_modules() if isinstance(module, LoRAModule)}
@@ -220,5 +221,4 @@ def save_lora_config(model: nn.Module, save_path, rank: int, alpha: float, dropo
     }
 
     config_path = save_path / "adapter_config.json"
-    with open(config_path, "w") as f:
-        json.dump(adapter_config, f, indent=2)
+    config_path.write_bytes(orjson.dumps(adapter_config, option=orjson.OPT_INDENT_2))

src/prime_rl/trainer/sft/data.py

@@ -1,11 +1,11 @@
-import json
 import time
 import uuid
 from collections import defaultdict
 from pathlib import Path
 from typing import Any, Callable, Literal, TypedDict, cast
 
 import numpy as np
+import orjson
 import torch
 from datasets import Dataset, interleave_datasets, load_dataset
 from huggingface_hub import snapshot_download
@@ -334,7 +334,7 @@ def resolve_messages(example: dict) -> list[dict]:
         # JSON-encoded string of one.
         tools = example.get(self.columns.tools) or []
         if isinstance(tools, str):
-            tools = json.loads(tools)
+            tools = orjson.loads(tools)
 
         def should_mask(message: dict) -> bool:
             assert "role" in message, "Message must have a role"

src/prime_rl/utils/chat_template.py

@@ -1,6 +1,7 @@
-import json
 from typing import Any
 
+import orjson
+
 
 def normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]:
     if messages is None:
@@ -27,7 +28,7 @@ def _deserialize_tool_call(tool_call: dict[str, Any]) -> dict[str, Any]:
         function = tool_call.get("function", {})
         arguments = function.get("arguments")
         if isinstance(arguments, str):
-            arguments = json.loads(arguments)
+            arguments = orjson.loads(arguments)
         return {
             **tool_call,
             "function": {**function, "arguments": arguments},