Repository: dask/distributed

Issue:
scheduler deserialization breaks PipInstall
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**: The `PipInstall` worker plugin installs packages on the worker, but recent changes to how the scheduler is deserializing might make code that tries to use this package break, if the scheduler and workers do not share a filesystem. I believe this is the same issue as described in #7834 but for `PipInstall` instead of `Client.upload_file`.

**Minimal Complete Verifiable Example**:
First a remote dask cluster. In our case, we are using a kuberentes-backed DaskGateway cluster. In the below example `client` is the `Client` object. And let's say, for example, that the docker image being used does not contain `s3fs`. We have pip installed s3fs locally and would like to use it in dask workflows as well

```python
from distributed import PipInstall
import s3fs

client.register_worker_plugin(PipInstall(["s3fs"]))

def testfunc():
    return s3fs.__version__

client.submit(testfunc).result()
```
Gives:
```python
ModuleNotFoundError                       Traceback (most recent call last)
File /srv/conda/envs/notebook/lib/python3.11/site-packages/distributed/scheduler.py:4545, in update_graph()
   4544 try:
-> 4545     graph = deserialize(graph_header, graph_frames).data
   4546     del graph_header, graph_frames

File /srv/conda/envs/notebook/lib/python3.11/site-packages/distributed/protocol/serialize.py:432, in deserialize()
    431 dumps, loads, wants_context = families[name]
--> 432 return loads(header, frames)

File /srv/conda/envs/notebook/lib/python3.11/site-packages/distributed/protocol/serialize.py:98, in pickle_loads()
     93 buffers = [
     94     memoryview(bytearray(mv) if w else bytes(mv)) if w == mv.readonly else mv
     95     for w, mv in zip(writeable, map(ensure_memoryview, buffers))
     96 ]
---> 98 return pickle.loads(x, buffers=buffers)

File /srv/conda/envs/notebook/lib/python3.11/site-packages/distributed/protocol/pickle.py:96, in loads()
     95     else:
---> 96         return pickle.loads(x)
     97 except Exception:

File /srv/conda/envs/notebook/lib/python3.11/site-packages/cloudpickle/cloudpickle.py:649, in subimport()
    644 def subimport(name):
    645     # We cannot do simply: `return __import__(name)`: Indeed, if ``name`` is
    646     # the name of a submodule, __import__ will return the top-level root module
    647     # of this submodule. For instance, __import__('os.path') returns the `os`
    648     # module.
--> 649     __import__(name)
    650     return sys.modules[name]

ModuleNotFoundError: No module named 's3fs'

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[21], line 1
----> 1 client.submit(tryimport).result()

File /srv/conda/envs/notebook/lib/python3.11/site-packages/distributed/client.py:327, in Future.result(self, timeout)
    325 if self.status == "error":
    326     typ, exc, tb = result
--> 327     raise exc.with_traceback(tb)
    328 elif self.status == "cancelled":
    329     raise result

RuntimeError: Error during deserialization of the task graph. This frequently occurs if the Scheduler and Client have different environments. For more information, see https://docs.dask.org/en/stable/deployment-considerations.html#consistent-software-environments
```

Note that the following would work (if the package is imported within the function). This is presumably b/c the scheduler no longer needs to deserialize s3fs into the python package:
```python
def testfunc():
    import s3fs
    return s3fs.__version__

client.submit(testfunc).result()
```

**Environment**:

- Dask version: 2023.8.0
- Python version: 3.11.4
- Operating System: Linux



Discussion:
Thanks for the issue @bolliger32 

For the `Client.upload_file` case we added a corresponding scheduler plugin to make sure files were also uploaded on the scheduler. We should probably do something similar for `PipInstall` / `CondaInstall`. That said, needing to manually register both types of plugins along the lintes of

```python
client.register_worker_plugin(PipInstall(["s3fs"]))
client.register_scheduler_plugin(PipInstallScheduler(["s3fs"]))    # a better name is welcome
```

is pretty verbose. It's a separate issue, but a couple of other follow-ups could be to:

1. Make a new type of plugin (`ClusterPlugin`?) that applies to both workers and the scheduler. This may be more complexity than it's worth
2. Add a new `Client.pip_install(...)` method that serves as the user-facing API for installing packages. This has been brought up before and would move the whole install on scheduler + workers complexity to an implementation detail users don't need to think about. 

