Add worker_class argument to LocalCluster

distributed/deploy/local.py

@@ -74,6 +74,8 @@ class LocalCluster(Cluster):
         Protocol to use like ``tcp://``, ``tls://``, ``inproc://``
         This defaults to sensible choice given other keyword arguments like
         ``processes`` and ``security``
+    worker_class: Worker
+        Worker class used to instantiate workers from.
 
     Examples
     --------
@@ -115,6 +117,7 @@ def __init__(
         security=None,
         protocol=None,
         blocked_handlers=None,
+        worker_class=None,
         **worker_kwargs
     ):
         if start is not None:
@@ -203,6 +206,10 @@ def __init__(
         if security:
             self.worker_kwargs["security"] = security
 
+        if not worker_class:
+            worker_class = Worker if not processes else Nanny
+        self.worker_class = worker_class
+
         self.start(ip=ip, n_workers=n_workers)
 
         clusters_to_close.add(self)
@@ -279,12 +286,9 @@ def _start_worker(self, death_timeout=60, **kwargs):
             return
 
         if self.processes:
-            W = Nanny
             kwargs["quiet"] = True
-        else:
-            W = Worker
 
-        w = yield W(
+        w = yield self.worker_class(
             self.scheduler.address,
             loop=self.loop,
             death_timeout=death_timeout,

distributed/deploy/tests/test_local.py

@@ -742,5 +742,35 @@ def test_protocol_ip(loop):
         assert cluster.scheduler.address.startswith("tcp://127.0.0.2")
 
 
+class MyWorker(Worker):
+    pass
+
+
+def test_worker_class_worker(loop):
+    with LocalCluster(
+        n_workers=2,
+        loop=loop,
+        worker_class=MyWorker,
+        processes=False,
+        scheduler_port=0,
+        dashboard_address=None,
+    ) as cluster:
+        assert all(isinstance(w, MyWorker) for w in cluster.workers)
+
+
+def test_worker_class_nanny(loop):
+    class MyNanny(Nanny):
+        pass
+
+    with LocalCluster(
+        n_workers=2,
+        loop=loop,
+        worker_class=MyNanny,
+        scheduler_port=0,
+        dashboard_address=None,
+    ) as cluster:
+        assert all(isinstance(w, MyNanny) for w in cluster.workers)
+
+
 if sys.version_info >= (3, 5):
     from distributed.deploy.tests.py3_test_deploy import *  # noqa F401