Repository: dask/distributed

Issue:
PipInstall plugin restart fails behind private pypi and when --upgrade pip option is specified
If you are installing from a private pypi server and you want the pip install plugin to restart the worker, the following lines of code cause the worker to restart constantly since the first line of text is "Looking in indexes..." not "Requirement already ..."

https://github.com/dask/distributed/blob/1fd07f03cacee6fde81d13282568a727bce789b9/distributed/diagnostics/plugin.py#L306-L312

Similarly, if you happen to pass --upgrade in the pip options list, you get similar behavior where the worker continually restarts, because the message from pip is "Requirement already up-to-date" instead of "Requirement already satisfied"


Discussion:
Hi @mdwgrogan, thanks for filing this issue. I have started a PR that enables installing from a private PyPI server in #7086. 

Regarding the problem with `--upgrade`, what version of `pip` are you using in your environment? I could not reproduce this issue without enabling the legacy resolver using `--use-deprecated=legacy-resolver` on `pip 22.2.2`.

Resolving patch:
diff --git a/distributed/diagnostics/plugin.py b/distributed/diagnostics/plugin.py
--- a/distributed/diagnostics/plugin.py
+++ b/distributed/diagnostics/plugin.py
@@ -254,10 +254,6 @@ class PipInstall(WorkerPlugin):
        libraries in the worker environment or image. This is
        primarily intended for experimentation and debugging.
 
-       Additional issues may arise if multiple workers share the same
-       file system. Each worker might try to install the packages
-       simultaneously.
-
     Parameters
     ----------
     packages : List[str]
@@ -282,34 +278,74 @@ def __init__(self, packages, pip_options=None, restart=False):
         self.packages = packages
         self.restart = restart
         self.pip_options = pip_options or []
+        self.id = f"pip-install-{uuid.uuid4()}"
 
     async def setup(self, worker):
-        from distributed.lock import Lock
-
-        async with Lock(socket.gethostname()):  # don't clobber one installation
-            logger.info("Pip installing the following packages: %s", self.packages)
-            proc = subprocess.Popen(
-                [sys.executable, "-m", "pip", "install"]
-                + self.pip_options
-                + self.packages,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-            )
-            stdout, stderr = proc.communicate()
-            returncode = proc.wait()
-
-            if returncode:
-                logger.error("Pip install failed with '%s'", stderr.decode().strip())
-                return
-
-            if self.restart and worker.nanny:
-                lines = stdout.strip().split(b"\n")
-                if not all(
-                    line.startswith(b"Requirement already satisfied") for line in lines
-                ):
-                    worker.loop.add_callback(
-                        worker.close_gracefully, restart=True
-                    )  # restart
+        from distributed.semaphore import Semaphore
+
+        async with (
+            await Semaphore(max_leases=1, name=socket.gethostname(), register=True)
+        ):
+            if not await self._is_installed(worker):
+                logger.info("Pip installing the following packages: %s", self.packages)
+                await self._set_installed(worker)
+                self._install()
+            else:
+                logger.info(
+                    "The following packages have already been installed: %s",
+                    self.packages,
+                )
+
+            if self.restart and worker.nanny and not await self._is_restarted(worker):
+                logger.info("Restarting worker to refresh interpreter.")
+                await self._set_restarted(worker)
+                worker.loop.add_callback(worker.close_gracefully, restart=True)
+
+    def _install(self):
+        proc = subprocess.Popen(
+            [sys.executable, "-m", "pip", "install"] + self.pip_options + self.packages,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        _, stderr = proc.communicate()
+        returncode = proc.wait()
+        if returncode != 0:
+            msg = f"Pip install failed with '{stderr.decode().strip()}'"
+            logger.error(msg)
+            raise RuntimeError(msg)
+
+    async def _is_installed(self, worker):
+        return await worker.client.get_metadata(
+            self._compose_installed_key(), default=False
+        )
+
+    async def _set_installed(self, worker):
+        await worker.client.set_metadata(
+            self._compose_installed_key(),
+            True,
+        )
+
+    def _compose_installed_key(self):
+        return [
+            self.id,
+            "installed",
+            socket.gethostname(),
+        ]
+
+    async def _is_restarted(self, worker):
+        return await worker.client.get_metadata(
+            self._compose_restarted_key(worker),
+            default=False,
+        )
+
+    async def _set_restarted(self, worker):
+        await worker.client.set_metadata(
+            self._compose_restarted_key(worker),
+            True,
+        )
+
+    def _compose_restarted_key(self, worker):
+        return [self.id, "restarted", worker.nanny]
 
 
 # Adapted from https://github.com/dask/distributed/issues/3560#issuecomment-596138522


Test patch:
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
@@ -15,6 +15,7 @@
 from concurrent.futures.process import BrokenProcessPool
 from numbers import Number
 from operator import add
+from textwrap import dedent
 from time import sleep
 from unittest import mock
 
@@ -1629,51 +1630,128 @@ def bad_startup(w):
         pass
 
 
-@gen_cluster(client=True)
-async def test_pip_install(c, s, a, b):
-    with mock.patch(
-        "distributed.diagnostics.plugin.subprocess.Popen.communicate",
-        return_value=(b"", b""),
-    ) as p1:
+@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=p1
-        ) as p2:
-            p1.communicate.return_value = b"", b""
-            p1.wait.return_value = 0
+            "distributed.diagnostics.plugin.subprocess.Popen", return_value=mocked
+        ) as Popen:
             await c.register_worker_plugin(
                 PipInstall(packages=["requests"], pip_options=["--upgrade"])
             )
 
