Faster message serialization (#1064)

* Add orjson_packer and orjson_unpacker to speed up serialization of session messages.

* Refactor orjson packer to use functools.partial.

* Add msgpack support for message serialization in Session

* Refactor packer/unpacker change handling for improved readability

* Fix test_serialize_objects datetime checks on ci to compare using datetime format.

* Fix datetime deserialization in test_serialize_objects to use dateutil.parser.isoparse

* Add PicklingError to exception handling in test_cannot_serialize

* Update api docs

* Replace dateutil.parser.isoparse with jsonutil.parse_date in test_serialize_objects for datetime validation

* Use rep in fstring

Co-authored-by: M Bussonnier <[email protected]>

* Add msgpack as a test dependency.

* Fallback to json_packer and json_unpacker for orjson to handle for better resilience.

* Fix: test_args checking for removed _default_pack_unpack and _default_pack_unpack.

* Add type annotation to orjson_packer.

* Add the other missing type annoatation.

* Change orjson from a dependency to an optional dependency. Test against the minimum version in CI.

* Double timeout for test_minimum_verisons

* Fix invalid argument name.

* Add orjson and msgpack as additional_dependencies for mypy in .pre-commit-config.yaml.

* Remove # type:ignore[import-not-found].

* Should return result of orjson.loads.

* fix: get mypy working with orjson/msgpack

Signed-off-by: Henry Schreiner <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix for previous refactor.

---------

Signed-off-by: Henry Schreiner <[email protected]>
Co-authored-by: M Bussonnier <[email protected]>
Co-authored-by: Alan <>
Co-authored-by: Henry Schreiner <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

.github/workflows/main.yml

@@ -126,7 +126,7 @@ jobs:
   test_minimum_verisons:
     name: Test Minimum Versions
     runs-on: ubuntu-latest
-    timeout-minutes: 10
+    timeout-minutes: 20
     steps:
       - uses: actions/checkout@v4
       - uses: jupyterlab/maintainer-tools/.github/actions/base-setup@v1
@@ -143,6 +143,11 @@ jobs:
         run: |
           hatch -vv run test:nowarn
 
+      - name: Run the unit tests with orjson installed
+        run: |
+          hatch -e test run pip install orjson
+          hatch -vv run test:nowarn
+
   test_prereleases:
     name: Test Prereleases
     timeout-minutes: 10

.pre-commit-config.yaml

@@ -37,14 +37,18 @@ repos:
         types_or: [yaml, html, json]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.18.2"
+    rev: "v1.19.0"
     hooks:
       - id: mypy
         files: jupyter_client
         stages: [manual]
         args: ["--install-types", "--non-interactive"]
         additional_dependencies:
-          ["traitlets>=5.13", "ipykernel>=6.26", "jupyter_core>=5.3.2"]
+          - traitlets>=5.13
+          - ipykernel>=6.26
+          - jupyter_core>=5.3.2
+          - orjson>=3.11.4
+          - msgpack-types
 
   - repo: https://github.com/adamchainz/blacken-docs
     rev: "1.20.0"

jupyter_client/session.py

@@ -13,6 +13,7 @@
 # Distributed under the terms of the Modified BSD License.
 from __future__ import annotations
 
+import functools
 import hashlib
 import hmac
 import json
@@ -33,6 +34,7 @@
 from traitlets import (
     Any,
     Bool,
+    Callable,
     CBytes,
     CUnicode,
     Dict,
@@ -125,15 +127,48 @@ def json_unpacker(s: str | bytes) -> t.Any:
     return json.loads(s)
 
 
+try:
+    import orjson
+except ModuleNotFoundError:
+    has_orjson = False
+    orjson_packer, orjson_unpacker = json_packer, json_unpacker
+else:
+    has_orjson = True
+
+    def orjson_packer(
+        obj: t.Any, *, option: int | None = orjson.OPT_NAIVE_UTC | orjson.OPT_UTC_Z
+    ) -> bytes:
+        """Convert a json object to a bytes using orjson with fallback to json_packer."""
+        try:
+            return orjson.dumps(obj, default=json_default, option=option)
+        except Exception:
+            return json_packer(obj)
+
+    def orjson_unpacker(s: str | bytes) -> t.Any:
+        """Convert a json bytes or string to an object using orjson with fallback to json_unpacker."""
+        try:
+            return orjson.loads(s)
+        except Exception:
+            return json_unpacker(s)
+
+
+try:
+    import msgpack
+except ModuleNotFoundError:
+    has_msgpack = False
+else:
+    has_msgpack = True
+    msgpack_packer = functools.partial(msgpack.packb, default=json_default)
+    msgpack_unpacker = msgpack.unpackb
+
+
 def pickle_packer(o: t.Any) -> bytes:
     """Pack an object using the pickle module."""
     return pickle.dumps(squash_dates(o), PICKLE_PROTOCOL)
 
 
 pickle_unpacker = pickle.loads
 
-default_packer = json_packer
-default_unpacker = json_unpacker
 
 DELIM = b"<IDS|MSG>"
 # singleton dummy tracker, which will always report as done
@@ -316,7 +351,7 @@ class Session(Configurable):
 
     debug : bool
         whether to trigger extra debugging statements
-    packer/unpacker : str : 'json', 'pickle' or import_string
+    packer/unpacker : str : 'orjson', 'json', 'pickle', 'msgpack' or import_string
         importstrings for methods to serialize message parts.  If just
         'json' or 'pickle', predefined JSON and pickle packers will be used.
         Otherwise, the entire importstring must be used.
@@ -351,48 +386,42 @@ class Session(Configurable):
         """,
     )
 
+    # serialization traits:
     packer = DottedObjectName(
-        "json",
+        "orjson" if has_orjson else "json",
         config=True,
         help="""The name of the packer for serializing messages.
             Should be one of 'json', 'pickle', or an import name
             for a custom callable serializer.""",
     )
-
-    @observe("packer")
-    def _packer_changed(self, change: t.Any) -> None:
-        new = change["new"]
-        if new.lower() == "json":
-            self.pack = json_packer
-            self.unpack = json_unpacker
-            self.unpacker = new
-        elif new.lower() == "pickle":
-            self.pack = pickle_packer
-            self.unpack = pickle_unpacker
-            self.unpacker = new
-        else:
-            self.pack = import_item(str(new))
-
     unpacker = DottedObjectName(
-        "json",
+        "orjson" if has_orjson else "json",
         config=True,
         help="""The name of the unpacker for unserializing messages.
         Only used with custom functions for `packer`.""",
     )
-
-    @observe("unpacker")
-    def _unpacker_changed(self, change: t.Any) -> None:
-        new = change["new"]
-        if new.lower() == "json":
-            self.pack = json_packer
-            self.unpack = json_unpacker
-            self.packer = new
-        elif new.lower() == "pickle":
-            self.pack = pickle_packer
-            self.unpack = pickle_unpacker
-            self.packer = new
+    pack = Callable(orjson_packer if has_orjson else json_packer)  # the actual packer function
+    unpack = Callable(
+        orjson_unpacker if has_orjson else json_unpacker
+    )  # the actual unpacker function
+
+    @observe("packer", "unpacker")
+    def _packer_unpacker_changed(self, change: t.Any) -> None:
+        new = change["new"].lower()
+        if new == "orjson" and has_orjson:
+            self.pack, self.unpack = orjson_packer, orjson_unpacker
+        elif new == "json" or new == "orjson":
+            self.pack, self.unpack = json_packer, json_unpacker
+        elif new == "pickle":
+            self.pack, self.unpack = pickle_packer, pickle_unpacker
+        elif new == "msgpack" and has_msgpack:
+            self.pack, self.unpack = msgpack_packer, msgpack_unpacker
         else:
-            self.unpack = import_item(str(new))
+            obj = import_item(str(change["new"]))
+            name = "pack" if change["name"] == "packer" else "unpack"
+            self.set_trait(name, obj)
+            return
+        self.packer = self.unpacker = change["new"]
 
     session = CUnicode("", config=True, help="""The UUID identifying this session.""")
 
@@ -417,8 +446,7 @@ def _session_changed(self, change: t.Any) -> None:
     metadata = Dict(
         {},
         config=True,
-        help="Metadata dictionary, which serves as the default top-level metadata dict for each "
-        "message.",
+        help="Metadata dictionary, which serves as the default top-level metadata dict for each message.",
     )
 
     # if 0, no adapting to do.
@@ -487,25 +515,6 @@ def _keyfile_changed(self, change: t.Any) -> None:
     # for protecting against sends from forks
     pid = Integer()
 
-    # serialization traits:
-
-    pack = Any(default_packer)  # the actual packer function
-
-    @observe("pack")
-    def _pack_changed(self, change: t.Any) -> None:
-        new = change["new"]
-        if not callable(new):
-            raise TypeError("packer must be callable, not %s" % type(new))
-
-    unpack = Any(default_unpacker)  # the actual packer function
-
-    @observe("unpack")
-    def _unpack_changed(self, change: t.Any) -> None:
-        # unpacker is not checked - it is assumed to be
-        new = change["new"]
-        if not callable(new):
-            raise TypeError("unpacker must be callable, not %s" % type(new))
-
     # thresholds:
     copy_threshold = Integer(
         2**16,
@@ -515,8 +524,7 @@ def _unpack_changed(self, change: t.Any) -> None:
     buffer_threshold = Integer(
         MAX_BYTES,
         config=True,
-        help="Threshold (in bytes) beyond which an object's buffer should be extracted to avoid "
-        "pickling.",
+        help="Threshold (in bytes) beyond which an object's buffer should be extracted to avoid pickling.",
     )
     item_threshold = Integer(
         MAX_ITEMS,
@@ -534,7 +542,7 @@ def __init__(self, **kwargs: t.Any) -> None:
 
         debug : bool
             whether to trigger extra debugging statements
-        packer/unpacker : str : 'json', 'pickle' or import_string
+        packer/unpacker : str : 'orjson', 'json', 'pickle', 'msgpack' or import_string
             importstrings for methods to serialize message parts.  If just
             'json' or 'pickle', predefined JSON and pickle packers will be used.
             Otherwise, the entire importstring must be used.
@@ -626,10 +634,7 @@ def _check_packers(self) -> None:
             unpacked = unpack(packed)
             assert unpacked == msg_list
         except Exception as e:
-            msg = (
-                f"unpacker '{self.unpacker}' could not handle output from packer"
-                f" '{self.packer}': {e}"
-            )
+            msg = f"unpacker {self.unpacker!r} could not handle output from packer {self.packer!r}: {e}"
             raise ValueError(msg) from e
 
         # check datetime support

pyproject.toml

@@ -55,6 +55,7 @@ test = [
     "pytest-jupyter[client]>=0.6.2",
     "pytest-cov",
     "pytest-timeout",
+    "msgpack"
 ]
 docs = [
     "ipykernel",
@@ -65,6 +66,7 @@ docs = [
     "sphinxcontrib-spelling",
     "sphinx-autodoc-typehints",
 ]
+orjson = ["orjson"] # When orjson is installed it will be used for faster pack and unpack
 
 [project.scripts]
 jupyter-kernelspec = "jupyter_client.kernelspecapp:KernelSpecApp.launch_instance"