Avoid stale CommContext in explicit comms

[TomAugspurger]

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

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[rjzamora]

This is really nice @TomAugspurger - Thanks!

I left a few minor comments. One additional question: Do you think we need to test an explicit-comms shuffle after the comms context has changed? These tests show that the comms context can be refreshed after we scale up or down workers, but I suppose we could also make sure that the refreshed context "works" as expected?

[TomAugspurger]

Do you think we need to test an explicit-comms shuffle after the comms context has changed

Yeah, that would be good. I'll add that.

[TomAugspurger]

The CI failure appears to be an unrelated timeout: https://github.com/rapidsai/dask-cuda/actions/runs/13332958081/job/37241713827?pr=1451#step:10:5333

But I'm going to spend a bit of time to try to understand what's going on.

[rjzamora]

But I'm going to spend a bit of time to try to understand what's going on.

I certainly appreciate that. It's also fine to try rerunning that check given that a Timeout error like that is probably unrelated to the PR (and spilling tests are known to be flaky in CI :/ ).

[TomAugspurger]

I didn't figure out much. This test logs stuff at various points in time:

1. cluster startup
2. Unmanaged memory warnings (I think from imports)
3. worker_assert (x2)
4. assert_host_chunks
5. assert_disk_chunks

Success logs from a run yesterday.

Failure logs from this run.

This table shows very roughly how long some stages took:

stage	failure	success
unmanaged memory	baseline	baseline
worker_assert	+6s	+1s
worker_assert2	+8	+0.5s
assert_host_chunks	+0	+0
assert_disk_chunks	+0	+0

So for whatever reason, the worker_asserts took a while on the failed run. That does involve an RPC between the client and the worker, but most of the time should be spent in device_host_file_size_matches. I'd need to look a bit more closely about where stuff could be taking time there (it does interact with the filesystem). Probably not today, so I'll restart that one failure.

[rjzamora]

LGTM - Thanks @TomAugspurger !

[TomAugspurger]

/merge