Avoid stale CommContext in explicit comms (#1451)

This PR updates the CommContext caching to be keyed by some information about the cluster, rather than a single global. This prevents us from using a stale comms object after the cluster changes (add or remove workers) or is recreated entirely.

Closes #1450

Authors:
  - Tom Augspurger (https://github.com/TomAugspurger)

Approvers:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

URL: #1451

dask_cuda/explicit_comms/comms.py

@@ -1,15 +1,21 @@
+# Copyright (c) 2021-2025 NVIDIA CORPORATION.
 import asyncio
 import concurrent.futures
 import contextlib
 import time
 import uuid
+import weakref
 from typing import Any, Dict, Hashable, Iterable, List, Optional
 
 import distributed.comm
+from dask.tokenize import tokenize
 from distributed import Client, Worker, default_client, get_worker
 from distributed.comm.addressing import parse_address, parse_host_port, unparse_address
 
-_default_comms = None
+# Mapping tokenize(client ID, [worker addresses]) to CommsContext
+_comms_cache: weakref.WeakValueDictionary[
+    str, "CommsContext"
+] = weakref.WeakValueDictionary()
 
 
 def get_multi_lock_or_null_context(multi_lock_context, *args, **kwargs):
@@ -38,9 +44,10 @@ def get_multi_lock_or_null_context(multi_lock_context, *args, **kwargs):
 
 
 def default_comms(client: Optional[Client] = None) -> "CommsContext":
-    """Return the default comms object
+    """Return the default comms object for ``client``.
 
-    Creates a new default comms object if no one exist.
+    Creates a new default comms object if one does not already exist
+    for ``client``.
 
     Parameters
     ----------
@@ -52,11 +59,31 @@ def default_comms(client: Optional[Client] = None) -> "CommsContext":
     -------
     comms: CommsContext
         The default comms object
+
+    Notes
+    -----
+    There are some subtle points around explicit-comms and the lifecycle
+    of a Dask Cluster.
+
+    A :class:`CommsContext` establishes explicit communication channels
+    between the workers *at the time it's created*. If workers are added
+    or removed, they will not be included in the communication channels
+    with the other workers.
+
+    If you need to refresh the explicit communications channels, then
+    create a new :class:`CommsContext` object or call ``default_comms``
+    again after workers have been added to or removed from the cluster.
     """
-    global _default_comms
-    if _default_comms is None:
-        _default_comms = CommsContext(client=client)
-    return _default_comms
+    # Comms are unique to a {client, [workers]} pair, so we key our
+    # cache by the token of that.
+    client = client or default_client()
+    token = tokenize(client.id, list(client.scheduler_info()["workers"].keys()))
+    maybe_comms = _comms_cache.get(token)
+    if maybe_comms is None:
+        maybe_comms = CommsContext(client=client)
+        _comms_cache[token] = maybe_comms
+
+    return maybe_comms
 
 
 def worker_state(sessionId: Optional[int] = None) -> dict:

dask_cuda/tests/test_explicit_comms.py

@@ -1,3 +1,5 @@
+# Copyright (c) 2021-2025 NVIDIA CORPORATION.
+
 import asyncio
 import multiprocessing as mp
 import os
@@ -415,3 +417,116 @@ def test_lock_workers():
             p.join()
 
         assert all(p.exitcode == 0 for p in ps)
+
+
+def test_create_destroy_create():
+    # https://github.com/rapidsai/dask-cuda/issues/1450
+    assert len(comms._comms_cache) == 0
+    with LocalCluster(n_workers=1) as cluster:
+        with Client(cluster) as client:
+            context = comms.default_comms()
+            scheduler_addresses_old = list(client.scheduler_info()["workers"].keys())
+            comms_addresses_old = list(comms.default_comms().worker_addresses)
+            assert comms.default_comms() is context
+            assert len(comms._comms_cache) == 1
+
+            # Add a worker, which should have a new comms object
+            cluster.scale(2)
+            client.wait_for_workers(2, timeout=5)
+            context2 = comms.default_comms()
+            assert context is not context2
+            assert len(comms._comms_cache) == 2
+
+    del context
+    del context2
+    assert len(comms._comms_cache) == 0
+    assert scheduler_addresses_old == comms_addresses_old
+
+    # A new cluster should have a new comms object. Previously, this failed
+    # because we referenced the old cluster's addresses.
+    with LocalCluster(n_workers=1) as cluster:
+        with Client(cluster) as client:
+            scheduler_addresses_new = list(client.scheduler_info()["workers"].keys())
+            comms_addresses_new = list(comms.default_comms().worker_addresses)
+
+    assert scheduler_addresses_new == comms_addresses_new
+
+
+def test_scaled_cluster_gets_new_comms_context():
+    # Ensure that if we create a CommsContext, scale the cluster,
+    # and create a new CommsContext, then the new CommsContext
+    # should include the new worker.
+    # https://github.com/rapidsai/dask-cuda/issues/1450
+
+    name = "explicit-comms-shuffle"
+    ddf = dd.from_pandas(pd.DataFrame({"key": np.arange(10)}), npartitions=2)
+
+    with LocalCluster(n_workers=2) as cluster:
+        with Client(cluster) as client:
+            context_1 = comms.default_comms()
+
+            def check(dask_worker, session_id: int):
+                has_state = hasattr(dask_worker, "_explicit_comm_state")
+                has_state_for_session = (
+                    has_state and session_id in dask_worker._explicit_comm_state
+                )
+                if has_state_for_session:
+                    n_workers = dask_worker._explicit_comm_state[session_id]["nworkers"]
+                else:
+                    n_workers = None
+                return {
+                    "has_state": has_state,
+                    "has_state_for_session": has_state_for_session,
+                    "n_workers": n_workers,
+                }
+
+            result_1 = client.run(check, session_id=context_1.sessionId)
+            expected_values = {
+                "has_state": True,
+                "has_state_for_session": True,
+                "n_workers": 2,
+            }
+            expected_1 = {
+                k: expected_values for k in client.scheduler_info()["workers"]
+            }
+            assert result_1 == expected_1
+
+            # Run a shuffle with the initial setup as a sanity test
+            with dask.config.set(explicit_comms=True):
+                shuffled = ddf.shuffle(on="key", npartitions=4)
+                assert any(name in str(key) for key in shuffled.dask)
+                result = shuffled.compute()
+
+            with dask.config.set(explicit_comms=False):
+                shuffled = ddf.shuffle(on="key", npartitions=4)
+                expected = shuffled.compute()
+
+            assert_eq(result, expected)
+
+            # --- Scale the cluster ---
+            cluster.scale(3)
+            client.wait_for_workers(3, timeout=5)
+
+            context_2 = comms.default_comms()
+            result_2 = client.run(check, session_id=context_2.sessionId)
+            expected_values = {
+                "has_state": True,
+                "has_state_for_session": True,
+                "n_workers": 3,
+            }
+            expected_2 = {
+                k: expected_values for k in client.scheduler_info()["workers"]
+            }
+            assert result_2 == expected_2
+
+            # Run a shuffle with the new setup
+            with dask.config.set(explicit_comms=True):
+                shuffled = ddf.shuffle(on="key", npartitions=4)
+                assert any(name in str(key) for key in shuffled.dask)
+                result = shuffled.compute()
+
+            with dask.config.set(explicit_comms=False):
+                shuffled = ddf.shuffle(on="key", npartitions=4)
+                expected = shuffled.compute()
+
+            assert_eq(result, expected)

docs/source/explicit_comms.rst

@@ -2,7 +2,7 @@ Explicit-comms
 ==============
 
 Communication and scheduling overhead can be a major bottleneck in Dask/Distributed. Dask-CUDA addresses this by introducing an API for explicit communication in Dask tasks.
-The idea is that Dask/Distributed spawns workers and distribute data as usually while the user can submit tasks on the workers that communicate explicitly.
+The idea is that Dask/Distributed spawns workers and distribute data as usual while the user can submit tasks on the workers that communicate explicitly.
 
 This makes it possible to bypass Distributed's scheduler and write hand-tuned computation and communication patterns. Currently, Dask-CUDA includes an explicit-comms
 implementation of the Dataframe `shuffle <../api/#dask_cuda.explicit_comms.dataframe.shuffle.shuffle>`_ operation used for merging and sorting.