> Note that the following would work (if the package is imported within the function). This is presumably b/c the scheduler no longer needs to deserialize s3fs into the python package

Yeah, that's correct. This would be my suggested temporary workaround
Thanks @jrbourbeau this makes a lot of sense to me. FWIW I like the idea of a Client.pip_install method that works like upload_file
Also, in case it's helpful, the workaround we're using right now goes something like this:

```python
pip_installer = PipInstall(
    packages=extra_pip_packages, pip_options=extra_pip_install_options
)
client.register_worker_plugin(pip_installer)
client.run_on_scheduler(pip_installer.install)
```
Probably not the most robust thing to do in the long run though.
FWIW, instead of introducing a `ClusterPlugin`, we could also turn the `PackageInstall` plugin into a `SchedulerPlugin` which then installs a corresponding worker plugin. This might be a good compromise for the short-/mid-term. The P2P scheduler plugin does something similar.
FWIW the same issue pops up with `UploadDirectory`, where the use case here is uploading a local development version of a package to the workers (and now to the scheduler). Here the workaround is not quite so straightforward, or at least I haven't figured it out. Curious if there are recommendations on a workaround for the time being

Resolving patch:
diff --git a/distributed/diagnostics/plugin.py b/distributed/diagnostics/plugin.py
--- a/distributed/diagnostics/plugin.py
+++ b/distributed/diagnostics/plugin.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import asyncio
 import contextlib
 import functools
 import logging
@@ -16,6 +17,8 @@
 from dask.typing import Key
 from dask.utils import funcname, tmpfile
 
+from distributed.protocol.pickle import dumps
+
 if TYPE_CHECKING:
     # circular imports
     from distributed.scheduler import Scheduler
@@ -345,10 +348,10 @@ async def start(self, scheduler: Scheduler) -> None:
         await scheduler.upload_file(self.filename, self.data, load=self.load)
 
 
-class PackageInstall(WorkerPlugin, abc.ABC):
-    """Abstract parent class for a worker plugin to install a set of packages
+class PackageInstall(SchedulerPlugin, abc.ABC):
+    """Abstract parent class for a scheduler plugin to install a set of packages
 
-    This accepts a set of packages to install on all workers.
+    This accepts a set of packages to install on the scheduler and all workers.
     You can also optionally ask for the worker to restart itself after
     performing this installation.
 
@@ -363,7 +366,7 @@ class PackageInstall(WorkerPlugin, abc.ABC):
     ----------
     packages
         A list of packages (with optional versions) to install
-    restart
+    restart_workers
         Whether or not to restart the worker after installing the packages
         Only functions if the worker has an attached nanny process
 
