Reverse precedence between collection and distributed default (#9869)

fjetter · web-flow · commit 579b191fddb3 · 2023-01-26T18:09:49.000+01:00
diff --git a/dask/base.py b/dask/base.py
@@ -1345,7 +1345,8 @@ def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
 
     1.  Passing in scheduler= parameters
     2.  Passing these into global configuration
-    3.  Using defaults of a dask collection
+    3.  Using a dask.distributed default Client
+    4.  Using defaults of a dask collection
 
     This function centralizes the logic to determine the right scheduler to use
     from those many options
@@ -1407,16 +1408,16 @@ def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
     if config.get("get", None):
         raise ValueError(get_err_msg)
 
-    if cls is not None:
-        return cls.__dask_scheduler__
-
     try:
         from distributed import get_client
 
         return get_client().get
     except (ImportError, ValueError):
         pass
 
+    if cls is not None:
+        return cls.__dask_scheduler__
+
     if collections:
         collections = [c for c in collections if c is not None]
     if collections:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
@@ -22,7 +22,7 @@
 import dask
 import dask.bag as db
 from dask import compute, delayed, persist
-from dask.base import get_scheduler
+from dask.base import compute_as_if_collection, get_scheduler
 from dask.blockwise import Blockwise
 from dask.delayed import Delayed
 from dask.distributed import futures_of, wait
@@ -148,6 +148,62 @@ def test_fused_blockwise_dataframe_merge(c, fuse):
     )
 
 
+@pytest.mark.parametrize(
+    "computation",
+    [
+        None,
+        "compute_as_if_collection",
+        "dask.compute",
+    ],
+)
+@pytest.mark.parametrize(
+    "scheduler, use_distributed",
+    [
+        (None, True),
+        # If scheduler is explicitly provided, this takes precedence
+        ("sync", False),
+    ],
+)
+def test_default_scheduler_on_worker(c, computation, use_distributed, scheduler):
+    """Should a collection use its default scheduler or the distributed
+    scheduler when being computed within a task?
+    """
+
+    pd = pytest.importorskip("pandas")
+    dd = pytest.importorskip("dask.dataframe")
+
+    def foo():
+        size = 10
+        df = pd.DataFrame({"x": range(size), "y": range(size)})
+        ddf = dd.from_pandas(df, npartitions=2)
+        if computation is None:
+            ddf.compute(scheduler=scheduler)
+        elif computation == "dask.compute":
+            dask.compute(ddf, scheduler=scheduler)
+        elif computation == "compute_as_if_collection":
+            compute_as_if_collection(
+                ddf.__class__, ddf.dask, list(ddf.dask), scheduler=scheduler
+            )
+        else:
+            assert False
+
+        return True
+
+    res = c.submit(foo)
+    assert res.result() is True
+    # Count how many submits/update-graph were received by the scheduler
+    assert (
+        c.run_on_scheduler(
+            lambda dask_scheduler: sum(
+                len(comp.code) for comp in dask_scheduler.computations
+            )
+        )
+        == 2
+        if use_distributed
+        else 1
+    )
+
+
 def test_futures_to_delayed_bag(c):
     L = [1, 2, 3]
 
