Repository: dask/distributed

Issue:
Worker initialization function
As discussed in #2110 it would be able to register a function with the scheduler which is guaranteed to run on a worker when it first connects.

> Eventually it would be nice to allow clients to register functions to be run at worker start time with the scheduler that could be passed to workers as they start up.
>
> 1. Client registers func with the scheduler as a preload operation
> 2. Scheduler holds on to func
> 3. Scheduler also sends func to all workers to have them run func
> 4. Workers run func
> 5. New worker arrives, tells Scheduler that it exists
> 6. As part of saying "Hi" the scheduler also gives the worker func and tells it to run it
>
> If anyone is interested in implementing this let me know and I'll point to the various locations in the client, scheduler, and worker, where these changes would have to be made. It's a modest amount of work and would be a good introduction to the distributed system.


Discussion:
It might also be useful to allow registering a `shutdown` function as well?
Agreed

On Tue, Jul 24, 2018 at 6:29 AM, Dave Hirschfeld <notifications@github.com>
wrote:

> It might also be useful to allow registering a shutdown function as well?
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/2140#issuecomment-407359248>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszGbdDx05KgjgCITxO4xKfzYupI-0ks5uJvbvgaJpZM4Vcey2>
> .
>

Hi @dhirschfeld, @mrocklin, I'd like to contribute to distributed directly, is this a good place to start?
Hi @guillaumeeb , yes I think that this could be an educational place to start.  FWIW I'd love to get more people engaged in the distributed scheduler, and I think that you would be a great help to the project.

This topic would be educational because it touches the client, scheduler, and worker, and will force you to become familiar with a bit of async programming.  However, it also isn't a trivial issue and will require some work to get right.  I would expect this to take some time and a fair amount of interaction (unless you're already familiar with the codebase)

I would expect this to require the following steps:

### Test

Build a `gen_cluster` style test that uses the client to submit a startup function, verify that that function affects existing workers (two are started up by default).  Then create a new worker and verify that that worker has also been affected.

I would place this in `tests/test_worker.py`, but this is somewhat arbitrary

### Write some code

Affect code the following three places:
    
Where the worker registers with the scheduler, and receives a response back from it.  This is a good time to run code after it has connected, but before it has had a chance to do any work.

https://github.com/dask/distributed/blob/770b4afa51bbc8e5fbe77d6f6b5d2fbe1f69a7bc/distributed/worker.py#L299-L320

Where the scheduler hears the worker's initial message, and where it sends data down to it

https://github.com/dask/distributed/blob/770b4afa51bbc8e5fbe77d6f6b5d2fbe1f69a7bc/distributed/scheduler.py#L1294-L1296

Separately the client will need to add a method that takes a function from the user and submits that to the scheduler for storage.  This will require changes on the client, I recommend looking at `_cancel` and `cancel` for a simple example of how to create a dual async/sync function on the client side.  This will also require changes on the scheduler, look at `self.handlers`

https://github.com/dask/distributed/blob/master/distributed/scheduler.py#L902

The method on the scheduler should be a coroutine that does the following:

1.  Uses the `Scheduler.broadcast` method to run the function on all of the existing workers (this will use the worker's `'run'` route, see `Client.run` for an example)
2.  Stores the serialized function on the scheduler in a place to be used later in `add_worker`

### Setting expectations

Getting all of that to work well on the first try is likely to be challenging.  I recommend checking in with a WIP PR early and often.  Given how you've stuck around and helped with dask-jobqueue I'm more than willing to spend time helping out here.  I suspect that if you try to do everything without asking for help then you will quickly become frustrated and burn out.
Thanks @mrocklin, I'll try to look at all that calmly next week! 
It might be nice to extend this functionality to upload_file (so that you can upload a file to all existing workers and have that file also be uploaded to all new workers). As far as I can tell (and I'm not an expert on Dask by any stretch of the imagination), `client.run` and `client.upload_file` are very similar (the only difference is the msg they ask the scheduler to broadcast), so maybe these two features are also very similar.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -3425,6 +3425,37 @@ def _get_task_stream(self, start=None, stop=None, count=None, plot=False,
         else:
             raise gen.Return(msgs)
 
+    @gen.coroutine
+    def _register_worker_callbacks(self, setup=None):
+        responses = yield self.scheduler.register_worker_callbacks(setup=dumps(setup))
+        results = {}
+        for key, resp in responses.items():
+            if resp['status'] == 'OK':
+                results[key] = resp['result']
+            elif resp['status'] == 'error':
+                six.reraise(*clean_exception(**resp))
+        raise gen.Return(results)
+
+    def register_worker_callbacks(self, setup=None):
+        """
+        Registers a setup callback function for all current and future workers.
+
+        This registers a new setup function for workers in this cluster. The
+        function will run immediately on all currently connected workers. It
+        will also be run upon connection by any workers that are added in the
+        future. Multiple setup functions can be registered - these will be
+        called in the order they were added.
+
+        If the function takes an input argument named ``dask_worker`` then
+        that variable will be populated with the worker itself.
+
+        Parameters
+        ----------
+        setup : callable(dask_worker: Worker) -> None
+            Function to register and run on all workers
+        """
+        return self.sync(self._register_worker_callbacks, setup=setup)
+
 
 class Executor(Client):
     """ Deprecated: see Client """
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -899,6 +899,7 @@ def __init__(
         self.plugins = []
         self.transition_log = deque(maxlen=dask.config.get('distributed.scheduler.transition-log-length'))
         self.log = deque(maxlen=dask.config.get('distributed.scheduler.transition-log-length'))
+        self.worker_setups = []
 
         worker_handlers = {
             'task-finished': self.handle_task_finished,
@@ -956,6 +957,7 @@ def __init__(
             'heartbeat_worker': self.heartbeat_worker,
             'get_task_status': self.get_task_status,
             'get_task_stream': self.get_task_stream,
+            'register_worker_callbacks': self.register_worker_callbacks
         }
 
         self._transitions = {
@@ -1330,7 +1332,8 @@ def add_worker(self, comm=None, address=None, keys=(), ncores=None,
 
             yield comm.write({'status': 'OK',
                               'time': time(),
-                              'heartbeat-interval': heartbeat_interval(len(self.workers))})
+                              'heartbeat-interval': heartbeat_interval(len(self.workers)),
+                              'worker-setups': self.worker_setups})
             yield self.handle_worker(comm=comm, worker=address)
 
     def update_graph(self, client=None, tasks=None, keys=None,
@@ -3035,6 +3038,17 @@ def get_task_stream(self, comm=None, start=None, stop=None, count=None):
         ts = [p for p in self.plugins if isinstance(p, TaskStreamPlugin)][0]
         return ts.collect(start=start, stop=stop, count=count)
 
+    @gen.coroutine
+    def register_worker_callbacks(self, comm, setup=None):
+        """ Registers a setup function, and call it on every worker """
+        if setup is None:
+            raise gen.Return({})
+
+        self.worker_setups.append(setup)
+
+        responses = yield self.broadcast(msg=dict(op='run', function=setup))
+        raise gen.Return(responses)
+
     #####################
     # State Transitions #
     #####################
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -342,6 +342,15 @@ def _register_with_scheduler(self):
             raise ValueError("Unexpected response from register: %r" %
                              (response,))
         else:
+            # Retrieve eventual init functions and run them
+            for function_bytes in response['worker-setups']:
+                setup_function = pickle.loads(function_bytes)
+                if has_arg(setup_function, 'dask_worker'):
+                    result = setup_function(dask_worker=self)
+                else:
+                    result = setup_function()
+                logger.info('Init function %s ran: output=%s' % (setup_function, result))
+
             logger.info('        Registered to: %26s', self.scheduler.address)
             logger.info('-' * 49)
 
@@ -567,10 +576,11 @@ def executor_submit(self, key, function, args=(), kwargs=None,
         # logger.info("Finish job %d, %s", i, key)
         raise gen.Return(result)
 
-    def run(self, comm, function, args=(), kwargs={}):
+    def run(self, comm, function, args=(), kwargs=None):
+        kwargs = kwargs or {}
         return run(self, comm, function=function, args=args, kwargs=kwargs)
 
-    def run_coroutine(self, comm, function, args=(), kwargs={}, wait=True):
+    def run_coroutine(self, comm, function, args=(), kwargs=None, wait=True):
         return run(self, comm, function=function, args=args, kwargs=kwargs,
                    is_coro=True, wait=wait)
 


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -1203,3 +1203,76 @@ def test_avoid_oversubscription(c, s, *workers):
 def test_custom_metrics(c, s, a, b):
     assert s.workers[a.address].metrics['my_port'] == a.port
     assert s.workers[b.address].metrics['my_port'] == b.port
+
+
+@gen_cluster(client=True)
+def test_register_worker_callbacks(c, s, a, b):
+    #preload function to run
+    def mystartup(dask_worker):
+        dask_worker.init_variable = 1
+
+    def mystartup2():
+        import os
+        os.environ['MY_ENV_VALUE'] = 'WORKER_ENV_VALUE'
+        return "Env set."
+
+    #Check that preload function has been run
+    def test_import(dask_worker):
+        return hasattr(dask_worker, 'init_variable')
+        #       and dask_worker.init_variable == 1
+
+    def test_startup2():
+        import os
+        return os.getenv('MY_ENV_VALUE', None) == 'WORKER_ENV_VALUE'
+
+    # Nothing has been run yet
+    assert len(s.worker_setups) == 0
+    result = yield c.run(test_import)
+    assert list(result.values()) == [False] * 2
+    result = yield c.run(test_startup2)
+    assert list(result.values()) == [False] * 2
+
+    # Start a worker and check that startup is not run
+    worker = Worker(s.address, loop=s.loop)
+    yield worker._start()
+    result = yield c.run(test_import, workers=[worker.address])
+    assert list(result.values()) == [False]
+    yield worker._close()
+
+    # Add a preload function
+    response = yield c.register_worker_callbacks(setup=mystartup)
+    assert len(response) == 2
+    assert len(s.worker_setups) == 1
+
+    # Check it has been ran on existing worker
+    result = yield c.run(test_import)
+    assert list(result.values()) == [True] * 2
+
+    # Start a worker and check it is ran on it
+    worker = Worker(s.address, loop=s.loop)
+    yield worker._start()
+    result = yield c.run(test_import, workers=[worker.address])
+    assert list(result.values()) == [True]
+    yield worker._close()
+
+    # Register another preload function
+    response = yield c.register_worker_callbacks(setup=mystartup2)
+    assert len(response) == 2
+    assert len(s.worker_setups) == 2
+
+    # Check it has been run
+    result = yield c.run(test_startup2)
+    assert list(result.values()) == [True] * 2
+
+    # Start a worker and check it is ran on it
+    worker = Worker(s.address, loop=s.loop)
+    yield worker._start()
+    result = yield c.run(test_import, workers=[worker.address])
+    assert list(result.values()) == [True]
+    result = yield c.run(test_startup2, workers=[worker.address])
+    assert list(result.values()) == [True]
+    yield worker._close()
+
+    # Final exception test
+    with pytest.raises(ZeroDivisionError):
+        yield c.register_worker_callbacks(setup=lambda: 1 / 0)