Repository: dask/distributed

Issue:
gather() should not remove unresponsive workers
`Scheduler.gather` removes workers it fails to connect to.
However, if it fails to connect to a worker, but then successfully retrieves the same key from another worker, this will not happen.

This is a bad idea to begin with, and the last remnant of an old design where various failed RPC calls would shut down workers.
One key failure point is that temporary GIL congestion or a minor glitch in the network can cause new RPC calls to randomly fail.
The current philosophy is that only a collapse of the Batched Comms - which, once established, is kept alive by the kernel and is immune to the GIL - should cause worker and scheduler to get disconnected.

Discussion:

Patch:
diff --git a/continuous_integration/environment-3.9.yaml b/continuous_integration/environment-3.9.yaml
index ad4d3408..d4359981 100644
--- a/continuous_integration/environment-3.9.yaml
+++ b/continuous_integration/environment-3.9.yaml
@@ -22,6 +22,7 @@ dependencies:
   - jinja2 >=2.10.3
   - locket >=1.0
   - lz4 >=0.23.1  # Only tested here
+  - memray # Only tested here
   - msgpack-python
   - netcdf4
   - paramiko
diff --git a/distributed/client.py b/distributed/client.py
index 4e424a73..5dfc4365 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2210,8 +2210,8 @@ class Client(SyncMethodMixin):
         if mismatched_futures:
             raise ValueError(
                 "Cannot gather Futures created by another client. "
-                f"These are the {len(mismatched_futures)} (out of {len(futures)}) mismatched Futures and their client IDs "
-                f"(this client is {self.id}): "
+                f"These are the {len(mismatched_futures)} (out of {len(futures)}) "
+                f"mismatched Futures and their client IDs (this client is {self.id}): "
                 f"{ {f: f.client.id for f in mismatched_futures} }"
             )
         keys = [stringify(future.key) for future in future_set]