-            args = p2.call_args[0][0]
+            args = Popen.call_args[0][0]
             assert "python" in args[0]
             assert args[1:] == ["-m", "pip", "install", "--upgrade", "requests"]
+            assert Popen.call_count == 1
+            logs = logger.getvalue()
+            assert "Pip installing" in logs
+            assert "failed" not in logs
+            assert "restart" not in logs
 
 
-@gen_cluster(client=True)
+@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:
-        with mock.patch(
-            "distributed.diagnostics.plugin.subprocess.Popen.communicate",
-            return_value=(b"", b"error"),
-        ) as p1:
-            with mock.patch(
-                "distributed.diagnostics.plugin.subprocess.Popen", return_value=p1
-            ) as p2:
-                p1.communicate.return_value = (
+        mocked = mock.Mock()
+        mocked.configure_mock(
+            **{
+                "communicate.return_value": (
                     b"",
                     b"Could not find a version that satisfies the requirement not-a-package",
-                )
-                p1.wait.return_value = 1
+                ),
+                "wait.return_value": 1,
+            }
+        )
+        with mock.patch(
+            "distributed.diagnostics.plugin.subprocess.Popen", return_value=mocked
+        ) as Popen:
+            with pytest.raises(RuntimeError):
                 await c.register_worker_plugin(PipInstall(packages=["not-a-package"]))
 
-                assert "not-a-package" in logger.getvalue()
+            assert Popen.call_count == 1
+            logs = logger.getvalue()
+            assert "install failed" in logs
+            assert "not-a-package" in logs
+
+
+@gen_cluster(client=True, nthreads=[])
+async def test_pip_install_restarts_on_nanny(c, s):
+    preload = dedent(
+        """\
+        from unittest import mock
+
+        mock.patch(
+            "distributed.diagnostics.plugin.PipInstall._install", return_value=None
+        ).start()
+        """
+    )
+    async with Nanny(s.address, preload=preload):
+        (addr,) = s.workers
+        await c.register_worker_plugin(
+            PipInstall(packages=["requests"], pip_options=["--upgrade"], restart=True)
+        )
+
+        # Wait until the worker is restarted
+        while len(s.workers) != 1 or set(s.workers) == {addr}:
+            await asyncio.sleep(0.01)
 
 
-#             args = p2.call_args[0][0]
-#             assert "python" in args[0]
-#             assert args[1:] == ["-m", "pip", "--upgrade", "install", "requests"]
+@gen_cluster(client=True, nthreads=[])
+async def test_pip_install_failing_does_not_restart_on_nanny(c, s):
+    preload = dedent(
+        """\
+        from unittest import mock
+
+        mock.patch(
+            "distributed.diagnostics.plugin.PipInstall._install", side_effect=RuntimeError
+        ).start()
+        """
+    )
+    async with Nanny(s.address, preload=preload) as n:
+        (addr,) = s.workers
+        with pytest.raises(RuntimeError):
+            await c.register_worker_plugin(
+                PipInstall(
+                    packages=["requests"], pip_options=["--upgrade"], restart=True
+                )
+            )
+        # Nanny does not restart
+        assert n.status is Status.running
+        assert set(s.workers) == {addr}
+
+
+@gen_cluster(client=True, nthreads=[("", 1), ("", 1)])
+async def test_pip_install_multiple_workers(c, s, a, b):
+    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_worker_plugin(
+                PipInstall(packages=["requests"], pip_options=["--upgrade"])
+            )
+
+            args = Popen.call_args[0][0]
+            assert "python" in args[0]
+            assert args[1:] == ["-m", "pip", "install", "--upgrade", "requests"]
+            assert Popen.call_count == 1
+            logs = logger.getvalue()
+            assert "Pip installing" in logs
+            assert "already been installed" in logs
 
 
 @gen_cluster(nthreads=[])