Repository: dask/distributed

Issue:
gpu CI failing pretty consistently with segfault 
I've noticed gpuCI has been failing pretty consistently (for example, [this build](https://gpuci.gpuopenanalytics.com/job/dask/job/distributed/job/prb/job/distributed-prb/7578/CUDA_VER=11.5,LINUX_VER=ubuntu18.04,PYTHON_VER=3.9,RAPIDS_VER=23.10/console) and [this build](https://gpuci.gpuopenanalytics.com/job/dask/job/distributed/job/prb/job/distributed-prb/7574/CUDA_VER=11.5,LINUX_VER=ubuntu18.04,PYTHON_VER=3.9,RAPIDS_VER=23.10/console)) with a segfault. Note that tests actually pass -- there must be some extra step where things are going awry 

```
09:26:31 SKIPPED [2] distributed/diagnostics/tests/test_nvml.py:107: Less than two GPUs available
09:26:31 SKIPPED [1] distributed/shuffle/tests/test_shuffle.py:134: could not import 'dask_cudf': No module named 'dask_cudf'
09:26:31 XFAIL distributed/comm/tests/test_ucx.py::test_cuda_context - If running on Docker, requires --pid=host
09:26:31 XFAIL distributed/diagnostics/tests/test_nvml.py::test_has_cuda_context - If running on Docker, requires --pid=host
09:26:31 XFAIL distributed/tests/test_nanny.py::test_no_unnecessary_imports_on_worker[pandas] - distributed#5723
09:26:31 ==== 372 passed, 8 skipped, 3702 deselected, 3 xfailed in 257.84s (0:04:17) ====
09:26:32 continuous_integration/gpuci/build.sh: line 58:   173 Segmentation fault      (core dumped) py.test distributed -v -m gpu --runslow --junitxml="$WORKSPACE/junit-distributed.xml"
09:26:33 Build step 'Execute shell' marked build as failure
```

cc @charlesbluca @quasiben @dask/gpu 

Discussion:
pentschev: rerun tests
pentschev: It seems gpuCI is failing to resolve github.com, I've raised an internal issue to report that.
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

       21 files  ±  0         21 suites  ±0   10h 53m 10s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "duration of all tests") + 36m 35s
  3 813 tests ±  0    3 701 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "passed tests")  -   3     107 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "skipped / disabled tests") ±0  5 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "failed tests") +3 
36 855 runs  +12  35 049 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "passed tests") +11  1 799 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "skipped / disabled tests")  - 4  7 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.9.0/README.md#the-symbols "failed tests") +5 