@@ -373,20 +376,111 @@ class PackageInstall(WorkerPlugin, abc.ABC):
     PipInstall
     """
 
-    INSTALLER: ClassVar[str]
+    _lock: ClassVar[asyncio.Lock | None] = None
+    WORKER_PLUGIN_CLASS: ClassVar[type]
 
+    installer: _PackageInstaller
     name: str
-    packages: list[str]
-    restart: bool
+    restart_workers: bool
+    _scheduler: Scheduler
 
     def __init__(
         self,
-        packages: list[str],
-        restart: bool,
+        installer: _PackageInstaller,
+        restart_workers: bool,
     ):
-        self.packages = packages
+        self.installer = installer
+        self.restart_workers = restart_workers
+        self.name = f"{self.installer.INSTALLER}-install-{uuid.uuid4()}"
+
+    async def start(self, scheduler: Scheduler) -> None:
+        self._scheduler = scheduler
+
+        if PackageInstall._lock is None:
+            PackageInstall._lock = asyncio.Lock()
+
+        async with PackageInstall._lock:
+            if not self._is_installed(scheduler):
+                logger.info(
+                    "%s installing the following packages on the scheduler: %s",
+                    self.installer.INSTALLER,
+                    self.installer.packages,
+                )
+                self.installer.install()
+                self._set_installed(scheduler)
+
+                worker_plugin = _PackageInstallWorker(
+                    self.installer, self.restart_workers, self.name
+                )
+                await scheduler.register_worker_plugin(
+                    comm=None, plugin=dumps(worker_plugin), name=self.name
+                )
+            else:
+                logger.info(
+                    "The following packages have already been installed on the scheduler: %s",
+                    self.installer.packages,
+                )
+
+    async def close(self) -> None:
+        assert PackageInstall._lock is not None
+
+        async with PackageInstall._lock:
+            await self._scheduler.unregister_worker_plugin(comm=None, name=self.name)
+
+    def _is_installed(self, scheduler: Scheduler) -> bool:
+        return scheduler.get_metadata(self._compose_installed_key(), default=False)
+
+    def _set_installed(self, scheduler: Scheduler) -> None:
+        scheduler.set_metadata(
+            self._compose_installed_key(),
+            True,
+        )
+
+    def _compose_installed_key(self) -> list[Key]:
+        return [
+            self.name,
+            "installed",
+            socket.gethostname(),
+        ]
+
+
+class _PackageInstallWorker(WorkerPlugin):
+    """Worker plugin to install a set of packages
+
+    This accepts an installer which will install a set of packages on all workers.
+    You can also optionally ask for the worker to restart itself after
+    performing this installation.
+
+    .. note::
+
+       This will increase the time it takes to start up
+       each worker. If possible, we recommend including the
+       libraries in the worker environment or image. This is
+       primarily intended for experimentation and debugging.
+
+    Parameters
+    ----------
+    installer
+        Installer that installs a set of packages
+    restart
+        Whether or not to restart the worker after installing the packages
+        Only functions if the worker has an attached nanny process
+    name:
+        Name of the plugin
+
+    See Also
+    --------
+    PackageInstall
+    """
+
+    installer: _PackageInstaller
+    name: str
+    restart: bool
+
+    def __init__(self, installer: _PackageInstaller, restart: bool, name: str):
+        self.installer = installer
         self.restart = restart
-        self.name = f"{self.INSTALLER}-install-{uuid.uuid4()}"
+        self.name = name
 
     async def setup(self, worker):
         from distributed.semaphore import Semaphore
@@ -403,15 +497,15 @@ async def setup(self, worker):
             if not await self._is_installed(worker):
                 logger.info(
                     "%s installing the following packages: %s",
-                    self.INSTALLER,
-                    self.packages,
+                    self.installer.INSTALLER,
+                    self.installer.packages,
                 )
                 await self._set_installed(worker)
-                self.install()
+                self.installer.install()
             else:
                 logger.info(
                     "The following packages have already been installed: %s",
-                    self.packages,
+                    self.installer.packages,
                 )
 
             if self.restart and worker.nanny and not await self._is_restarted(worker):
@@ -421,10 +515,6 @@ async def setup(self, worker):
                     worker.close_gracefully, restart=True, reason=f"{self.name}-setup"
                 )
 
-    @abc.abstractmethod
-    def install(self) -> None:
-        """Install the requested packages"""
-
     async def _is_installed(self, worker):
         return await worker.client.get_metadata(
             self._compose_installed_key(), default=False
@@ -459,6 +549,18 @@ def _compose_restarted_key(self, worker):
         return [self.name, "restarted", worker.nanny]
 
 
+class _PackageInstaller(abc.ABC):
+    INSTALLER: ClassVar[str]
+    packages: list[str]
+
+    def __init__(self, packages: list[str]):
+        self.packages = packages
+
+    @abc.abstractmethod
+    def install(self) -> None:
+        """Install the requested packages"""
+
+
 class CondaInstall(PackageInstall):
     """A Worker Plugin to conda install a set of packages
 
@@ -480,7 +582,7 @@ class CondaInstall(PackageInstall):
         A list of packages (with optional versions) to install using conda
     conda_options
         Additional options to pass to conda
-    restart
+    restart_workers
         Whether or not to restart the worker after installing the packages
         Only functions if the worker has an attached nanny process
 
@@ -497,17 +599,23 @@ class CondaInstall(PackageInstall):
     PipInstall
     """
 
-    INSTALLER = "conda"
-
-    conda_options: list[str]
-
     def __init__(
         self,
         packages: list[str],
         conda_options: list[str] | None = None,
-        restart: bool = False,
+        restart_workers: bool = False,
     ):
-        super().__init__(packages, restart=restart)
+        installer = _CondaInstaller(packages, conda_options)
+        super().__init__(installer, restart_workers=restart_workers)
+
+
+class _CondaInstaller(_PackageInstaller):
+    INSTALLER = "conda"
+
+    conda_options: list[str]
+
+    def __init__(self, packages: list[str], conda_options: list[str] | None = None):
+        super().__init__(packages)
         self.conda_options = conda_options or []
 
     def install(self) -> None:
@@ -556,7 +664,7 @@ class PipInstall(PackageInstall):
         A list of packages (with optional versions) to install using pip
     pip_options
         Additional options to pass to pip
-    restart
+    restart_workers
         Whether or not to restart the worker after installing the packages
         Only functions if the worker has an attached nanny process
 
@@ -573,17 +681,23 @@ class PipInstall(PackageInstall):
     CondaInstall
     """
 