diff --git a/distributed/diagnostics/memray.py b/distributed/diagnostics/memray.py
new file mode 100644
index 00000000..fc96a34d
--- /dev/null
+++ b/distributed/diagnostics/memray.py
@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import contextlib
+import logging
+import pathlib
+import subprocess
+import time
+import uuid
+from collections.abc import Iterator, Sequence
+from typing import Any, Literal
+from urllib.parse import quote
+
+from toolz.itertoolz import partition
+
+from distributed import get_client
+from distributed.worker import Worker
+
+try:
+    import memray
+except ImportError:
+    raise ImportError("You have to install memray to use this module.")
+
+logger = logging.getLogger(__name__)
+
+
+def _start_memray(dask_worker: Worker, filename: str, **kwargs: Any) -> bool:
+    """Start the memray Tracker on a Server"""
+    if hasattr(dask_worker, "_memray"):
+        dask_worker._memray.close()
+
+    path = pathlib.Path(dask_worker.local_directory) / (filename + str(dask_worker.id))
+    if path.exists():
+        path.rmdir()
+
+    dask_worker._memray = contextlib.ExitStack()  # type: ignore[attr-defined]
+    dask_worker._memray.enter_context(  # type: ignore[attr-defined]
+        memray.Tracker(path, native_traces=True, **kwargs)
+    )
+
+    return True
+
+
+def _fetch_memray_profile(
+    dask_worker: Worker, filename: str, report_args: Sequence[str] | Literal[False]
+) -> bytes:
+    """Generate and fetch the memray report"""
+    if not hasattr(dask_worker, "_memray"):
+        return b""
+    path = pathlib.Path(dask_worker.local_directory) / (filename + str(dask_worker.id))
+    dask_worker._memray.close()
+    del dask_worker._memray
+
+    if not report_args:
+        with open(path, "rb") as fd:
+            return fd.read()
+
+    report_filename = path.with_suffix(".html")
+    if not report_args[0] == "memray":
+        report_args = ["memray"] + list(report_args)
+    assert "-f" not in report_args, "Cannot provide filename for report generation"
+    assert (
+        "-o" not in report_args
+    ), "Cannot provide output filename for report generation"
+    report_args = list(report_args) + ["-f", str(path), "-o", str(report_filename)]
+    subprocess.run(report_args)
+    with open(report_filename, "rb") as fd:
+        return fd.read()
+
+
+@contextlib.contextmanager
+def memray_workers(
+    directory: str | pathlib.Path = "memray-profiles",
+    workers: int | None | list[str] = None,
+    report_args: Sequence[str]
+    | Literal[False] = ("flamegraph", "--temporal", "--leaks"),
+    fetch_reports_parallel: bool | int = True,
+    **memray_kwargs: Any,
+) -> Iterator[None]:
+    """Generate a Memray profile on the workers and download the generated report.
+
+    Example::
+
+        with memray_workers():
+            client.submit(my_function).result()
+
+        # Or even while the computation is already running
+
+        fut = client.submit(my_function)
+
+        with memray_workers():
+            time.sleep(10)
+
+        fut.result()
+
+    Parameters
+    ----------
+    directory : str
+        The directory to save the reports to.
+    workers : int | None | list[str]
+        The workers to profile. If int, the first n workers will be used.
+        If None, all workers will be used.
+        If list[str], the workers with the given addresses will be used.
+    report_args : tuple[str]
+        Particularly for native_traces=True, the reports have to be
+        generated on the same host using the same Python interpreter as the
+        profile was generated. Otherwise, native traces will yield unusable
+        results. Therefore, we're generating the reports on the workers and
+        download them afterwards. You can modify the report generation by
+        providing additional arguments and we will generate the reports as::
+
+            memray *report_args -f <filename> -o <filename>.html
+
+        If the raw data should be fetched instead of the report, set this to
+        False.
+
+    fetch_reports_parallel : bool | int
+        Fetching results is sometimes slow and it's sometimes not desired to
+        wait for all workers to finish before receiving the first reports.
+        This controls how many workers are fetched concurrently.
+
+            int: Number of workers to fetch concurrently
+            True: All workers concurrently
+            False: One worker at a time
+
+    **memray_kwargs
+        Keyword arguments to be passed to memray.Tracker, e.g.
+        {"native_traces": True}
+    """
+    directory = pathlib.Path(directory)
+    client = get_client()
+    scheduler_info = client.scheduler_info()
+    worker_addr = scheduler_info["workers"]
+    worker_names = {
+        addr: winfo["name"] for addr, winfo in scheduler_info["workers"].items()
+    }
+    if not workers or isinstance(workers, int):
+        nworkers = len(worker_addr)
+        if isinstance(workers, int):
+            nworkers = workers
+        workers = list(worker_addr)[:nworkers]
+    workers = list(workers)
+    filename = uuid.uuid4().hex
+    assert all(client.run(_start_memray, filename=filename, **memray_kwargs).values())
+    # Sleep for a brief moment such that we get
+    # a clear profiling signal when everything starts
+    time.sleep(0.1)
+    yield
+    directory.mkdir(exist_ok=True)
+
+    client = get_client()
+    if fetch_reports_parallel is True:
+        fetch_parallel = len(workers)
+    elif fetch_reports_parallel is False:
+        fetch_parallel = 1
+    else:
+        fetch_parallel = fetch_reports_parallel
+
+    for w in partition(fetch_parallel, workers):
+        try:
+            profiles = client.run(
+                _fetch_memray_profile,
+                filename=filename,
+                report_args=report_args,
+                workers=w,
+            )
+            for worker_addr, profile in profiles.items():
+                path = directory / quote(worker_names[worker_addr], safe="")
+                if report_args:
+                    suffix = ".html"
+                else:
+                    suffix = ".memray"
+                with open(str(path) + suffix, "wb") as fd:
+                    fd.write(profile)
+
+        except Exception:
+            logger.exception("Exception during report downloading from worker %s", w)
+
+
+@contextlib.contextmanager
+def memray_scheduler(
+    directory: str | pathlib.Path = "memray-profiles",
+    report_args: Sequence[str]
+    | Literal[False] = ("flamegraph", "--temporal", "--leaks"),
+    **memray_kwargs: Any,
+) -> Iterator[None]:
+    """Generate a Memray profile on the Scheduler and download the generated report.
+
+    Example::
+
+        with memray_scheduler():
+            client.submit(my_function).result()
+
+        # Or even while the computation is already running
+
+        fut = client.submit(my_function)
+
+        with memray_scheduler():
+            time.sleep(10)
+
+        fut.result()
+
+    Parameters
+    ----------
+    directory : str
+        The directory to save the reports to.
+    report_args : tuple[str]
+        Particularly for native_traces=True, the reports have to be
+        generated on the same host using the same Python interpreter as the
+        profile was generated. Otherwise, native traces will yield unusable
+        results. Therefore, we're generating the reports on the Scheduler and
+        download them afterwards. You can modify the report generation by
+        providing additional arguments and we will generate the reports as::
+
+            memray *report_args -f <filename> -o <filename>.html
+
+        If the raw data should be fetched instead of the report, set this to
+        False.
+    **memray_kwargs
+        Keyword arguments to be passed to memray.Tracker, e.g.
+        {"native_traces": True}
+    """
+    directory = pathlib.Path(directory)
+    client = get_client()
+    filename = uuid.uuid4().hex
+    assert client.run_on_scheduler(_start_memray, filename=filename, **memray_kwargs)
+    # Sleep for a brief moment such that we get
+    # a clear profiling signal when everything starts
+    time.sleep(0.1)
+    yield
+    directory.mkdir(exist_ok=True)
+
+    client = get_client()
+
+    profile = client.run_on_scheduler(
+        _fetch_memray_profile,
+        filename=filename,
+        report_args=report_args,
+    )
+    path = directory / "scheduler"
+    if report_args:
+        suffix = ".html"
+    else:
+        suffix = ".memray"
+    with open(str(path) + suffix, "wb") as fd:
+        fd.write(profile)
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 3d59b159..e83e97c7 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -5887,7 +5887,6 @@ class Scheduler(SchedulerState, ServerNode):
         self, keys: Collection[str], serializers: list[str] | None = None
     ) -> dict[str, Any]:
         """Collect data from workers to the scheduler"""
