Ensure memray profiler runs on all workers (#9095)

Author: jrbourbeau

distributed/diagnostics/memray.py

@@ -131,7 +131,7 @@ def memray_workers(
     """
     directory = pathlib.Path(directory)
     client = get_client()
-    scheduler_info = client.scheduler_info()
+    scheduler_info = client.scheduler_info(n_workers=-1)
     worker_addr = scheduler_info["workers"]
     worker_names = {
         addr: winfo["name"] for addr, winfo in scheduler_info["workers"].items()

distributed/diagnostics/tests/test_memray.py

@@ -6,9 +6,20 @@
 
 pytestmark = pytest.mark.extra_packages
 
+from distributed import Client, LocalCluster
 from distributed.diagnostics.memray import memray_scheduler, memray_workers
 
 
+def test_all_workers(tmp_path, loop):
+    n_workers = 7
+    with LocalCluster(n_workers=n_workers, loop=loop) as cluster:
+        with Client(cluster, loop=loop) as client:
+            with memray_workers(tmp_path):
+                pass
+
+            assert len(list(tmp_path.glob("*.html"))) == n_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):