Repository: dask/distributed

Issue:
dask-scheduler --jupyter fails to start
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:
I have been using `Scheduler(jupyter=True)` for a while without issue; I just tried the CLI version and it seems to fail.


**Minimal Complete Verifiable Example**:

```
(.venv)  ➜  distributed git:(main) dask-scheduler --jupyter
/Users/brett/distributed/distributed/cli/dask_scheduler.py:140: FutureWarning: dask-scheduler is deprecated and will be removed in a future release; use `dask scheduler` instead
  warnings.warn(
2023-02-21 18:29:56,114 - distributed.scheduler - INFO - -----------------------------------------------
2023-02-21 18:29:56,130 - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
usage: dask-scheduler [-h] [--debug] [--show-config] [--show-config-json] [--generate-config] [-y]
                      [--allow-root] [--no-browser] [--autoreload] [--script] [--no-script]
                      [--log-level ServerApp.log_level] [--config ServerApp.config_file]
                      [--ip ServerApp.ip] [--port ServerApp.port]
                      [--port-retries ServerApp.port_retries] [--sock ServerApp.sock]
                      [--sock-mode ServerApp.sock_mode] [--transport KernelManager.transport]
                      [--keyfile ServerApp.keyfile] [--certfile ServerApp.certfile]
                      [--client-ca ServerApp.client_ca] [--notebook-dir ServerApp.root_dir]
                      [--preferred-dir ServerApp.preferred_dir] [--browser ServerApp.browser]
                      [--pylab ServerApp.pylab] [--gateway-url GatewayClient.url]
                      [extra_args ...]
dask-scheduler: error: argument --jupyter: expected one argument
```

