refactor: replace pysimdjson with orjson

pysimdjson does not build from source on Python 3.14, macOS arm64.
orjson seems better maintained.
orjson is about 20% slower than pysimdjson.

Author: riddell-stan

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
 python = "^3.12"
 aiohttp = "^3.6"
 httpstan = "~4.17"
-pysimdjson = ">=5.0.2"
+orjson = ">=3.10"
 numpy = ">=1.19"
 clikit = "^0.6"
 setuptools = "*"

stan/common.py

@@ -6,7 +6,7 @@
 import aiohttp
 import aiohttp.web
 import httpstan.app
-import simdjson
+import orjson
 
 
 def unused_tcp_port():
@@ -22,8 +22,7 @@ class HTTPResponse(typing.NamedTuple):
     content: bytes
 
     def json(self) -> dict:
-        # mypy 0.961 complains that simdjson lacks a `loads`.
-        return simdjson.loads(self.content)  # type: ignore
+        return orjson.loads(self.content)
 
 
 class HttpstanClient:

stan/fit.py

@@ -4,7 +4,7 @@
 from typing import Generator, Tuple, cast
 
 import numpy as np
-import simdjson
+import orjson
 
 
 class Fit(collections.abc.Mapping):
@@ -65,43 +65,38 @@ def __init__(
         # _draws is an ndarray with shape (num_sample_and_sampler_params + num_flat_params, num_draws, num_chains)
         self._draws: np.ndarray
 
-        parser = simdjson.Parser()
         for chain_index, stan_output in zip(range(self.num_chains), self.stan_outputs):
             draw_index = 0
             for line in stan_output.splitlines():
                 try:
-                    msg = cast(simdjson.Object, parser.parse(line))
-                except ValueError:
-                    # Occurs when draws contain an nan or infinity. simdjson cannot parse such values.
+                    msg = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    # Occurs when draws contain a NaN or infinity. orjson cannot parse such values.
                     msg = json.loads(line)
-                try:
-                    if msg["topic"] == "sample":
-                        # Ignore sample message which is mixed together with proper draws.
-                        if not isinstance(msg["values"], (simdjson.Object, dict)):
-                            continue
-
-                        # for the first draw: collect sample and sampler parameter names.
-                        if not hasattr(self, "_draws"):
-                            feature_names = cast(Tuple[str, ...], tuple(msg["values"].keys()))
-                            self.sample_and_sampler_param_names = tuple(
-                                name for name in feature_names if name.endswith("__")
+                if msg["topic"] == "sample":
+                    # Ignore sample message which is mixed together with proper draws.
+                    if not isinstance(msg["values"], dict):
+                        continue
+
+                    # for the first draw: collect sample and sampler parameter names.
+                    if not hasattr(self, "_draws"):
+                        feature_names = cast(Tuple[str, ...], tuple(msg["values"].keys()))
+                        self.sample_and_sampler_param_names = tuple(
+                            name for name in feature_names if name.endswith("__")
+                        )
+                        num_rows = len(self.sample_and_sampler_param_names) + num_flat_params
+                        # column-major order ("F") aligns with how the draws are stored (in cols).
+                        self._draws = np.empty((num_rows, num_samples_saved, num_chains), order="F")
+                        # rudimentary check of parameter order (sample & sampler params must be first)
+                        if num_flat_params and feature_names[-1].endswith("__"):
+                            raise RuntimeError(
+                                f"Expected last parameter name to be one declared in program code, found `{feature_names[-1]}`"
                             )
-                            num_rows = len(self.sample_and_sampler_param_names) + num_flat_params
-                            # column-major order ("F") aligns with how the draws are stored (in cols).
-                            self._draws = np.empty((num_rows, num_samples_saved, num_chains), order="F")
-                            # rudimentary check of parameter order (sample & sampler params must be first)
-                            if num_flat_params and feature_names[-1].endswith("__"):
-                                raise RuntimeError(
-                                    f"Expected last parameter name to be one declared in program code, found `{feature_names[-1]}`"
-                                )
-
-                        draw_row = tuple(msg["values"].values())  # a "row" of values from a single draw from Stan C++
-                        draw_row = cast(Tuple[float, ...], draw_row)
-                        self._draws[:, draw_index, chain_index] = draw_row
-                        draw_index += 1
-                finally:
-                    # clean up `Object`s produced by parser, required by simdjson
-                    del msg
+
+                    draw_row = tuple(msg["values"].values())  # a "row" of values from a single draw from Stan C++
+                    draw_row = cast(Tuple[float, ...], draw_row)
+                    self._draws[:, draw_index, chain_index] = draw_row
+                    draw_index += 1
             assert draw_index == num_samples_saved
         assert self.sample_and_sampler_param_names and self._draws.size
         self._draws.flags["WRITEABLE"] = False  # type: ignore

stan/model.py

@@ -10,7 +10,7 @@
 import httpstan.services.arguments as arguments
 import httpstan.utils
 import numpy as np
-import simdjson
+import orjson
 from clikit.io import ConsoleIO
 
 import stan.common
@@ -257,10 +257,10 @@ async def go():
 
             stan_outputs = tuple(stan_outputs)  # Fit constructor expects a tuple.
 
-            def is_nonempty_logger_message(msg: simdjson.Object):
+            def is_nonempty_logger_message(msg: dict):
                 return msg["topic"] == "logger" and msg["values"][0] != "info:"  # type: ignore
 
-            def is_iteration_or_elapsed_time_logger_message(msg: simdjson.Object):
+            def is_iteration_or_elapsed_time_logger_message(msg: dict):
                 # Assumes `msg` is a message with topic `logger`.
                 text = msg["values"][0]  # type: ignore
                 text = cast(str, text)
@@ -271,19 +271,16 @@ def is_iteration_or_elapsed_time_logger_message(msg: simdjson.Object):
                     or text.startswith("info:" + " " * 15)
                 )
 
-            parser = simdjson.Parser()
             nonstandard_logger_messages = []
             for stan_output in stan_outputs:
                 for line in stan_output.splitlines():
                     # Do not attempt to parse non-logger messages. Draws could contain nan or inf values.
-                    # simdjson cannot parse lines containing such values.
+                    # orjson cannot parse lines containing such values.
                     if b'"logger"' not in line:
                         continue
-                    msg = parser.parse(line)
+                    msg = orjson.loads(line)
                     if is_nonempty_logger_message(msg) and not is_iteration_or_elapsed_time_logger_message(msg):
-                        nonstandard_logger_messages.append(msg.as_dict())
-                    del msg
-            del parser  # simdjson.Parser is no longer used at this point.
+                        nonstandard_logger_messages.append(msg)
 
             if nonstandard_logger_messages:
                 io.error_line("<comment>Messages received during sampling:</comment>")