Repository: dask/distributed

Issue:
Upload file to new workers only
Currently, `upload_file` will only send files to workers connected to scheduler at an instant = t.
If a worker connects to the scheduler afterwards, it does not have the data send with `upload_file` beforehand.

This is a problem because then we cannot assume that all the workers have the file available.

1. One simple workaround is to store the identity of the connected workers just before calling `upload_file` and only send jobs requiring this file to these workers in the future. But this prevent scalability.

2. An other possibility would be to register a `SchedulerPlugin` and use its `add_worker` event function to trigger there resending the file with  `upload_file` to all the workers when a new worker joins. If this solution is preferred, then it might be good that `upload_file` provides a `worker` parameter, so that we can target the newly connected worker and avoid reuploading the file to other workers already having the file.

3. As an extension of solution 2, add a `future_workers` bool parameter to `upload_file` and extend `upload_file` to automatically register the `SchedulerPlugin` described above.

4. Other ideas?


Discussion:
This could be a worker plugin, which already has the machinery to track new
workers

On Sat, Mar 7, 2020 at 10:54 AM H4dr1en <notifications@github.com> wrote:

> Currently, upload_file will only send files to workers connected to
> scheduler at an instant = t.
> If a worker connects to the scheduler afterwards, it does not have the
> data send with upload_file beforehand.
> This is a problem because then we cannot assume that all the workers have
> the file uploaded.
>
>    1.
>
>    One simple workaround is to store the identity of the workers just
>    before calling upload_file and only send jobs requiring this file to
>    these workers in the future. But this prevent scalability.
>    2.
>
>    An other possibility would be to use client.add_worker and trigger
>    there resending the file with upload_file to all the workers. If this
>    solution is preferred, then it might be good that upload_file provide
>    a worker parameter, so that we can target the newly connected worker
>    and avoid reuploading the file to other workers already having the file.
>    3.
>
>    Other ideas?
>
> —
> 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/3560?email_source=notifications&email_token=AACKZTCCRIN7AIQWRCRGYWLRGJU3FA5CNFSM4LDQW4EKYY3PNVWWK3TUL52HS4DFUVEXG43VMWVGG33NNVSW45C7NFSM4ITJW5SQ>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTEILQLTKDIKDUARKRTRGJU3FANCNFSM4LDQW4EA>
> .
>

@H4dr1en FWIW I've used the following `WorkerPlugin` in the past:

```python

import os
import asyncio
from distributed.diagnostics.plugin import WorkerPlugin

class UploadFile(WorkerPlugin):
    """ Upload local package to workers

    Parameters
    ----------
    filepath : str, List[str]
        Filename of .py, .egg or .zip file to send to workers

    Examples
    --------
    Upload a local Python module, for example ``foo.py``, to workers:

    >>> plugin = UploadFile("foo.py")
    >>> client.register_worker_plugin(plugin)
    """

    name = "upload_file"

    def __init__(self, filepath):
        if isinstance(filepath, str):
            filepath = [filepath]

        self.data = {}
        for file_ in filepath:
            with open(file_, "rb") as f:
                filename = os.path.basename(file_)
                self.data[filename] = f.read()

    async def setup(self, worker):

        responses = await asyncio.gather(
            *[
                worker.upload_file(comm=None, filename=filename, data=data, load=True)
                for filename, data in self.data.items()
            ]
        )

        assert all(
            len(data) == r["nbytes"] for r, data in zip(responses, self.data.values())
        )
```

It may, or may not, serve as a good starting point for your use case 
Very cool! Thanks @jrbourbeau for the code, I was not aware of the `worker.upload_file` (not appearing in the docs when searching):
https://distributed.dask.org/en/latest/search.html?q=upload_file&check_keywords=yes&area=default
@mrocklin Is there a reason that the default behavior isn't to initialize new workers with files that have been uploaded to the Client? In my case I went from manually starting schedulers and workers on the command line, in this case uploading to the client once works. But now I am transporting my code onto a HPC system using the PBSqueue tool and the upload behavior is now quite different. Having a look at stack overflow, it seems that this is a common pitfall for Dask users. If this step was automated, the users would never have to consider in which order workers arrive compared to the client and in IMO the Api would be more consistent across all of the ways of initializing schedulers clients etc... Given how generally robust Dask is to adding and subtracting workers, this might help less sophisticated users (like me) avoid a nasty bug. Just a friendly question from long time fan of the project. 
I'm open to that change.  Presumably we would change the `client.upload_file` method to use a worker plugin much like @jrbourbeau has provided above.  This may slightly disrupt existing users of upload_file due to the changed semantics, but  I suspect that this would fix more problems than it would create. 