**Anything else we need to know?**:
Adding `sys.argv.remove("--jupyter")` right before `Scheduler()` gets called [here](https://github.com/dask/distributed/blob/main/distributed/cli/dask_scheduler.py#L209) resolves the issue
```
(.venv)  ➜  distributed git:(main) ✗ dask-scheduler --jupyter
/Users/brett/distributed/distributed/cli/dask_scheduler.py:140: FutureWarning: dask-scheduler is deprecated and will be removed in a future release; use `dask scheduler` instead
  warnings.warn(
2023-02-21 18:32:27,072 - distributed.scheduler - INFO - -----------------------------------------------
2023-02-21 18:32:27,081 - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
[I 2023-02-21 18:32:27.752 ServerApp] jupyterlab | extension was successfully linked.
[I 2023-02-21 18:32:27.754 ServerApp] nbclassic | extension was successfully linked.
[I 2023-02-21 18:32:27.861 ServerApp] notebook_shim | extension was successfully linked.
[W 2023-02-21 18:32:27.880 ServerApp] All authentication is disabled.  Anyone who can connect to this server will be able to run code.
[I 2023-02-21 18:32:27.881 ServerApp] notebook_shim | extension was successfully loaded.
[I 2023-02-21 18:32:27.881 LabApp] JupyterLab extension loaded from /Users/brett/distributed/.venv/lib/python3.10/site-packages/jupyterlab
[I 2023-02-21 18:32:27.881 LabApp] JupyterLab application directory is /Users/brett/distributed/.venv/share/jupyter/lab
[I 2023-02-21 18:32:27.883 ServerApp] jupyterlab | extension was successfully loaded.
```
So it seems like something down in the Jupyter initialization stack is getting confused? I also tried on an old Docker image from back when this feature was first released and it seems like it didn't work then either, so I don't think it's anything to do with recent updates in `jupyter-server` etc.

**Environment**:

- Dask version: current main `9e8876dba90e4fac5101a1a3dba0f658b279a243`
- Python version: 3.10

Discussion:
mrocklin: This seems fine to me.  I'll merge later this afternoon if there are no further comments.

Patch:
diff --git a/.github/workflows/update-gpuci.yaml b/.github/workflows/update-gpuci.yaml
index d40aca0c..ab15bbc5 100644
--- a/.github/workflows/update-gpuci.yaml
+++ b/.github/workflows/update-gpuci.yaml
@@ -42,8 +42,8 @@ jobs:
         run: |
           echo RAPIDS_VER=${{ steps.rapids_current.outputs.RAPIDS_VER_0 }} >> $GITHUB_ENV
           echo UCX_PY_VER=$(curl -sL https://version.gpuci.io/rapids/${{ steps.rapids_current.outputs.RAPIDS_VER_0 }}) >> $GITHUB_ENV
-          echo NEW_RAPIDS_VER=${FULL_RAPIDS_VER::-10} >> $GITHUB_ENV
-          echo NEW_UCX_PY_VER=${FULL_UCX_PY_VER::-10} >> $GITHUB_ENV
+          echo NEW_RAPIDS_VER=${FULL_RAPIDS_VER::-4} >> $GITHUB_ENV
+          echo NEW_UCX_PY_VER=${FULL_UCX_PY_VER::-4} >> $GITHUB_ENV
 
       - name: Update RAPIDS version
         uses: jacobtomlinson/gha-find-replace@v3
diff --git a/distributed/core.py b/distributed/core.py
index a72b2ed0..50ccfa41 100644
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -11,10 +11,10 @@ import uuid
 import warnings
 import weakref
 from collections import defaultdict, deque
-from collections.abc import Container, Coroutine
+from collections.abc import Callable, Container, Coroutine, Generator
 from enum import Enum
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, TypedDict, TypeVar, final
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, TypeVar, final
 
 import tblib
 from tlz import merge
@@ -46,7 +46,7 @@ from distributed.utils import (
 )
 
 if TYPE_CHECKING:
-    from typing_extensions import ParamSpec
+    from typing_extensions import ParamSpec, Self
 
     P = ParamSpec("P")
     R = TypeVar("R")
@@ -1014,13 +1014,19 @@ async def send_recv(  # type: ignore[no-untyped-def]
     return response
 
 
-def addr_from_args(addr=None, ip=None, port=None):
+def addr_from_args(
+    addr: str | tuple[str, int | None] | None = None,
+    ip: str | None = None,
+    port: int | None = None,
+) -> str:
     if addr is None:
-        addr = (ip, port)
-    else:
-        assert ip is None and port is None
+        assert ip is not None
+        return normalize_address(unparse_host_port(ip, port))
+
+    assert ip is None and port is None
     if isinstance(addr, tuple):
-        addr = unparse_host_port(*addr)
+        return normalize_address(unparse_host_port(*addr))
+
     return normalize_address(addr)
 
 
@@ -1288,20 +1294,20 @@ class ConnectionPool:
 
     def __init__(
         self,
-        limit=512,
-        deserialize=True,
-        serializers=None,
-        allow_offload=True,
-        deserializers=None,
-        connection_args=None,
-        timeout=None,
-        server=None,
-    ):
+        limit: int = 512,
+        deserialize: bool = True,
+        serializers: list[str] | None = None,
+        allow_offload: bool = True,
+        deserializers: list[str] | None = None,
+        connection_args: dict[str, object] | None = None,
+        timeout: float | None = None,
+        server: object = None,
+    ) -> None:
         self.limit = limit  # Max number of open comms
         # Invariant: len(available) == open - active
-        self.available = defaultdict(set)
+        self.available: defaultdict[str, set[Comm]] = defaultdict(set)
         # Invariant: len(occupied) == active
-        self.occupied = defaultdict(set)
+        self.occupied: defaultdict[str, set[Comm]] = defaultdict(set)
         self.allow_offload = allow_offload
         self.deserialize = deserialize
         self.serializers = serializers
@@ -1309,18 +1315,21 @@ class ConnectionPool:
         self.connection_args = connection_args or {}
         self.timeout = timeout
         self.server = weakref.ref(server) if server else None
-        self._created = weakref.WeakSet()
+        self._created: weakref.WeakSet[Comm] = weakref.WeakSet()
         self._instances.add(self)
         # _n_connecting and _connecting have subtle different semantics. The set
         # _connecting contains futures actively trying to establish a connection
         # while the _n_connecting also accounts for connection attempts which
         # are waiting due to the connection limit
-        self._connecting = set()
+        self._connecting: defaultdict[str, set[asyncio.Task[Comm]]] = defaultdict(set)
         self._pending_count = 0
         self._connecting_count = 0
         self.status = Status.init
+        self._reasons: weakref.WeakKeyDictionary[
+            asyncio.Task[Any], str
+        ] = weakref.WeakKeyDictionary()
 
-    def _validate(self):
+    def _validate(self) -> None:
         """
         Validate important invariants of this class
 
@@ -1329,35 +1338,40 @@ class ConnectionPool:
         assert self.semaphore._value == self.limit - self.open - self._n_connecting
 
     @property
-    def active(self):
+    def active(self) -> int:
         return sum(map(len, self.occupied.values()))
 
     @property
-    def open(self):
+    def open(self) -> int:
         return self.active + sum(map(len, self.available.values()))
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "<ConnectionPool: open=%d, active=%d, connecting=%d>" % (
             self.open,
             self.active,
             len(self._connecting),
         )
 
-    def __call__(self, addr=None, ip=None, port=None):
+    def __call__(
+        self,
+        addr: str | tuple[str, int | None] | None = None,
+        ip: str | None = None,
+        port: int | None = None,
+    ) -> PooledRPCCall:
         """Cached rpc objects"""
         addr = addr_from_args(addr=addr, ip=ip, port=port)
         return PooledRPCCall(
             addr, self, serializers=self.serializers, deserializers=self.deserializers
         )
 
-    def __await__(self):
-        async def _():
+    def __await__(self) -> Generator[Any, Any, Self]:
+        async def _() -> Self:
             await self.start()
             return self
 
         return _().__await__()
 
-    async def start(self):
+    async def start(self) -> None:
         # Invariant: semaphore._value == limit - open - _n_connecting
         self.semaphore = asyncio.Semaphore(self.limit)
         self.status = Status.running
@@ -1366,7 +1380,7 @@ class ConnectionPool:
     def _n_connecting(self) -> int:
         return self._connecting_count
 
-    async def _connect(self, addr, timeout=None):
+    async def _connect(self, addr: str, timeout: float | None = None) -> Comm:
         self._pending_count += 1
         try:
             await self.semaphore.acquire()
@@ -1392,11 +1406,14 @@ class ConnectionPool:
             finally:
                 self._connecting_count -= 1
         except asyncio.CancelledError:
-            raise CommClosedError("ConnectionPool closing.")
+            current_task = asyncio.current_task()
+            assert current_task
+            reason = self._reasons.pop(current_task, "ConnectionPool closing.")
+            raise CommClosedError(reason)
         finally:
             self._pending_count -= 1
 
-    async def connect(self, addr, timeout=None):
+    async def connect(self, addr: str, timeout: float | None = None) -> Comm:
         """
         Get a Comm to the given address.  For internal use.
         """
@@ -1422,9 +1439,21 @@ class ConnectionPool:
         # it to propagate
         connect_attempt = asyncio.create_task(self._connect(addr, timeout))
         done = asyncio.Event()
-        self._connecting.add(connect_attempt)
-        connect_attempt.add_done_callback(lambda _: done.set())
-        connect_attempt.add_done_callback(self._connecting.discard)
+        connecting = self._connecting[addr]
+        connecting.add(connect_attempt)
+
+        def callback(task: asyncio.Task[Comm]) -> None:
+            done.set()
+            connecting = self._connecting[addr]
+            connecting.discard(task)
+
+            if not connecting:
+                try:
+                    del self._connecting[addr]
+                except KeyError:  # pragma: no cover
+                    pass
+
+        connect_attempt.add_done_callback(callback)
 
         try:
             await done.wait()
@@ -1438,7 +1467,7 @@ class ConnectionPool:
             raise
         return await connect_attempt
 
-    def reuse(self, addr, comm):
+    def reuse(self, addr: str, comm: Comm) -> None:
         """
         Reuse an open communication to the given address.  For internal use.
         """
@@ -1457,7 +1486,7 @@ class ConnectionPool:
                 if self.semaphore.locked() and self._pending_count:
                     self.collect()
 
-    def collect(self):
+    def collect(self) -> None:
         """
         Collect open but unused communications, to allow opening other ones.
         """
@@ -1473,7 +1502,7 @@ class ConnectionPool:
                 self.semaphore.release()
             comms.clear()
 
-    def remove(self, addr):
+    def remove(self, addr: str, *, reason: str = "Address removed.") -> None:
         """
         Remove all Comms to a given address.
         """
@@ -1489,13 +1518,20 @@ class ConnectionPool:
                 IOLoop.current().add_callback(comm.close)
                 self.semaphore.release()
 
-    async def close(self):
+        if addr in self._connecting:
+            tasks = self._connecting[addr]
+            for task in tasks:
+                self._reasons[task] = reason
+                task.cancel()
+
+    async def close(self) -> None:
         """
         Close all communications
         """
         self.status = Status.closed
-        for conn_fut in self._connecting:
-            conn_fut.cancel()
+        for tasks in self._connecting.values():
+            for task in tasks:
+                task.cancel()
         for d in [self.available, self.occupied]:
             comms = set()
             while d:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index b9bc05a6..7daeb640 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -3510,6 +3510,7 @@ class Scheduler(SchedulerState, ServerNode):
             )
             j.initialize(
                 new_httpserver=False,
+                argv=[],
             )
             self._jupyter_server_application = j
             self.http_application.add_application(j.web_app)
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
index f0dc722e..bf1aae98 100644
--- a/distributed/shuffle/_arrow.py
+++ b/distributed/shuffle/_arrow.py
@@ -4,9 +4,30 @@ from io import BytesIO
 from typing import TYPE_CHECKING, BinaryIO
 
 if TYPE_CHECKING:
+    import pandas as pd
     import pyarrow as pa
 
 
+def check_dtype_support(meta_input: pd.DataFrame) -> None:
+    import pandas as pd
+
+    for name in meta_input:
+        column = meta_input[name]
+        # FIXME: PyArrow does not support complex numbers: https://issues.apache.org/jira/browse/ARROW-638
+        if pd.api.types.is_complex_dtype(column):
+            raise TypeError(
+                f"p2p does not support data of type '{column.dtype}' found in column '{name}'."
+            )
+        # FIXME: Serializing custom objects to PyArrow is not supported in P2P shuffling
+        if pd.api.types.is_object_dtype(column):
+            raise TypeError(
+                f"p2p does not support custom objects found in column '{name}'."
+            )
+        # FIXME: PyArrow does not support sparse data: https://issues.apache.org/jira/browse/ARROW-8679
+        if pd.api.types.is_sparse(column):
+            raise TypeError("p2p does not support sparse data found in column '{name}'")
+
+
 def dump_shards(shards: list[bytes], file: BinaryIO) -> None:
     """
     Write multiple shard tables to the file
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
index 8aa0adf7..b1eee5bc 100644
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -7,6 +7,8 @@ from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import SimpleShuffleLayer
 
+from distributed.shuffle._arrow import check_dtype_support
+
 logger = logging.getLogger("distributed.shuffle")
 if TYPE_CHECKING:
     import pandas as pd
@@ -94,6 +96,7 @@ def rearrange_by_column_p2p(
 ) -> DataFrame:
     from dask.dataframe import DataFrame
 
+    check_dtype_support(df._meta)
     npartitions = npartitions or df.npartitions
     token = tokenize(df, column, npartitions)
 
@@ -103,12 +106,6 @@ def rearrange_by_column_p2p(
         raise TypeError(
             f"p2p requires all column names to be str, found: {unsupported}",
         )
-    for c, dt in empty.dtypes.items():
-        if dt == object:
-            empty[c] = empty[c].astype(
-                "string"
-            )  # TODO: we fail at non-string object dtypes
-    empty[column] = empty[column].astype("int64")  # TODO: this shouldn't be necesssary
 
     name = f"shuffle-p2p-{token}"
     layer = P2PShuffleLayer(
diff --git a/distributed/utils_comm.py b/distributed/utils_comm.py
index 7a11b4c3..e78f303b 100644
--- a/distributed/utils_comm.py
+++ b/distributed/utils_comm.py
@@ -4,9 +4,10 @@ import asyncio
 import logging
 import random
 from collections import defaultdict
+from collections.abc import Callable, Collection, Coroutine, Mapping
 from functools import partial
 from itertools import cycle
-from typing import Any, Callable
+from typing import Any, TypeVar
 
 from tlz import concat, drop, groupby, merge
 
@@ -14,13 +15,19 @@ import dask.config
 from dask.optimization import SubgraphCallable
 from dask.utils import is_namedtuple_instance, parse_timedelta, stringify
 
-from distributed.core import rpc
+from distributed.core import ConnectionPool, rpc
 from distributed.utils import All
 
 logger = logging.getLogger(__name__)
 
 
-async def gather_from_workers(who_has, rpc, close=True, serializers=None, who=None):
+async def gather_from_workers(
+    who_has: Mapping[str, Collection[str]],
+    rpc: ConnectionPool,
+    close: bool = True,
+    serializers: list[str] | None = None,
+    who: str | None = None,
+) -> tuple[dict[str, object], dict[str, list[str]], list[str]]:
     """Gather data directly from peers
 
     Parameters
@@ -38,18 +45,18 @@ async def gather_from_workers(who_has, rpc, close=True, serializers=None, who=No
     """
     from distributed.worker import get_data_from_worker
 
-    bad_addresses = set()
+    bad_addresses: set[str] = set()
     missing_workers = set()
     original_who_has = who_has
-    who_has = {k: set(v) for k, v in who_has.items()}
-    results = dict()
-    all_bad_keys = set()
+    new_who_has = {k: set(v) for k, v in who_has.items()}
+    results: dict[str, object] = {}
+    all_bad_keys: set[str] = set()
 
     while len(results) + len(all_bad_keys) < len(who_has):
         d = defaultdict(list)
         rev = dict()
         bad_keys = set()
-        for key, addresses in who_has.items():
+        for key, addresses in new_who_has.items():
             if key in results:
                 continue
             try:
@@ -62,19 +69,23 @@ async def gather_from_workers(who_has, rpc, close=True, serializers=None, who=No
             all_bad_keys |= bad_keys
         coroutines = {
             address: asyncio.create_task(
-                get_data_from_worker(
-                    rpc,
-                    keys,
-                    address,
-                    who=who,
-                    serializers=serializers,
-                    max_connections=False,
+                retry_operation(
+                    partial(
+                        get_data_from_worker,
+                        rpc,
+                        keys,
+                        address,
+                        who=who,
+                        serializers=serializers,
+                        max_connections=False,
+                    ),
+                    operation="get_data_from_worker",
                 ),
                 name=f"get-data-from-{address}",
             )
             for address, keys in d.items()
         }
-        response = {}
+        response: dict[str, object] = {}
         for worker, c in coroutines.items():
             try:
                 r = await c
@@ -91,8 +102,11 @@ async def gather_from_workers(who_has, rpc, close=True, serializers=None, who=No
         bad_addresses |= {v for k, v in rev.items() if k not in response}
         results.update(response)
 
-    bad_keys = {k: list(original_who_has[k]) for k in all_bad_keys}
-    return (results, bad_keys, list(missing_workers))
+    return (
+        results,
+        {k: list(original_who_has[k]) for k in all_bad_keys},
+        list(missing_workers),
+    )
 
 
 class WrappedKey:
@@ -337,15 +351,19 @@ def subs_multiple(o, d):
             return o
 
 
+T = TypeVar("T")
+
+
 async def retry(
-    coro,
-    count,
-    delay_min,
-    delay_max,
-    jitter_fraction=0.1,
-    retry_on_exceptions=(EnvironmentError, IOError),
-    operation=None,
-):
+    coro: Callable[[], Coroutine[Any, Any, T]],
+    count: int,
+    delay_min: float,
+    delay_max: float,
+    jitter_fraction: float = 0.1,
+    retry_on_exceptions: type[BaseException]
+    | tuple[type[BaseException], ...] = (EnvironmentError, IOError),
+    operation: str | None = None,
+) -> T:
     """
     Return the result of ``await coro()``, re-trying in case of exceptions
 
@@ -395,7 +413,13 @@ async def retry(
     return await coro()
 
 
-async def retry_operation(coro, *args, operation=None, **kwargs):
+# ParamSpec is not supported here due to the additional "operation" kwarg
+async def retry_operation(
+    coro: Callable[..., Coroutine[Any, Any, T]],
+    *args: object,
+    operation: str | None = None,
+    **kwargs: object,
+) -> T:
     """
     Retry an operation using the configuration values for the retry parameters
     """
diff --git a/distributed/worker.py b/distributed/worker.py
index c6cc60a6..c29cd2bd 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -28,7 +28,17 @@ from concurrent.futures import Executor
 from contextlib import suppress
 from datetime import timedelta
 from inspect import isawaitable
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, TextIO, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Literal,
+    TextIO,
+    TypedDict,
+    TypeVar,
+    cast,
+    overload,
+)
 
 from tlz import first, keymap, pluck
 from tornado.ioloop import IOLoop
@@ -166,6 +176,15 @@ WORKER_ANY_RUNNING = {
 }
 
 
+class GetDataBusy(TypedDict):
+    status: Literal["busy"]
+
+
+class GetDataSuccess(TypedDict):
+    status: Literal["OK"]
+    data: dict[str, object]
+
+
 def fail_hard(method: Callable[P, T]) -> Callable[P, T]:
     """
     Decorator to close the worker if this method encounters an exception.
@@ -1304,7 +1323,7 @@ class Worker(BaseWorker, ServerNode):
             if k not in self.data
         }
         result, missing_keys, missing_workers = await gather_from_workers(
-            who_has, rpc=self.rpc, who=self.address
+            who_has=who_has, rpc=self.rpc, who=self.address
         )
         self.update_data(data=result)
         if missing_keys:
@@ -1693,13 +1712,21 @@ class Worker(BaseWorker, ServerNode):
         self.stream_comms[address].send(msg)
 
     async def get_data(
-        self, comm, keys=None, who=None, serializers=None, max_connections=None
-    ) -> dict | Status:
+        self,
+        comm: Comm,
+        keys: Collection[str] | None = None,
+        who: str | None = None,
+        serializers: list[str] | None = None,
+        max_connections: int | None = None,
+    ) -> GetDataBusy | Literal[Status.dont_reply]:
         start = time()
 
         if max_connections is None:
             max_connections = self.transfer_outgoing_count_limit
 
+        if keys is None:
+            keys = set()
+
         # Allow same-host connections more liberally
         if (
             max_connections
@@ -2052,7 +2079,7 @@ class Worker(BaseWorker, ServerNode):
 
             start = time()
             response = await get_data_from_worker(
-                self.rpc, to_gather, worker, who=self.address
+                rpc=self.rpc, keys=to_gather, worker=worker, who=self.address
             )
             stop = time()
             if response["status"] == "busy":
@@ -2816,15 +2843,44 @@ class Reschedule(Exception):
     """
 
 
+@overload
 async def get_data_from_worker(
-    rpc,
-    keys,
-    worker,
-    who=None,
-    max_connections=None,
-    serializers=None,
-    deserializers=None,
-):
+    rpc: ConnectionPool,
+    keys: Collection[str],
+    worker: str,
+    *,
+    who: str | None = None,
+    max_connections: Literal[False],
+    serializers: list[str] | None = None,
+    deserializers: list[str] | None = None,
+) -> GetDataSuccess:
+    ...
+
+
+@overload
+async def get_data_from_worker(
+    rpc: ConnectionPool,
+    keys: Collection[str],
+    worker: str,
+    *,
+    who: str | None = None,
+    max_connections: bool | int | None = None,
+    serializers: list[str] | None = None,
+    deserializers: list[str] | None = None,
+) -> GetDataBusy | GetDataSuccess:
+    ...
+
+
+async def get_data_from_worker(
+    rpc: ConnectionPool,
+    keys: Collection[str],
+    worker: str,
+    *,
+    who: str | None = None,
+    max_connections: bool | int | None = None,
+    serializers: list[str] | None = None,
+    deserializers: list[str] | None = None,
+) -> GetDataBusy | GetDataSuccess:
     """Get keys from worker
 
     The worker has a two step handshake to acknowledge when data has been fully
@@ -2841,31 +2897,28 @@ async def get_data_from_worker(
     if deserializers is None:
         deserializers = rpc.deserializers
 
-    async def _get_data():
-        comm = await rpc.connect(worker)
-        comm.name = "Ephemeral Worker->Worker for gather"
+    comm = await rpc.connect(worker)
+    comm.name = "Ephemeral Worker->Worker for gather"
+    try:
+        response = await send_recv(
+            comm,
+            serializers=serializers,
+            deserializers=deserializers,
+            op="get_data",
+            keys=keys,
+            who=who,
+            max_connections=max_connections,
+        )
         try:
-            response = await send_recv(
-                comm,
-                serializers=serializers,
-                deserializers=deserializers,
-                op="get_data",
-                keys=keys,
-                who=who,
-                max_connections=max_connections,
-            )
-            try:
-                status = response["status"]
-            except KeyError:  # pragma: no cover
-                raise ValueError("Unexpected response", response)
-            else:
-                if status == "OK":
-                    await comm.write("OK")
-            return response
-        finally:
-            rpc.reuse(worker, comm)
-
-    return await retry_operation(_get_data, operation="get_data_from_worker")
+            status = response["status"]
+        except KeyError:  # pragma: no cover
+            raise ValueError("Unexpected response", response)
+        else:
+            if status == "OK":
+                await comm.write("OK")
+        return response
+    finally:
+        rpc.reuse(worker, comm)
 
 
 job_counter = [0]


Test patch:
diff --git a/distributed/dashboard/tests/test_scheduler_bokeh.py b/distributed/dashboard/tests/test_scheduler_bokeh.py
index 80b0ba29..909b122c 100644
--- a/distributed/dashboard/tests/test_scheduler_bokeh.py
+++ b/distributed/dashboard/tests/test_scheduler_bokeh.py
@@ -1118,6 +1118,7 @@ async def test_shuffling(c, s, a, b):
     ss = Shuffling(s)
 
     df = dask.datasets.timeseries()
+    df["name"] = df["name"].astype("string[python]")
     df2 = dd.shuffle.shuffle(df, "x", shuffle="p2p").persist()
     start = time()
     while not ss.source.data["disk_read"]:
diff --git a/distributed/shuffle/tests/test_graph.py b/distributed/shuffle/tests/test_graph.py
index 4f2d6add..6f129928 100644
--- a/distributed/shuffle/tests/test_graph.py
+++ b/distributed/shuffle/tests/test_graph.py
@@ -29,6 +29,38 @@ def test_basic(client):
     # ^ NOTE: this works because `assert_eq` sorts the rows before comparing
 
 
+@pytest.mark.parametrize("dtype", ["csingle", "cdouble", "clongdouble"])
+def test_raise_on_complex_numbers(dtype):
+    df = dd.from_pandas(
+        pd.DataFrame({"x": pd.array(range(10), dtype=dtype)}), npartitions=5
+    )
+    with pytest.raises(
+        TypeError, match=f"p2p does not support data of type '{df.x.dtype}'"
+    ):
+        df.shuffle("x", shuffle="p2p")
+
+
+def test_raise_on_custom_objects():
+    class Stub:
+        def __init__(self, value: int) -> None:
+            self.value = value
+
+    df = dd.from_pandas(
+        pd.DataFrame({"x": pd.array([Stub(i) for i in range(10)], dtype="object")}),
+        npartitions=5,
+    )
+    with pytest.raises(TypeError, match="p2p does not support custom objects"):
+        df.shuffle("x", shuffle="p2p")
+
+
+def test_raise_on_sparse_data():
+    df = dd.from_pandas(
+        pd.DataFrame({"x": pd.array(range(10), dtype="Sparse[float64]")}), npartitions=5
+    )
+    with pytest.raises(TypeError, match="p2p does not support sparse data"):
+        df.shuffle("x", shuffle="p2p")
+
+
 def test_raise_on_non_string_column_name():
     df = dd.from_pandas(pd.DataFrame({"a": range(10), 1: range(10)}), npartitions=5)
     with pytest.raises(TypeError, match="p2p requires all column names to be str"):
@@ -43,6 +75,7 @@ def test_does_not_raise_on_stringified_numeric_column_name():
 @gen_cluster([("", 2)] * 4, client=True)
 async def test_basic_state(c, s, *workers):
     df = dd.demo.make_timeseries(freq="15D", partition_freq="30D")
+    df["name"] = df["name"].astype("string[python]")
     shuffled = df.shuffle("id", shuffle="p2p")
 
     exts = [w.extensions["shuffle"] for w in workers]
diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py
index 54c62fb2..d3721d55 100644
--- a/distributed/tests/test_core.py
+++ b/distributed/tests/test_core.py
@@ -7,6 +7,7 @@ import socket
 import threading
 import time as timemod
 import weakref
+from unittest import mock
 
 import pytest
 from tornado.ioloop import IOLoop
@@ -724,10 +725,12 @@ async def test_connection_pool():
         *(rpc(ip="127.0.0.1", port=s.port).ping() for s in servers[:5])
     )
     await asyncio.gather(*(rpc(s.address).ping() for s in servers[:5]))
-    await asyncio.gather(*(rpc("127.0.0.1:%d" % s.port).ping() for s in servers[:5]))
+    await asyncio.gather(*(rpc(f"127.0.0.1:{s.port}").ping() for s in servers[:5]))
     await asyncio.gather(
         *(rpc(ip="127.0.0.1", port=s.port).ping() for s in servers[:5])
     )
+    await asyncio.gather(*(rpc(("127.0.0.1", s.port)).ping() for s in servers[:5]))
+
     assert sum(map(len, rpc.available.values())) == 5
     assert sum(map(len, rpc.occupied.values())) == 0
     assert rpc.active == 0
@@ -839,6 +842,34 @@ async def test_connection_pool_outside_cancellation(monkeypatch):
         assert all(t.cancelled() for t in tasks)
 
 
+@gen_test()
+async def test_remove_cancels_connect_attempts():
+    loop = asyncio.get_running_loop()
+    connect_started = asyncio.Event()
+    connect_finished = loop.create_future()
+
+    async def connect(*args, **kwargs):
+        connect_started.set()
+        await connect_finished
+
+    async def connect_to_server():
+        with pytest.raises(CommClosedError, match="Address removed."):
+            await rpc.connect("tcp://0.0.0.0")
+
+    async def remove_address():
+        await connect_started.wait()
+        rpc.remove("tcp://0.0.0.0")
+
+    rpc = await ConnectionPool(limit=1)
+
+    with mock.patch("distributed.core.connect", connect):
+        await asyncio.gather(
+            connect_to_server(),
+            remove_address(),
+        )
+    assert connect_finished.cancelled()
+
+
 @gen_test()
 async def test_connection_pool_respects_limit():
     limit = 5
diff --git a/distributed/tests/test_jupyter.py b/distributed/tests/test_jupyter.py
index ccd41ada..9ea2521d 100644
--- a/distributed/tests/test_jupyter.py
+++ b/distributed/tests/test_jupyter.py
@@ -2,12 +2,19 @@ from __future__ import annotations
 
 import pytest
 
+pytest.importorskip("requests")
+
+import requests
+
+from distributed import Client
+
 pytest.importorskip("jupyter_server")
 
 from tornado.httpclient import AsyncHTTPClient
 
 from distributed import Scheduler
-from distributed.utils_test import gen_test
+from distributed.utils import open_port
+from distributed.utils_test import gen_test, popen
 
 
 @gen_test()
@@ -18,3 +25,22 @@ async def test_jupyter_server():
             f"http://localhost:{s.http_server.port}/jupyter/api/status"
         )
         assert response.code == 200
+
+
+@pytest.mark.slow
+def test_jupyter_cli(loop):
+    port = open_port()
+    with popen(
+        [
+            "dask",
+            "scheduler",
+            "--jupyter",
+            "--no-dashboard",
+            "--host",
+            f"127.0.0.1:{port}",
+        ],
+        capture_output=True,
+    ):
+        with Client(f"127.0.0.1:{port}", loop=loop):
+            response = requests.get("http://127.0.0.1:8787/jupyter/api/status")
+            assert response.status_code == 200
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index 109f035d..5ff1234f 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -13,6 +13,7 @@ from itertools import product
 from textwrap import dedent
 from time import sleep
 from typing import Collection
+from unittest import mock
 
 import cloudpickle
 import psutil
@@ -2744,23 +2745,26 @@ class FlakyConnectionPool(ConnectionPool):
 
 
 @gen_cluster(client=True)
-async def test_gather_failing_cnn_recover(c, s, a, b):
-    orig_rpc = s.rpc
+async def test_gather_failing_cnn_recover(c, s, a, b, caplog):
     x = await c.scatter({"x": 1}, workers=a.address)
-
-    s.rpc = await FlakyConnectionPool(failing_connections=1)
-    with dask.config.set({"distributed.comm.retry.count": 1}):
+    rpc = await FlakyConnectionPool(failing_connections=1)
+    with mock.patch.object(s, "rpc", rpc), dask.config.set(
+        {"distributed.comm.retry.count": 1}
+    ), caplog.at_level(logging.INFO):
+        caplog.clear()
         res = await s.gather(keys=["x"])
+    assert [
+        record.message for record in caplog.records if record.levelno == logging.INFO
+    ] == ["Retrying get_data_from_worker after exception in attempt 0/1: "]
     assert res["status"] == "OK"
 
 
 @gen_cluster(client=True)
 async def test_gather_failing_cnn_error(c, s, a, b):
-    orig_rpc = s.rpc
     x = await c.scatter({"x": 1}, workers=a.address)
-
-    s.rpc = await FlakyConnectionPool(failing_connections=10)
-    res = await s.gather(keys=["x"])
+    rpc = await FlakyConnectionPool(failing_connections=10)
+    with mock.patch.object(s, "rpc", rpc):
+        res = await s.gather(keys=["x"])
     assert res["status"] == "error"
     assert list(res["keys"]) == ["x"]