-    INSTALLER = "pip"
-
-    pip_options: list[str]
-
     def __init__(
         self,
         packages: list[str],
         pip_options: list[str] | None = None,
-        restart: bool = False,
+        restart_workers: bool = False,
     ):
-        super().__init__(packages, restart=restart)
+        installer = _PipInstaller(packages, pip_options)
+        super().__init__(installer, restart_workers=restart_workers)
+
+
+class _PipInstaller(_PackageInstaller):
+    INSTALLER = "pip"
+
+    pip_options: list[str]
+
+    def __init__(self, packages: list[str], pip_options: list[str] | None = None):
+        super().__init__(packages)
         self.pip_options = pip_options or []
 
     def install(self) -> None:


Test patch:
diff --git a/distributed/diagnostics/tests/test_packageinstall.py b/distributed/diagnostics/tests/test_packageinstall.py
new file mode 100644
--- /dev/null
+++ b/distributed/diagnostics/tests/test_packageinstall.py
@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from unittest import mock
+
+import pytest
+
+from distributed import Nanny
+from distributed.core import Status
+from distributed.diagnostics.plugin import (
+    CondaInstall,
+    PackageInstall,
+    PipInstall,
+    _PackageInstaller,
+)
+from distributed.utils_test import captured_logger, gen_cluster
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_pip_install(c, s, a):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.INFO
+    ) as logger:
+        mocked = mock.Mock()
+        mocked.configure_mock(
+            **{"communicate.return_value": (b"", b""), "wait.return_value": 0}
+        )
+        with mock.patch(
+            "distributed.diagnostics.plugin.subprocess.Popen", return_value=mocked
+        ) as Popen:
+            await c.register_plugin(
+                PipInstall(packages=["requests"], pip_options=["--upgrade"])
+            )
+            assert Popen.call_count == 1
+            args = Popen.call_args[0][0]
+            assert "python" in args[0]
+            assert args[1:] == ["-m", "pip", "install", "--upgrade", "requests"]
+            logs = logger.getvalue()
+            assert "pip installing" in logs
+            assert "failed" not in logs
+            assert "restart" not in logs
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_conda_install(c, s, a):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.INFO
+    ) as logger:
+        run_command_mock = mock.Mock(name="run_command_mock")
+        run_command_mock.configure_mock(return_value=(b"", b"", 0))
+        module_mock = mock.Mock(name="conda_cli_python_api_mock")
+        module_mock.run_command = run_command_mock
+        module_mock.Commands.INSTALL = "INSTALL"
+        with mock.patch.dict("sys.modules", {"conda.cli.python_api": module_mock}):
+            await c.register_plugin(
+                CondaInstall(packages=["requests"], conda_options=["--update-deps"])
+            )
+            assert run_command_mock.call_count == 1
+            command = run_command_mock.call_args[0][0]
+            assert command == "INSTALL"
+            arguments = run_command_mock.call_args[0][1]
+            assert arguments == ["--update-deps", "requests"]
+            logs = logger.getvalue()
+            assert "conda installing" in logs
+            assert "failed" not in logs
+            assert "restart" not in logs
+
+
+@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
+async def test_pip_install_fails(c, s, a, b):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.ERROR
+    ) as logger:
+        mocked = mock.Mock()
+        mocked.configure_mock(
+            **{
+                "communicate.return_value": (
+                    b"",
+                    b"Could not find a version that satisfies the requirement not-a-package",
+                ),
+                "wait.return_value": 1,
+            }
+        )
+        with mock.patch(
+            "distributed.diagnostics.plugin.subprocess.Popen", return_value=mocked
+        ) as Popen:
+            with pytest.raises(RuntimeError):
+                await c.register_plugin(PipInstall(packages=["not-a-package"]))
+
+            assert Popen.call_count == 1
+            logs = logger.getvalue()
+            assert "install failed" in logs
+            assert "not-a-package" in logs
+
+
+@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
+async def test_conda_install_fails_when_conda_not_found(c, s, a, b):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.ERROR
+    ) as logger:
+        with mock.patch.dict("sys.modules", {"conda": None}):
+            with pytest.raises(RuntimeError):
+                await c.register_plugin(CondaInstall(packages=["not-a-package"]))
+            logs = logger.getvalue()
+            assert "install failed" in logs
+            assert "conda could not be found" in logs
+
+
+@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
+async def test_conda_install_fails_when_conda_raises(c, s, a, b):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.ERROR
+    ) as logger:
+        run_command_mock = mock.Mock(name="run_command_mock")
+        run_command_mock.configure_mock(side_effect=RuntimeError)
+        module_mock = mock.Mock(name="conda_cli_python_api_mock")
+        module_mock.run_command = run_command_mock
+        module_mock.Commands.INSTALL = "INSTALL"
+        with mock.patch.dict("sys.modules", {"conda.cli.python_api": module_mock}):
+            with pytest.raises(RuntimeError):
+                await c.register_plugin(CondaInstall(packages=["not-a-package"]))
+            assert run_command_mock.call_count == 1
+            logs = logger.getvalue()
+            assert "install failed" in logs
+
+
+@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
+async def test_conda_install_fails_on_returncode(c, s, a, b):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.ERROR
+    ) as logger:
+        run_command_mock = mock.Mock(name="run_command_mock")
+        run_command_mock.configure_mock(return_value=(b"", b"", 1))
+        module_mock = mock.Mock(name="conda_cli_python_api_mock")
+        module_mock.run_command = run_command_mock
+        module_mock.Commands.INSTALL = "INSTALL"
+        with mock.patch.dict("sys.modules", {"conda.cli.python_api": module_mock}):
+            with pytest.raises(RuntimeError):
+                await c.register_plugin(CondaInstall(packages=["not-a-package"]))
+            assert run_command_mock.call_count == 1
+            logs = logger.getvalue()
+            assert "install failed" in logs
+
+
+class _StubInstaller(_PackageInstaller):
+    INSTALLER = "stub"
+
+    def __init__(self, packages: list[str]) -> None:
+        super().__init__(packages)
+
+    def install(self) -> None:
+        pass
+
+
+class StubInstall(PackageInstall):
+    def __init__(self, packages: list[str], restart_workers: bool = False):
+        installer = _StubInstaller(packages)
+        super().__init__(installer=installer, restart_workers=restart_workers)
+
+
+@gen_cluster(client=True, nthreads=[("", 1), ("", 1)])
+async def test_package_install_installs_once_with_multiple_workers(c, s, a, b):
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.INFO
+    ) as logger:
+        install_mock = mock.Mock(name="install")
+        with mock.patch.object(_StubInstaller, "install", install_mock):
+            await c.register_plugin(
+                StubInstall(
+                    packages=["requests"],
+                )
+            )
+            assert install_mock.call_count == 1
+            logs = logger.getvalue()
+            assert "already been installed" in logs
+
+
+@gen_cluster(client=True, nthreads=[("", 1), ("", 1)])
+async def test_package_install_installs_once_when_reregistered(c, s, a, b):
+    stub_install = StubInstall(
+        packages=["requests"],
+    )
+    with captured_logger(
+        "distributed.diagnostics.plugin", level=logging.INFO
+    ) as logger:
+        install_mock = mock.Mock(name="install")
+        with mock.patch.object(_StubInstaller, "install", install_mock):
+            await c.register_plugin(stub_install)
+            with pytest.warns(
+                UserWarning,
+                match=r"Scheduler already contains a plugin with name stub-install-.*; overwriting.",
+            ):
+                await c.register_plugin(stub_install)
+            assert install_mock.call_count == 1
+            logs = logger.getvalue()
+            assert "already been installed on the scheduler" in logs
+            # doesn't re-regsiter install on workers
+            assert logs.count("already been installed") == 3
+
+
+@pytest.mark.slow
+@gen_cluster(client=True, nthreads=[("", 1)], Worker=Nanny)
+async def test_package_install_restarts_on_nanny(c, s, a):
+    (addr,) = s.workers
+    await c.register_plugin(
+        StubInstall(
+            packages=["requests"],
+            restart_workers=True,
+        )
+    )
+    # Wait until the worker is restarted
+    while len(s.workers) != 1 or set(s.workers) == {addr}:
+        await asyncio.sleep(0.01)
+
+
+class _FailingInstaller(_PackageInstaller):
+    INSTALLER = "fail"
+
+    def __init__(self, packages: list[str]) -> None:
+        super().__init__(packages)
+
+    def install(self) -> None:
+        raise RuntimeError()
+
+
+class FailingInstall(PackageInstall):
+    def __init__(self, packages: list[str], restart_workers: bool = False):
+        installer = _FailingInstaller(packages=packages)
+        super().__init__(installer=installer, restart_workers=restart_workers)
+
+
+@gen_cluster(client=True, nthreads=[("", 1)], Worker=Nanny)
+async def test_package_install_failing_does_not_restart_on_nanny(c, s, a):
+    (addr,) = s.workers
+    with pytest.raises(RuntimeError):
+        await c.register_plugin(
+            FailingInstall(
+                packages=["requests"],
+                restart_workers=True,
+            )
+        )
+    # Nanny does not restart
+    assert a.status is Status.running
+    assert set(s.workers) == {addr}
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
@@ -47,12 +47,7 @@
 from distributed.compatibility import LINUX, WINDOWS
 from distributed.core import CommClosedError, Status, rpc
 from distributed.diagnostics import nvml
