[WIP] Enable Skorch+Dask-ML

distributed/protocol/__init__.py

@@ -119,3 +119,9 @@ def _register_cudf():
 def _register_cuml():
     with suppress(ImportError):
         from cuml.comm import serialize
+
+
+@dask_serialize.register_lazy("skorch")
+@dask_deserialize.register_lazy("skorch")
+def _register_skorch():
+    from . import skorch

distributed/protocol/skorch.py

@@ -0,0 +1,42 @@
+import skorch
+
+from . import pickle
+from .serialize import dask_deserialize, dask_serialize
+
+
+@dask_serialize.register(skorch.NeuralNet)
+def serialize_skorch(x, context=None):
+    protocol = (context or {}).get("pickle-protocol", None)
+    headers = {}
+    has_module = hasattr(x, "module_")
+    if has_module:
+        module = x.__dict__.pop("module_")
+        # module's is an interactively defined class on client so its namespace is often `__main__` .
+        # Pickle has problems pickling when interactively defined classes  when they are
+        # set as an attributes of another object.
+        # By pickling it on its own we are able to serialize successfully
+        frames = [None]
+        buffer_callback = lambda f: frames.append(memoryview(f))
+        frames[0] = pickle.dumps(x, buffer_callback=buffer_callback, protocol=protocol)
+        headers["subframe-split"] = i = len(frames)
+        frames.append(None)
+        frames[i] = pickle.dumps(
+            module, buffer_callback=buffer_callback, protocol=protocol
+        )
+        x.__dict__["module_"] = module
+    else:
+        frames = [None]
+        buffer_callback = lambda f: frames.append(memoryview(f))
+        frames[0] = pickle.dumps(x, buffer_callback=buffer_callback, protocol=protocol)
+
+    return headers, frames
+
+
+@dask_deserialize.register(skorch.NeuralNet)
+def deserialize_skorch(header, frames):
+    i = header.get("subframe-split")
+    model = pickle.loads(frames[0], buffers=frames[1:i])
+    if i is not None:
+        module = pickle.loads(frames[i], buffers=frames[i + 1 :])
+        model.module_ = module
+    return model

distributed/protocol/tests/test_skorch.py

@@ -0,0 +1,42 @@
+import pytest
+
+skorch = pytest.importorskip("skorch")
+torch = pytest.importorskip("torch")
+
+from distributed import Client
+from distributed.protocol import deserialize, serialize
+
+
+def test_serialize_deserialize_skorch_model():
+
+    client = Client(processes=True, n_workers=1)
+
+    class MyModule(torch.nn.Module):
+        def __init__(self, num_units=10):
+            super().__init__()
+            self.dense0 = torch.nn.Linear(20, num_units)
+
+        def forward(self, X, **kwargs):
+            return self.dense0(X)
+
+    net = skorch.NeuralNetClassifier(
+        MyModule,
+        max_epochs=10,
+        iterator_train__shuffle=True,
+    )
+
+    def test_serialize_skorch(net):
+        net = net.initialize()
+        return deserialize(*serialize(net))
+
+    # We test on a different worker to ensure that
+    # errors skorch serialization faces on a different process
+    # other than the client due to lack of __main__ context
+    # are actually resolved
+    # See isssue for context
+    # https://github.com/dask/dask-ml/issues/549#issuecomment-669924762
+
+    deserialized_net = list(client.run(test_serialize_skorch, net).values())[0]
+    assert isinstance(deserialized_net.module_, MyModule)
+
+    client.close()