Dask 2025.4.0 scheduler info compatibility

[jrbourbeau]

Closes #11461 per @trivialfis suggestion here #11461 (comment)

[jrbourbeau]

@trivialfis this test

xgboost/tests/python/test_tracker.py

Lines 160 to 181 in 9ad4e24

	@pytest.mark.skipif(**tm.no_dask())
	def test_rank_assignment() -> None:
	from distributed import Client, LocalCluster
	from xgboost import dask as dxgb
	from xgboost.testing.dask import get_rabit_args
	def local_test(worker_id):
	with dxgb.CommunicatorContext(**args) as ctx:
	task_id = ctx["DMLC_TASK_ID"]
	matched = re.search(".*-([0-9]).*", task_id)
	rank = collective.get_rank()
	# As long as the number of workers is lesser than 10, rank and worker id
	# should be the same
	assert rank == int(matched.group(1))
	with LocalCluster(n_workers=8) as cluster:
	with Client(cluster) as client:
	workers = tm.dask.get_client_workers(client)
	args = get_rabit_args(client, len(workers))
	futures = client.map(local_test, range(len(workers)), workers=workers)
	client.gather(futures)

is explicitly checking the previous worker id is in the task id used for determining the rank. Do you have a suggestion on how to update this test? Or should it be removed?

[trivialfis]

Do you have a suggestion on how to update this test? Or should it be removed?

Based on #11461 (comment) , can we use worker.name as a replacement for the current wid?
This way, the test should pass since it uses a local cluster, which has an integer name by default. In addition, XGBoost can continue to sort the workers by ID.

[jrbourbeau]

Based on #11461 (comment) , can we use worker.name as a replacement for the current wid?

Yeah, fair point. I just want to reiterate whether the current wid (and the equivalent worker.name) are integers like 0, 1, 2, 3, ... depends on how the workers were created. distributed.LocalCluster sets them to be integers but other deployments can set those names to be strings (or any other hashable object).

That said, that's a separate issue. I've gone ahead and updated this PR to avoid the scheduler_info() call altogether by using worker.name for the current wid. That should keep the current behavior.

[jrbourbeau]

Hmm the build failure here looks unrelated (probably just some transient HTTP error). I'll push an empty commit to rerun CI to see if that helps.

[trivialfis]

Please ignore the canceled error. I don't know how it happens, github action is dark magic to me.

On the other hand, we have seen this https://github.com/dmlc/xgboost/actions/runs/15050046963/job/42303211874?pr=11462 a couple of times before. Would you like to help make a guess? (not expecting you to spend time debugging it). It's a flaky error that looks like dask is having trouble with parameter packs:

     |   File "/workspace/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py", line 194, in run_gpu_hist
    |     history = dxgb.train(
    |   File "/home/runner/.local/lib/python3.10/site-packages/xgboost/core.py", line 772, in inner_f
    |     return func(**kwargs)
    |   File "/home/runner/.local/lib/python3.10/site-packages/xgboost/dask/__init__.py", line 923, in train
    |     return client.sync(
    |   File "/opt/miniforge/envs/gpu_test/lib/python3.10/site-packages/distributed/utils.py", line 363, in sync
    |     return sync(
    |   File "/opt/miniforge/envs/gpu_test/lib/python3.10/site-packages/distributed/utils.py", line 439, in sync
    |     raise error
    |   File "/opt/miniforge/envs/gpu_test/lib/python3.10/site-packages/distributed/utils.py", line 409, in f
    |     awaitable = func(*args, **kwargs)
    | TypeError: _train_async() got an unexpected keyword argument 'args'
    +------------------------------------

[jrbourbeau]

Would you like to help make a guess? (not expecting you to spend time debugging it). It's a flaky error that looks like dask is having trouble with parameter packs:

Hmm yeah that's a weird one. My guess is that it has something to do with args = locals() here

xgboost/python-package/xgboost/dask/__init__.py

Line 926 in 9ad4e24

args = locals()

somehow including args in the returned dictionary. To be clear, I'm not sure how that would happen, but, much like GitHub actions, locals() can sometimes seem like dark magic. A change like this

diff --git a/python-package/xgboost/dask/__init__.py b/python-package/xgboost/dask/__init__.py
index 7f38647fa..e87e85126 100644
--- a/python-package/xgboost/dask/__init__.py
+++ b/python-package/xgboost/dask/__init__.py
@@ -923,12 +923,11 @@ def train(  # pylint: disable=unused-argument

     """
     client = _get_client(client)
-    args = locals()
     return client.sync(
         _train_async,
         global_config=config.get_config(),
         dconfig=_get_dask_config(),
-        **args,
+        **locals(),
     )

might? fix the issue. At a minimum, it seems like a harmless change (and it's how locals() is used in other places in this module)

[jrbourbeau]

Thanks @trivialfis

[jrbourbeau]

Btw, do you have a sense for when this change might be included in a release? I see it just missed the 3.0.1 patch release unfortunately

[trivialfis]

We can do another one if necessary, cc @hcho3 .

[jrbourbeau]

I'm not sure how much of a burden it is to release, but I'd certainly appreciate a release with this update as soon as makes sense for you all (most xgboost.dask usage that I'm aware of in the wild runs on clusters with more than 5 workers, so would benefit from this update).