Experiment with environment syncing using conda-pack

[jacobtomlinson]

• Refactored the agent to create workers in subprocesses
• Allowed changing the executable via the reprovision_nodes call
• Added helper to sync environment
  • Environment is packaged with conda-pack
  • Client sends the environment to the scheduler which distributes it to the agents
  • The agents unpack the environment and update the worker executable path
  • The client reprovisions the workers which causes them to use the synced environment

Quick example using Dask Kubernetes

In [1]: from dask_kubernetes.operator import KubeCluster
   ...: from dask.distributed import Client
   ...: cluster = KubeCluster(name="agent", n_workers=3, worker_command="dask-agent", env={"EXTRA_PIP_PACKAGES": "git+https://github.com/jacobtomlinson/dask-agent.git@environment-sync"})
   ...: client = await Client(cluster.scheduler_address, asynchronous=True)
+-------------+----------------+----------------+----------------+
| Package     | Client         | Scheduler      | Workers        |
+-------------+----------------+----------------+----------------+
| cloudpickle | 2.0.0          | 2.2.0          | 2.2.0          |
| lz4         | 3.1.3          | 4.0.2          | 4.0.2          |
| msgpack     | 1.0.3          | 1.0.4          | 1.0.4          |
| numpy       | 1.23.3         | 1.23.4         | 1.23.4         |
| pandas      | 1.4.4          | 1.5.1          | 1.5.1          |
| python      | 3.10.6.final.0 | 3.8.13.final.0 | 3.8.13.final.0 |
+-------------+----------------+----------------+----------------+

Note that the Python versions between the client and workers are mismatched.

In [2]: from dask_agent import sync_env
   ...: await sync_env(client)
Packaging env into /tmp/tmpww5mppg4.tgz...
Packaged
Uploading env...
Uploaded
Reprovisioning nodes...
Reprovisioned

+-------------+----------------+----------------+----------------+
| Package     | Client         | Scheduler      | Workers        |
+-------------+----------------+----------------+----------------+
| cloudpickle | 2.0.0          | 2.2.0          | 2.0.0          |
| dask        | 2022.10.2      | 2022.10.2      | 2022.7.0       |
| distributed | 2022.10.2      | 2022.10.2      | 2022.7.0       |
| lz4         | 3.1.3          | 4.0.2          | 3.1.3          |
| msgpack     | 1.0.3          | 1.0.4          | 1.0.3          |
| python      | 3.10.6.final.0 | 3.8.13.final.0 | 3.10.6.final.0 |
+-------------+----------------+----------------+----------------+

Note that the Python versions are now in sync (but weirdly the dask versions are now out of sync).

Problems

• Conda pack doesn't always seem to do the right thing. In the example above the workers end up with dask/distributed 2022.7.0 got installed with pip when I installed dask-agent, not 2022.10.2 which I installed with conda (and seems to be picked up by default).
• The scheduler isn't included so version mismatches could be problematic there
• Local/editable installs seem to be ignored/broken
• The scheduler sends the environment to agents sequentially.
• It's a bit slow