Implementing a WorkerPlugin for uploading a directory containing Python code or at least updating docs

[guillaumeeb]

Hi everyone,

Just wanted to discuss the UploadDirectory plugin, and propose some changes at least to documentation regarding it.

This is both an issue and a PR, as I dug a bit into plugins and tried implementing one, I thought it would be easier to discuss with a proposal.

So I am in a situation where I work on an under development Python project (it's actually https://github.com/pangeo-data/benchmarking that I'm reviving), I need to test some code on a dask-gateway cluster. This project was developed on HPC Centers, so it assumed a shared folder where the code is available for every workers, which is not the case in the Cloud. The project folder look like this:

benchmarking/
-- benckmarks/
---- datasets.py
---- ops.py
---- utils.py

Quite simple!

What I needed was to be able to run something like:

def f():
    from benchmarks.ops import DevNullStore
    return DevNullStore()

On the workers.

Searching around, I came into https://distributed.dask.org/en/stable/plugins.html#distributed.diagnostics.plugin.UploadDirectory and some explanations in https://docs.coiled.io/user_guide/tutorials/upload_file_to_coiled.html#upload-a-local-directory.

I though that was it (and I discovered later this was actually it 😄) but I was not able to get it working. I was either trying to upload the wrong directory (e.g. benchmarks), either not giving the correct arguments (no update_path and restart kwargs). So I developed the UploadDirectoryWorker version proposed here as an intermediate.

When preparing this PR, I discovered on the distributed tests that I did not try everything, the documentation I sought was here: https://github.com/dask/distributed/blob/main/distributed/tests/test_client.py#L7032.

In order to make my uploaded code able to be imported, I have to use:

plugin = UploadDirectory('/some/path/to/benchmarking', restart=True, update_path=True)

With the proposed WorkerPlugin, it also works with:

plugin = UploadDirectory('/some/path/to/benchmarking/benchmarks')

So I'm happy to just propose some clarifications in the docs (but I'm not sure were except in the docstring?), and I wanted to hear if you think adding this UploadDirectoryWorker plugin is useful or not.

[github-actions]

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

       15 files  ±0         15 suites  ±0   5h 58m 30s ⏱️ - 26m 56s
  3 113 tests +1    3 026 ✔️ ±0    85 💤 ±0  2 ❌ +1 
23 042 runs  +8  22 127 ✔️ ±0  906 💤 ±0  9 ❌ +8 

For more details on these failures, see this check.

Results for commit 4c66702. ± Comparison against base commit 17c70e9.