-        stimulus_id = f"gather-{time()}"
         data = {}
         missing_keys = list(keys)
         failed_keys: list[str] = []
@@ -5924,20 +5923,6 @@ class Scheduler(SchedulerState, ServerNode):
             for key in failed_keys
         }
         logger.error("Couldn't gather keys: %s", failed_states)
-
-        if missing_workers:
-            with log_errors():
-                # Remove suspicious workers from the scheduler and shut them down.
-                await asyncio.gather(
-                    *(
-                        self.remove_worker(
-                            address=worker, close=True, stimulus_id=stimulus_id
-                        )
-                        for worker in missing_workers
-                    )
-                )
-                logger.error("Shut down unresponsive workers:: %s", missing_workers)
-
         return {"status": "error", "keys": list(failed_keys)}
 
     @log_errors
diff --git a/distributed/utils_comm.py b/distributed/utils_comm.py
index 1ba22e3d..04266a97 100644
--- a/distributed/utils_comm.py
+++ b/distributed/utils_comm.py
@@ -76,8 +76,8 @@ async def gather_from_workers(
 
             return data, list(to_gather), failed_keys, list(missing_workers)
 
-        tasks = {
-            address: asyncio.create_task(
+        tasks = [
+            asyncio.create_task(
                 retry_operation(
                     partial(
                         get_data_from_worker,
@@ -92,23 +92,28 @@ async def gather_from_workers(
                 name=f"get-data-from-{address}",
             )
             for address, keys in d.items()
-        }
-        for address, task in tasks.items():
-            try:
-                r = await task
-            except OSError:
+        ]
+
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        for address, r in zip(d, results):
+            if isinstance(r, OSError):
                 missing_workers.add(address)
-            except Exception:
+            elif isinstance(r, Exception):
                 # For example, deserialization error
-                logger.exception(
+                logger.error(
                     "Unexpected error while collecting tasks %s from %s",
                     d[address],
                     address,
+                    exc_info=r,
                 )
                 for key in d[address]:
                     failed_keys.append(key)
                     del to_gather[key]
+            elif isinstance(r, BaseException):  # pragma: nocover
+                # for example, asyncio.CancelledError
+                raise r
             else:
+                assert isinstance(r, dict), r
                 if r["status"] == "busy":
                     busy_workers.add(address)
                     continue


Test patch:
diff --git a/distributed/diagnostics/tests/test_memray.py b/distributed/diagnostics/tests/test_memray.py
new file mode 100644
index 00000000..3ebb595d
--- /dev/null
+++ b/distributed/diagnostics/tests/test_memray.py
@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import pytest
+
+memray = pytest.importorskip("memray")
+
+
+from distributed.diagnostics.memray import memray_scheduler, memray_workers
+
+
+@pytest.mark.parametrize("fetch_reports_parallel", [True, False, 1])
+def test_basic_integration_workers(client, tmp_path, fetch_reports_parallel):
+    with memray_workers(tmp_path, fetch_reports_parallel=fetch_reports_parallel):
+        pass
+
+    assert len(list(tmp_path.glob("*.html"))) == 2
+
+
+@pytest.mark.parametrize("report_args", [("flamegraph", "--leaks"), False])
+def test_basic_integration_workers_report_args(client, tmp_path, report_args):
+    with memray_workers(tmp_path, report_args=report_args):
+        pass
+
+    if report_args:
+        assert len(list(tmp_path.glob("*.html"))) == 2
+    else:
+        assert len(list(tmp_path.glob("*.html"))) == 0
+        assert len(list(tmp_path.glob("*.memray"))) == 2
+
+
+def test_basic_integration_scheduler(client, tmp_path):
+    with memray_scheduler(tmp_path):
+        pass
+
+    assert len(list(tmp_path.glob("*.html"))) == 1
+
+
+@pytest.mark.parametrize("report_args", [("flamegraph", "--leaks"), False])
+def test_basic_integration_scheduler_report_args(client, tmp_path, report_args):
+    with memray_scheduler(tmp_path, report_args=report_args):
+        pass
+
+    if report_args:
+        assert len(list(tmp_path.glob("*.html"))) == 1
+    else:
+        assert len(list(tmp_path.glob("*.html"))) == 0
+        assert len(list(tmp_path.glob("*.memray"))) == 1
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index 890b4d8d..95d0ee98 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -74,7 +74,7 @@ from distributed.utils_test import (
     varying,
     wait_for_state,
 )
-from distributed.worker import dumps_function, get_worker, secede
+from distributed.worker import dumps_function, secede
 
 pytestmark = pytest.mark.ci1
 
@@ -2860,56 +2860,41 @@ async def test_gather_no_workers(c, s, a, b):
     assert list(res["keys"]) == ["x"]
 
 
-@gen_cluster(client=True, client_kwargs={"direct_to_workers": False})
-async def test_gather_bad_worker_removed(c, s, a, b):
-    """
-    Upon connection failure or missing expected keys during gather, a worker is
-    shut down. The tasks should be rescheduled onto different workers, transparently
-    to `client.gather`.
+@pytest.mark.parametrize("direct", [False, True])
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)],
+    # This behaviour is independent of retries.
+    # Disable them to reduce the complexity of this test.
+    config={"distributed.comm.retry.count": 0},
+)
+async def test_gather_bad_worker(c, s, a, direct):
+    """Upon connection failure, gather() tries again indefinitely and transparently,
+    for as long as the batched comms channel is active.
     """
-    x = c.submit(slowinc, 1, workers=[a.address], allow_other_workers=True)
-
-    def finalizer(*args):
-        return get_worker().address
+    x = c.submit(inc, 1, key="x")
+    c.rpc = await FlakyConnectionPool(failing_connections=3)
+    s.rpc = await FlakyConnectionPool(failing_connections=1)
 
-    fin = c.submit(
-        finalizer, x, key="final", workers=[a.address], allow_other_workers=True
-    )
+    with captured_logger("distributed.scheduler") as sched_logger:
+        with captured_logger("distributed.client") as client_logger:
+            assert await c.gather(x, direct=direct) == 2
 
-    s.rpc = await FlakyConnectionPool(failing_connections=1)
+    assert "Couldn't gather keys: {'x': 'memory'}" in sched_logger.getvalue()
+    assert "Couldn't gather 1 keys, rescheduling ('x',)" in client_logger.getvalue()
 
-    # This behaviour is independent of retries. Remove them to reduce complexity
-    # of this setup
-    with dask.config.set({"distributed.comm.retry.count": 0}):
-        with captured_logger(
-            logging.getLogger("distributed.scheduler")
-        ) as sched_logger, captured_logger(
-            logging.getLogger("distributed.client")
-        ) as client_logger:
-            # Gather using the client (as an ordinary user would)
-            # Upon a missing key, the client will remove the bad worker and
-            # reschedule the computations
-
-            # Both tasks are rescheduled onto `b`, since `a` was removed.
-            assert await fin == b.address
-
-            await a.finished()
-            assert list(s.workers) == [b.address]
-
-            sched_logger = sched_logger.getvalue()
-            client_logger = client_logger.getvalue()
-            assert "Shut down unresponsive workers" in sched_logger
-
-            assert "Couldn't gather 1 keys, rescheduling" in client_logger
-
-            assert s.tasks[fin.key].who_has == {s.workers[b.address]}
-            assert a.state.executed_count == 2
-            assert b.state.executed_count >= 1
-            # ^ leave room for a future switch from `remove_worker` to `retire_workers`
-
-    # Ensure that the communication was done via the scheduler, i.e. we actually hit a
-    # bad connection
-    assert s.rpc.cnn_count > 0
+    if direct:
+        # 1. try direct=True; fail
+        # 2. fall back to direct=False; fail
+        # 3. try direct=True again; fail
+        # 4. fall back to direct=False again; success
+        assert c.rpc.cnn_count == 2
+        assert s.rpc.cnn_count == 2
+    else:
+        # 1. try direct=False; fail
+        # 2. try again direct=False; success
+        assert c.rpc.cnn_count == 0
+        assert s.rpc.cnn_count == 2
 
 
 @gen_cluster(client=True)