Add a package to Dask workers?

[rsignell-usgs]

Is there any way to add a package to the environment used by the Planetary Computer hub Dask workers?

I guess one way would be to set up my own hub and create custom environments there, but hoping there is an easier way...

[TomAugspurger]

The easiest way, at least for pure-python packages, is to use Dask's PipInstall plugin to install additional packages at runtime. See https://planetarycomputer.microsoft.com/docs/overview/environment/#installing-packages-on-dask-clusters-at-runtime and the linked docs.

>>> from dask.distributed import PipInstall
>>> plugin = PipInstall(packages=["scikit-learn"], pip_options=["--upgrade"])
>>> client.register_worker_plugin(plugin)

A few more complicated options are to use worker plugins to execute shell commands like mamba install ..., or specifying a custom docker image to use when you create the cluster (https://planetarycomputer.microsoft.com/docs/quickstarts/scale-with-dask/#Customize-your-cluster).

[rsignell-usgs]

Awesome @TomAugspurger ! I'm a bit embarrassed I didn't find that when I was looking in the docs, but this is great! I didn't know you could specify a custom docker image!