@MatthewLennie are you interested in doing this work?
@mrocklin It's a bit above my skill level as it currently stands, having just transitioned from mechanical engineering to SE. Interested yes, capable not yet. Sorry to be no help yet. 

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -69,7 +69,7 @@
 from .sizeof import sizeof
 from .threadpoolexecutor import rejoin
 from .worker import get_client, get_worker, secede
-from .diagnostics.plugin import WorkerPlugin
+from .diagnostics.plugin import UploadFile, WorkerPlugin
 from .utils import (
     All,
     sync,
@@ -3034,23 +3034,6 @@ def restart(self, **kwargs):
         """
         return self.sync(self._restart, **kwargs)
 
-    async def _upload_file(self, filename, raise_on_error=True):
-        with open(filename, "rb") as f:
-            data = f.read()
-        _, fn = os.path.split(filename)
-        d = await self.scheduler.broadcast(
-            msg={"op": "upload_file", "filename": fn, "data": to_serialize(data)}
-        )
-
-        if any(v["status"] == "error" for v in d.values()):
-            exceptions = [v["exception"] for v in d.values() if v["status"] == "error"]
-            if raise_on_error:
-                raise exceptions[0]
-            else:
-                return exceptions[0]
-
-        assert all(len(data) == v["nbytes"] for v in d.values())
-
     async def _upload_large_file(self, local_filename, remote_filename=None):
         if remote_filename is None:
             remote_filename = os.path.split(local_filename)[1]
@@ -3094,13 +3077,10 @@ def upload_file(self, filename, **kwargs):
         >>> from mylibrary import myfunc  # doctest: +SKIP
         >>> L = client.map(myfunc, seq)  # doctest: +SKIP
         """
-        result = self.sync(
-            self._upload_file, filename, raise_on_error=self.asynchronous, **kwargs
+        return self.register_worker_plugin(
+            UploadFile(filename),
+            name=filename + str(uuid.uuid4()),
         )
-        if isinstance(result, Exception):
-            raise result
-        else:
-            return result
 
     async def _rebalance(self, futures=None, workers=None):
         await _wait(futures)
@@ -4109,7 +4089,7 @@ def register_worker_plugin(self, plugin=None, name=None, **kwargs):
 
         If the plugin has a ``name`` attribute, or if the ``name=`` keyword is
         used then that will control idempotency.  If a plugin with that name has
-        already registered then any future plugins will not run.
+        already been registered then any future plugins will not run.
 
         For alternatives to plugins, you may also wish to look into preload
         scripts.
diff --git a/distributed/diagnostics/plugin.py b/distributed/diagnostics/plugin.py
--- a/distributed/diagnostics/plugin.py
+++ b/distributed/diagnostics/plugin.py
@@ -1,4 +1,5 @@
 import logging
+import os
 import socket
 import subprocess
 import sys
@@ -258,3 +259,36 @@ async def setup(self, worker):
                     worker.loop.add_callback(
                         worker.close_gracefully, restart=True
                     )  # restart
+
+
+# Adapted from https://github.com/dask/distributed/issues/3560#issuecomment-596138522
+class UploadFile(WorkerPlugin):
+    """A WorkerPlugin to upload a local file to workers.
+
+    Parameters
+    ----------
+    filepath: str
+        A path to the file (.py, egg, or zip) to upload
+
+    Examples
+    --------
+    >>> from distributed.diagnostics.plugin import UploadFile
+
+    >>> client.register_worker_plugin(UploadFile("/path/to/file.py"))  # doctest: +SKIP
+    """
+
+    name = "upload_file"
+
+    def __init__(self, filepath):
+        """
+        Initialize the plugin by reading in the data from the given file.
+        """
+        self.filename = os.path.basename(filepath)
+        with open(filepath, "rb") as f:
+            self.data = f.read()
+
+    async def setup(self, worker):
+        response = await worker.upload_file(
+            comm=None, filename=self.filename, data=self.data, load=True
+        )
+        assert len(self.data) == response["nbytes"]
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -999,7 +999,7 @@ def func(data):
                 cache_loads.data.clear()
             except Exception as e:
                 logger.exception(e)
-                return {"status": "error", "exception": to_serialize(e)}
+                raise e
 
         return {"status": "OK", "nbytes": len(data)}
 


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -1684,6 +1684,21 @@ def test_upload_file_exception_sync(c):
             c.upload_file(fn)
 
 
+@gen_cluster(client=True, nthreads=[])
+async def test_upload_file_new_worker(c, s):
+    def g():
+        import myfile
+
+        return myfile.x
+
+    with tmp_text("myfile.py", "x = 123") as fn:
+        await c.upload_file(fn)
+        async with Worker(s.address):
+            x = await c.submit(g)
+
+        assert x == 123
+
+
 @pytest.mark.skip
 @gen_cluster()
 async def test_multiple_clients(s, a, b):