For more details on these failures, see [this check](https://github.com/dask/distributed/runs/17002514123).

Results for commit 40285e73. ± Comparison against base commit 28589303.

[test-results]:data:application/gzip;base64,H4sIAJE5DGUC/03MywrDIBCF4VcJrrsYLxO1L1PETECaSzG6Kn33Tkpiuvy/A+ctxjTRJu6dkrdObDWVq4aaQ0nrwq299MDEa9l37aQ+87HVGHezIC97phebBNtoDGliwgaU85pZ9uNcl99v7xCPbL8Ixjc7fq1vdPzas/9v4zrPqXAIA8ohWU2DctJCbxSRUaBxxAGRM5oQyaD4fAHYOmoNFAEAAA==

pentschev: rerun tests
pentschev: > It seems gpuCI is failing to resolve github.com, I've raised an internal issue to report that.

This is now resolved.
pentschev: rerun tests
pentschev: rerun tests
pentschev: rerun tests

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index d02327ca..1422394b 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -35,6 +35,7 @@ from dask.base import collections_to_dsk, normalize_token, tokenize
 from dask.core import flatten, validate_key
 from dask.highlevelgraph import HighLevelGraph
 from dask.optimization import SubgraphCallable
+from dask.typing import no_default
 from dask.utils import (
     apply,
     ensure_dict,
@@ -101,7 +102,6 @@ from distributed.utils import (
     import_term,
     is_python_shutting_down,
     log_errors,
-    no_default,
     sync,
     thread_state,
 )
@@ -854,7 +854,7 @@ class Client(SyncMethodMixin):
         connection_limit=512,
         **kwargs,
     ):
-        if timeout == no_default:
+        if timeout is no_default:
             timeout = dask.config.get("distributed.comm.timeouts.connect")
         if timeout is not None:
             timeout = parse_timedelta(timeout, "s")
@@ -1248,7 +1248,7 @@ class Client(SyncMethodMixin):
 
         await self.rpc.start()
 
-        if timeout == no_default:
+        if timeout is no_default:
             timeout = self._timeout
         if timeout is not None:
             timeout = parse_timedelta(timeout, "s")
@@ -1753,7 +1753,7 @@ class Client(SyncMethodMixin):
         --------
         Client.restart
         """
-        if timeout == no_default:
+        if timeout is no_default:
             timeout = self._timeout * 2
         # XXX handling of self.status here is not thread-safe
         if self.status in ["closed", "newly-created"]:
@@ -2399,7 +2399,7 @@ class Client(SyncMethodMixin):
         timeout=no_default,
         hash=True,
     ):
-        if timeout == no_default:
+        if timeout is no_default:
             timeout = self._timeout
         if isinstance(workers, (str, Number)):
             workers = [workers]
@@ -2588,7 +2588,7 @@ class Client(SyncMethodMixin):
         --------
         Client.gather : Gather data back to local process
         """
-        if timeout == no_default:
+        if timeout is no_default:
             timeout = self._timeout
         if isinstance(data, pyQueue) or isinstance(data, Iterator):
             raise TypeError(
@@ -3577,7 +3577,7 @@ class Client(SyncMethodMixin):
             return result
 
     async def _restart(self, timeout=no_default, wait_for_workers=True):
-        if timeout == no_default:
+        if timeout is no_default:
             timeout = self._timeout * 4
         if timeout is not None:
             timeout = parse_timedelta(timeout, "s")
diff --git a/distributed/deploy/cluster.py b/distributed/deploy/cluster.py
index 514b9745..6a5d0238 100644
--- a/distributed/deploy/cluster.py
+++ b/distributed/deploy/cluster.py
@@ -33,9 +33,6 @@ from distributed.utils import (
 logger = logging.getLogger(__name__)
 
 
-no_default = "__no_default__"
-
-
 class Cluster(SyncMethodMixin):
     """Superclass for cluster objects
 
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 98297edc..63f5092b 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -54,6 +54,7 @@ from tornado.ioloop import IOLoop
 import dask
 import dask.utils
 from dask.core import get_deps, validate_key
+from dask.typing import no_default
 from dask.utils import (
     format_bytes,
     format_time,
@@ -119,7 +120,6 @@ from distributed.utils import (
     get_fileno_limit,
     key_split_group,
     log_errors,
-    no_default,
     offload,
     recursive_to_dict,
     wait_for,
@@ -7419,7 +7419,7 @@ class Scheduler(SchedulerState, ServerNode):
                 metadata = metadata[key]
             return metadata
         except KeyError:
-            if default != no_default:
+            if default is not no_default:
                 return default
             else:
                 raise
@@ -8449,8 +8449,8 @@ class KilledWorker(Exception):
 
     def __str__(self) -> str:
         return (
-            f"Attempted to run task {self.task} on {self.allowed_failures} different "
-            "workers, but all those workers died while running it. "
+            f"Attempted to run task {self.task} on {self.allowed_failures + 1} "
+            "different workers, but all those workers died while running it. "
             f"The last worker that attempt to run the task was {self.last_worker.address}. "
             "Inspecting worker logs is often a good next step to diagnose what went wrong. "
             "For more information see https://distributed.dask.org/en/stable/killed.html."
diff --git a/distributed/utils.py b/distributed/utils.py
index 45f5e48d..093fda23 100644
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -95,8 +95,6 @@ if TYPE_CHECKING:
     P = ParamSpec("P")
     T = TypeVar("T")
 
-no_default = "__no_default__"
-
 _forkserver_preload_set = False
 
 
@@ -1604,7 +1602,9 @@ def clean_dashboard_address(addrs: AnyType, default_listen_ip: str = "") -> list
     return addresses
 
 
-_deprecations: dict[str, str] = {}
+_deprecations = {
+    "no_default": "dask.typing.no_default",
+}
 
 
 def __getattr__(name):


Test patch:
diff --git a/distributed/comm/tests/test_ucx.py b/distributed/comm/tests/test_ucx.py
index 04d4cd18..5ecefcd6 100644
--- a/distributed/comm/tests/test_ucx.py
+++ b/distributed/comm/tests/test_ucx.py
@@ -179,7 +179,7 @@ async def test_ucx_deserialize(ucx_loop):
     "g",
     [
         lambda cudf: cudf.Series([1, 2, 3]),
-        lambda cudf: cudf.Series([]),
+        lambda cudf: cudf.Series([], dtype=object),
         lambda cudf: cudf.DataFrame([]),
         lambda cudf: cudf.DataFrame([1]).head(0),
         lambda cudf: cudf.DataFrame([1.0]).head(0),
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index 19b19ce9..8d5f68f4 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -4244,8 +4244,8 @@ async def test_KilledWorker_informative_message(s, a, b):
     with pytest.raises(KilledWorker) as excinfo:
         raise ex
     msg = str(excinfo.value)
-    assert "Attempted to run task foo-bar" in msg
-    assert str(s.allowed_failures) in msg
+    assert "Attempted to run task foo-bar on 667 different workers" in msg
+    assert a.address in msg
     assert "worker logs" in msg
     assert "https://distributed.dask.org/en/stable/killed.html" in msg
 
diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py
index 1be18142..39378d3b 100644
--- a/distributed/tests/test_utils.py
+++ b/distributed/tests/test_utils.py
@@ -700,6 +700,12 @@ async def test_offload_preserves_contextvars():
     await asyncio.gather(set_var("foo"), set_var("bar"))
 
 
+def test_no_default_deprecated():
+    with pytest.warns(FutureWarning, match="no_default is deprecated"):
+        from distributed.utils import no_default
+    assert no_default is dask.typing.no_default
+
+
 def test_iscoroutinefunction_unhashable_input():
     # Ensure iscoroutinefunction can handle unhashable callables
     assert not iscoroutinefunction(_UnhashableCallable())
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index 69e4d113..92aa0089 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -2156,17 +2156,6 @@ def ucx_loop():
     import distributed.comm.ucx
 
     distributed.comm.ucx.ucp = None
-    # If the test created a context, clean it up.
-    # TODO: should we check if there's already a context _before_ the test runs?
-    # I think that would be useful.
-    from distributed.diagnostics.nvml import has_cuda_context
-
-    ctx = has_cuda_context()
-    if ctx.has_context:
-        import numba.cuda
-
-        ctx = numba.cuda.current_context()
-        ctx.device.reset()
 
 
 def wait_for_log_line(