From 766c070a6c472cc589699d6539fca7847055c3c9 Mon Sep 17 00:00:00 2001 From: John Kirkham Date: Fri, 22 Jan 2021 13:14:53 -0800 Subject: [PATCH 1/2] Build `serialize` with Cython Enables Cythonization of `serialize`, this should improve functions like `extract_serialize`, which have already been annotated for this purpose. --- setup.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/setup.py b/setup.py index 9160e1ffdc9..b36d6f65d26 100755 --- a/setup.py +++ b/setup.py @@ -47,6 +47,10 @@ "distributed.scheduler", sources=["distributed/scheduler.py"], ), + Extension( + "distributed.protocol.serialize", + sources=["distributed/protocol/serialize.py"], + ), ] for e in cyext_modules: e.cython_directives = { From 2cd35244109ffee5d201c59ad3c1715baafb6faa Mon Sep 17 00:00:00 2001 From: Matthew Rocklin Date: Fri, 22 Jan 2021 15:17:27 -0800 Subject: [PATCH 2/2] add context to register serialization function We're no longer able to inspect the cythonized functions --- distributed/protocol/serialize.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py index 1bea7ca496a..e504e59d3b5 100644 --- a/distributed/protocol/serialize.py +++ b/distributed/protocol/serialize.py @@ -157,14 +157,16 @@ def serialization_error_loads(header, frames): families = {} -def register_serialization_family(name, dumps, loads): - families[name] = (dumps, loads, dumps and has_keyword(dumps, "context")) +def register_serialization_family(name, dumps, loads, context=None): + if context is None: + context = has_keyword(dumps, "context") + families[name] = (dumps, loads, dumps and context) -register_serialization_family("dask", dask_dumps, dask_loads) -register_serialization_family("pickle", pickle_dumps, pickle_loads) -register_serialization_family("msgpack", msgpack_dumps, msgpack_loads) -register_serialization_family("error", None, serialization_error_loads) +register_serialization_family("dask", dask_dumps, dask_loads, True) +register_serialization_family("pickle", pickle_dumps, pickle_loads, True) +register_serialization_family("msgpack", msgpack_dumps, msgpack_loads, False) +register_serialization_family("error", None, serialization_error_loads, False) def check_dask_serializable(x):