Patch wait_for_workers (#1681)

Author: fjetter

.github/workflows/tests.yml

@@ -4,7 +4,6 @@ on:
     # Runs at 00:00 on Sunday, Tuesday, and Thursday. (see https://crontab.guru)
     - cron: "0 0 * * 0,2,4"
   workflow_dispatch:
-
 concurrency:
   # Include `github.event_name` to avoid pushes to `main` and
   # scheduled jobs canceling one another

detect_regressions.py

@@ -87,7 +87,7 @@ def detect_regressions(database_file, is_pr=False):
                                 df_test[metric].iloc[-3] * units_norm,
                             )
                             reg = (
-                                f"{runtime = }, {name = }, {category = }, "
+                                f"{runtime=}, {name=}, {category=}, "
                                 f"last_three_{metric} {u} = "
                                 f"{last_three}, "
                                 f"{metric}_threshold {u} = {metric_threshold * units_norm} \n"

tests/benchmarks/test_parquet.py

@@ -40,11 +40,11 @@ def parquet_cluster(dask_env_variables, cluster_kwargs, github_cluster_tags):
 
 
 @pytest.fixture
-def parquet_client(parquet_cluster, cluster_kwargs, benchmark_all):
+def parquet_client(parquet_cluster, cluster_kwargs, benchmark_all, wait_for_workers):
     n_workers = cluster_kwargs["parquet_cluster"]["n_workers"]
     with distributed.Client(parquet_cluster) as client:
         parquet_cluster.scale(n_workers)
-        client.wait_for_workers(n_workers, timeout=600)
+        wait_for_workers(client, n_workers, timeout=600)
         client.restart()
         with benchmark_all(client):
             yield client

tests/benchmarks/test_spill.py

@@ -35,11 +35,11 @@ def spill_cluster(dask_env_variables, cluster_kwargs, github_cluster_tags):
 
 
 @pytest.fixture
-def spill_client(spill_cluster, cluster_kwargs, benchmark_all):
+def spill_client(spill_cluster, cluster_kwargs, benchmark_all, wait_for_workers):
     n_workers = cluster_kwargs["spill_cluster"]["n_workers"]
     with Client(spill_cluster) as client:
         spill_cluster.scale(n_workers)
-        client.wait_for_workers(n_workers, timeout=600)
+        wait_for_workers(client, n_workers, timeout=600)
         client.restart()
         with benchmark_all(client):
             yield client

tests/benchmarks/test_xarray.py

@@ -36,11 +36,13 @@ def group_reduction_cluster(dask_env_variables, cluster_kwargs, github_cluster_t
 
 
 @pytest.fixture
-def group_reduction_client(group_reduction_cluster, cluster_kwargs, benchmark_all):
+def group_reduction_client(
+    group_reduction_cluster, cluster_kwargs, benchmark_all, wait_for_workers
+):
     n_workers = cluster_kwargs["group_reduction_cluster"]["n_workers"]
     with Client(group_reduction_cluster) as client:
         group_reduction_cluster.scale(n_workers)
-        client.wait_for_workers(n_workers, timeout=600)
+        wait_for_workers(client, n_workers, timeout=600)
         client.restart()
         with benchmark_all(client):
             yield client

tests/conftest.py

@@ -15,6 +15,7 @@
 from functools import lru_cache
 
 import adlfs
+import coiled
 import dask
 import dask.array as da
 import distributed
@@ -565,14 +566,15 @@ def small_client(
     small_cluster,
     cluster_kwargs,
     benchmark_all,
+    wait_for_workers,
 ):
     n_workers = cluster_kwargs["small_cluster"]["n_workers"]
     test_label = f"{request.node.name}, session_id={testrun_uid}"
     with Client(small_cluster) as client:
         log_on_scheduler(client, "Starting client setup of %s", test_label)
         client.restart()
         small_cluster.scale(n_workers)
-        client.wait_for_workers(n_workers, timeout=600)
+        wait_for_workers(client, n_workers, timeout=600)
 
         log_on_scheduler(client, "Finished client setup of %s", test_label)
 
@@ -895,3 +897,15 @@ def _performance_report():
             test_run_benchmark.performance_report_url = filename
 
     yield _performance_report
+
+
+@pytest.fixture
+def wait_for_workers():
+    def _(client, n_workers, timeout):
+        # https://github.com/coiled/platform/pull/8226
+        if Version(coiled.__version__) <= Version("1.87"):
+            client.sync(client._wait_for_workers, n_workers, timeout=timeout)
+        else:
+            client.wait_for_workers(n_workers, timeout=timeout)
+
+    return _