Use distributed default clients even if no config is set (#9808)

fjetter · web-flow · commit dfcefd889131 · 2023-01-12T13:16:08.000-06:00
diff --git a/dask/base.py b/dask/base.py
@@ -26,7 +26,6 @@
 
 from dask import config, local
 from dask.compatibility import _EMSCRIPTEN, _PY_VERSION
-from dask.context import thread_state
 from dask.core import flatten
 from dask.core import get as simple_get
 from dask.core import literal, quote
@@ -1362,14 +1361,25 @@ def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
         elif isinstance(scheduler, str):
             scheduler = scheduler.lower()
 
+            try:
+                from distributed import default_client
+
+                default_client()
+                client_available = True
+            except (ImportError, ValueError):
+                client_available = False
             if scheduler in named_schedulers:
-                if config.get("scheduler", None) in ("dask.distributed", "distributed"):
+                if client_available:
                     warnings.warn(
                         "Running on a single-machine scheduler when a distributed client "
                         "is active might lead to unexpected results."
                     )
                 return named_schedulers[scheduler]
             elif scheduler in ("dask.distributed", "distributed"):
+                if not client_available:
+                    raise RuntimeError(
+                        f"Requested {scheduler} scheduler but no Client active."
+                    )
                 from distributed.worker import get_client
 
                 return get_client().get
@@ -1397,14 +1407,16 @@ def get_scheduler(get=None, scheduler=None, collections=None, cls=None):
     if config.get("get", None):
         raise ValueError(get_err_msg)
 
-    if getattr(thread_state, "key", False):
-        from distributed.worker import get_worker
-
-        return get_worker().client.get
-
     if cls is not None:
         return cls.__dask_scheduler__
 
+    try:
+        from distributed import get_client
+
+        return get_client().get
+    except (ImportError, ValueError):
+        pass
+
     if collections:
         collections = [c for c in collections if c is not None]
     if collections:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
@@ -1557,19 +1557,6 @@ def test_get_scheduler():
     assert get_scheduler() is None
 
 
-def test_get_scheduler_with_distributed_active():
-
-    with dask.config.set(scheduler="dask.distributed"):
-        warning_message = (
-            "Running on a single-machine scheduler when a distributed client "
-            "is active might lead to unexpected results."
-        )
-        with pytest.warns(UserWarning, match=warning_message) as user_warnings_a:
-            get_scheduler(scheduler="threads")
-            get_scheduler(scheduler="sync")
-        assert len(user_warnings_a) == 2
-
-
 def test_callable_scheduler():
     called = [False]
 
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
@@ -22,6 +22,7 @@
 import dask
 import dask.bag as db
 from dask import compute, delayed, persist
+from dask.base import get_scheduler
 from dask.blockwise import Blockwise
 from dask.delayed import Delayed
 from dask.distributed import futures_of, wait
@@ -804,3 +805,33 @@ def test_set_index_no_resursion_error(c):
         ddf.compute()
     except RecursionError:
         pytest.fail("dd.set_index triggered a recursion error")
+
+
+def test_get_scheduler_without_distributed_raises():
+    msg = "no Client"
+    with pytest.raises(RuntimeError, match=msg):
+        get_scheduler(scheduler="dask.distributed")
+
+    with pytest.raises(RuntimeError, match=msg):
+        get_scheduler(scheduler="distributed")
+
+
+def test_get_scheduler_with_distributed_active(c):
+    assert get_scheduler() == c.get
+    warning_message = (
+        "Running on a single-machine scheduler when a distributed client "
+        "is active might lead to unexpected results."
+    )
+    with pytest.warns(UserWarning, match=warning_message) as user_warnings_a:
+        get_scheduler(scheduler="threads")
+        get_scheduler(scheduler="sync")
+    assert len(user_warnings_a) == 2
+
+
+def test_get_scheduler_with_distributed_active_reset_config(c):
+    assert get_scheduler() == c.get
+    with dask.config.set(scheduler="threads"):
+        with pytest.warns(UserWarning):
+            assert get_scheduler() != c.get
+        with dask.config.set(scheduler=None):
+            assert get_scheduler() == c.get