-from distributed.diagnostics.plugin import (
-    CondaInstall,
-    ForwardOutput,
-    PackageInstall,
-    PipInstall,
-)
+from distributed.diagnostics.plugin import ForwardOutput
 from distributed.metrics import time
 from distributed.protocol import pickle
 from distributed.scheduler import KilledWorker, Scheduler
@@ -1687,199 +1682,6 @@ def bad_startup(w):
         pass
 
 
-@gen_cluster(client=True, nthreads=[("", 1)])
-async def test_pip_install(c, s, a):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.INFO
-    ) as logger:
-        mocked = mock.Mock()
-        mocked.configure_mock(
-            **{"communicate.return_value": (b"", b""), "wait.return_value": 0}
-        )
-        with mock.patch(
-            "distributed.diagnostics.plugin.subprocess.Popen", return_value=mocked
-        ) as Popen:
-            await c.register_plugin(
-                PipInstall(packages=["requests"], pip_options=["--upgrade"])
-            )
-            assert Popen.call_count == 1
-            args = Popen.call_args[0][0]
-            assert "python" in args[0]
-            assert args[1:] == ["-m", "pip", "install", "--upgrade", "requests"]
-            logs = logger.getvalue()
-            assert "pip installing" in logs
-            assert "failed" not in logs
-            assert "restart" not in logs
-
-
-@gen_cluster(client=True, nthreads=[("", 1)])
-async def test_conda_install(c, s, a):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.INFO
-    ) as logger:
-        run_command_mock = mock.Mock(name="run_command_mock")
-        run_command_mock.configure_mock(return_value=(b"", b"", 0))
-        module_mock = mock.Mock(name="conda_cli_python_api_mock")
-        module_mock.run_command = run_command_mock
-        module_mock.Commands.INSTALL = "INSTALL"
-        with mock.patch.dict("sys.modules", {"conda.cli.python_api": module_mock}):
-            await c.register_plugin(
-                CondaInstall(packages=["requests"], conda_options=["--update-deps"])
-            )
-            assert run_command_mock.call_count == 1
-            command = run_command_mock.call_args[0][0]
-            assert command == "INSTALL"
-            arguments = run_command_mock.call_args[0][1]
-            assert arguments == ["--update-deps", "requests"]
-            logs = logger.getvalue()
-            assert "conda installing" in logs
-            assert "failed" not in logs
-            assert "restart" not in logs
-
-
-@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
-async def test_pip_install_fails(c, s, a, b):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.ERROR
-    ) as logger:
-        mocked = mock.Mock()
-        mocked.configure_mock(
-            **{
-                "communicate.return_value": (
-                    b"",
-                    b"Could not find a version that satisfies the requirement not-a-package",
-                ),
-                "wait.return_value": 1,
-            }
-        )
-        with mock.patch(
-            "distributed.diagnostics.plugin.subprocess.Popen", return_value=mocked
-        ) as Popen:
-            with pytest.raises(RuntimeError):
-                await c.register_plugin(PipInstall(packages=["not-a-package"]))
-
-            assert Popen.call_count == 1
-            logs = logger.getvalue()
-            assert "install failed" in logs
-            assert "not-a-package" in logs
-
-
-@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
-async def test_conda_install_fails_when_conda_not_found(c, s, a, b):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.ERROR
-    ) as logger:
-        with mock.patch.dict("sys.modules", {"conda": None}):
-            with pytest.raises(RuntimeError):
-                await c.register_plugin(CondaInstall(packages=["not-a-package"]))
-            logs = logger.getvalue()
-            assert "install failed" in logs
-            assert "conda could not be found" in logs
-
-
-@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
-async def test_conda_install_fails_when_conda_raises(c, s, a, b):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.ERROR
-    ) as logger:
-        run_command_mock = mock.Mock(name="run_command_mock")
-        run_command_mock.configure_mock(side_effect=RuntimeError)
-        module_mock = mock.Mock(name="conda_cli_python_api_mock")
-        module_mock.run_command = run_command_mock
-        module_mock.Commands.INSTALL = "INSTALL"
-        with mock.patch.dict("sys.modules", {"conda.cli.python_api": module_mock}):
-            with pytest.raises(RuntimeError):
-                await c.register_plugin(CondaInstall(packages=["not-a-package"]))
-            assert run_command_mock.call_count == 1
-            logs = logger.getvalue()
-            assert "install failed" in logs
-
-
-@gen_cluster(client=True, nthreads=[("", 2), ("", 2)])
-async def test_conda_install_fails_on_returncode(c, s, a, b):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.ERROR
-    ) as logger:
-        run_command_mock = mock.Mock(name="run_command_mock")
-        run_command_mock.configure_mock(return_value=(b"", b"", 1))
-        module_mock = mock.Mock(name="conda_cli_python_api_mock")
-        module_mock.run_command = run_command_mock
-        module_mock.Commands.INSTALL = "INSTALL"
-        with mock.patch.dict("sys.modules", {"conda.cli.python_api": module_mock}):
-            with pytest.raises(RuntimeError):
-                await c.register_plugin(CondaInstall(packages=["not-a-package"]))
-            assert run_command_mock.call_count == 1
-            logs = logger.getvalue()
-            assert "install failed" in logs
-
-
-class StubInstall(PackageInstall):
-    INSTALLER = "stub"
-
-    def __init__(self, packages: list[str], restart: bool = False):
-        super().__init__(packages=packages, restart=restart)
-
-    def install(self) -> None:
-        pass
-
-
-@gen_cluster(client=True, nthreads=[("", 1), ("", 1)])
-async def test_package_install_installs_once_with_multiple_workers(c, s, a, b):
-    with captured_logger(
-        "distributed.diagnostics.plugin", level=logging.INFO
-    ) as logger:
-        install_mock = mock.Mock(name="install")
-        with mock.patch.object(StubInstall, "install", install_mock):
-            await c.register_plugin(
-                StubInstall(
-                    packages=["requests"],
-                )
-            )
-            assert install_mock.call_count == 1
-            logs = logger.getvalue()
-            assert "already been installed" in logs
-
-
-@pytest.mark.slow
-@gen_cluster(client=True, nthreads=[("", 1)], Worker=Nanny)
-async def test_package_install_restarts_on_nanny(c, s, a):
-    (addr,) = s.workers
-    await c.register_plugin(
-        StubInstall(
-            packages=["requests"],
-            restart=True,
-        )
-    )
-    # Wait until the worker is restarted
-    while len(s.workers) != 1 or set(s.workers) == {addr}:
-        await asyncio.sleep(0.01)
-
-
-class FailingInstall(PackageInstall):
-    INSTALLER = "fail"
-
-    def __init__(self, packages: list[str], restart: bool = False):
-        super().__init__(packages=packages, restart=restart)
-
-    def install(self) -> None:
-        raise RuntimeError()
-
-
-@gen_cluster(client=True, nthreads=[("", 1)], Worker=Nanny)
-async def test_package_install_failing_does_not_restart_on_nanny(c, s, a):
-    (addr,) = s.workers
-    with pytest.raises(RuntimeError):
-        await c.register_plugin(
-            FailingInstall(
-                packages=["requests"],
-                restart=True,
-            )
-        )
-    # Nanny does not restart
-    assert a.status is Status.running
-    assert set(s.workers) == {addr}
-
-
 @gen_cluster(nthreads=[])
 async def test_update_latency(s):
     async with Worker(s.address) as w: