Skip to content

Dask worker removal causes crash after training completes #11804

Description

@adamreeve

I consistently get a crash after the final training iteration when using Dask. This is reproducible using the example code at https://xgboost.readthedocs.io/en/stable/python/dask-examples/gpu_training.html

When running this, I get:

Using DaskQuantileDMatrix
[20:47:31] Task [xgboost.dask-0]:tcp://127.0.0.1:39803 got rank 0
[20:47:31] Task [xgboost.dask-1]:tcp://127.0.0.1:41247 got rank 1
[20:47:33] [0]  train-rmse:2.84356
[20:47:33] [1]  train-rmse:2.79851
[20:47:33] [2]  train-rmse:2.75492
[20:47:33] [3]  train-rmse:2.71145
2025-11-11 20:47:35,620 - distributed.scheduler - WARNING - Removing worker 'tcp://127.0.0.1:39803' caused the cluster to lose already computed task(s), which will be recomputed elsewhere: {('uniform-from-dask-array-1df0a43faf0d656b0d4794e39ba42714', 4), 'dict-5ae98066-0603-498d-975e-c7a9fabec6fc', 'dict-4169c838-7a09-479a-b53a-3435de0a492a', 'dict-b6d54058-0c05-4645-8461-419017411166', 'dict-7d911c65-146a-4dc0-bff2-27444a132580', ('from-dask-array-1df0a43faf0d656b0d4794e39ba42714', 7), 'dict-5f48dca7-2b8e-46a8-858d-7f804f071678', 'dict-5ccb8d20-6cf0-4c16-87e3-a59cb23c51ad', 'dict-13825add-3269-45e2-b57b-7fb58bac9560', ('uniform-from-dask-array-1df0a43faf0d656b0d4794e39ba42714', 1), 'dict-881f8ed8-2aef-4d71-8966-c92beac004c1'} (stimulus_id='handle-worker-cleanup-1762894055.6203396')
2025-11-11 20:47:35,621 - distributed.scheduler - ERROR - Removing worker 'tcp://127.0.0.1:39803' caused the cluster to lose scattered data, which can't be recovered: {'Booster-317a20c46f8943fbac42423eeda8e498'} (stimulus_id='handle-worker-cleanup-1762894055.6203396')
Traceback (most recent call last):
  File "/home/adam/xgb-dask/main.py", line 83, in <module>
    from_ddqdm = using_quantile_device_dmatrix(client, X, y).compute()
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/dask/base.py", line 373, in compute
    (result,) = compute(self, traverse=False, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/dask_cuda/explicit_comms/dataframe/shuffle.py", line 665, in _patched_compute
    return _base_compute(
           ^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/dask/base.py", line 681, in compute
    results = schedule(expr, keys, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/distributed/client.py", line 2416, in _gather
    raise exception.with_traceback(traceback)
distributed.client.FutureCancelledError: ('repartitiontofewer-cc359db8dc9f6766b59d7812991e267c', 0) cancelled for reason: already forgotten.

I also get a similar but slightly different error on a larger training run:

Details
[21:04:47] Task [xgboost.dask-0]:tcp://127.0.0.1:40689 got rank 0
[21:04:47] Task [xgboost.dask-1]:tcp://127.0.0.1:34101 got rank 1
[21:05:21] [0]  dt-rmse:34901.35366
[21:05:21] [1]  dt-rmse:34873.19654
[21:05:22] [2]  dt-rmse:34842.21352
[21:05:22] [3]  dt-rmse:34808.43900
[21:05:22] [4]  dt-rmse:34777.91744
[21:05:22] [5]  dt-rmse:34749.22733
[21:05:23] [6]  dt-rmse:34715.53136
[21:05:23] [7]  dt-rmse:34684.79391
[21:05:23] [8]  dt-rmse:34651.51852
[21:05:23] [9]  dt-rmse:34617.87699
[21:05:24] [10] dt-rmse:34584.33838
2025-11-11 21:05:26,575 - distributed.scheduler - WARNING - Removing worker 'tcp://127.0.0.1:40689' caused the cluster to lose already computed task(s), which will be recomputed elsewhere: {'dict-c61bc932-9a43-45c7-8b34-c0a6717a70aa'} (stimulus_id='handle-worker-cleanup-1762895126.574773')
2025-11-11 21:05:26,758 - distributed.nanny - WARNING - Restarting worker
2025-11-11 21:05:27,656 - distributed.scheduler - WARNING - Removing worker 'tcp://127.0.0.1:34101' caused the cluster to lose already computed task(s), which will be recomputed elsewhere: {'dict-1cdb34f6-2459-46a8-a17e-adb372e774e9'} (stimulus_id='handle-worker-cleanup-1762895127.6557338')
2025-11-11 21:05:27,787 - distributed.nanny - WARNING - Restarting worker
2025-11-11 21:05:55,828 - distributed.worker - ERROR - Compute Failed
Key:       fn-009ab801-487a-4db6-97a9-5e937b0bb91b
State:     executing
Task:  <Task 'fn-009ab801-487a-4db6-97a9-5e937b0bb91b' fn(...)>
Exception: "ValueError('Invalid worker address: tcp://127.0.0.1:42895, expecting tcp://127.0.0.1:40689. This is likely caused by one of the workers died and Dask re-scheduled a different one. Resilience is not yet supported.')"
Traceback: '  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/xgboost/dask/__init__.py", line 536, in fn\n    raise ValueError(\n'

2025-11-11 21:05:55,847 - distributed.worker.state_machine - WARNING - Async instruction for <Task cancelled name="execute(('read_parquet-fused-2a3f1d7277bc5387a324b35f580bad87', 0))" coro=<Worker.execute() done, defined at /home/adam/xgb-dask/.venv/lib/python3.12/site-packages/distributed/worker_state_machine.py:3607>> ended with CancelledError
Traceback (most recent call last):
  File "/home/adam/xgb-dask/train_xgboost_dask.py", line 110, in <module>
    train(client)
  File "/home/adam/xgb-dask/train_xgboost_dask.py", line 61, in train
    output = dxgb.train(
             ^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/xgboost/core.py", line 774, in inner_f
    return func(**kwargs)
           ^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/xgboost/dask/__init__.py", line 876, in train
    return client.sync(
           ^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/distributed/utils.py", line 376, in sync
    return sync(
           ^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/distributed/utils.py", line 452, in sync
    raise error
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/distributed/utils.py", line 426, in f
    result = yield future
             ^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/tornado/gen.py", line 783, in run
    value = future.result()
            ^^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/xgboost/dask/__init__.py", line 810, in _train_async
    result = await map_worker_partitions(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/xgboost/dask/__init__.py", line 578, in map_worker_partitions
    result = await client.compute(fut).result()
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/distributed/client.py", line 404, in _result
    raise exc.with_traceback(tb)
  File "/home/adam/xgb-dask/.venv/lib/python3.12/site-packages/xgboost/dask/__init__.py", line 536, in fn
    raise ValueError(
ValueError: Invalid worker address: tcp://127.0.0.1:42895, expecting tcp://127.0.0.1:40689. This is likely caused by one of the workers died and Dask re-scheduled a different one. Resilience is not yet supported.

This is from an x86_64 machine running Ubuntu 24.04 with Nvidia H100 GPUs, CUDA 12.7, and Nvidia driver 565.57.01. The full set of dependencies from uv tree is:

Details
xgb-dask v0.1.0                                               
├── dask-cuda v25.10.0          
│   ├── click v8.3.0            
│   ├── cuda-core v0.3.2    
│   │   └── numpy v2.2.6     
│   ├── numba-cuda v0.19.1       
│   │   ├── numba v0.61.2              
│   │   │   ├── llvmlite v0.44.0    
│   │   │   └── numpy v2.2.6             
│   │   ├── cuda-bindings v12.9.4 (extra: cu12)
│   │   │   └── cuda-pathfinder v1.3.2
│   │   ├── cuda-core v0.3.2 (extra: cu12) (*)
│   │   ├── cuda-python v12.9.4 (extra: cu12)
│   │   │   └── cuda-bindings v12.9.4 (*)
│   │   ├── nvidia-cuda-cccl-cu12 v12.9.27 (extra: cu12)
│   │   ├── nvidia-cuda-nvcc-cu12 v12.9.86 (extra: cu12)
│   │   ├── nvidia-cuda-nvrtc-cu12 v12.9.86 (extra: cu12)
│   │   ├── nvidia-cuda-runtime-cu12 v12.9.79 (extra: cu12)
│   │   └── nvidia-nvjitlink-cu12 v12.9.86 (extra: cu12)
│   ├── numpy v2.2.6                     
│   ├── nvidia-ml-py v13.580.82
│   ├── pandas v2.3.3          
│   │   ├── numpy v2.2.6         
│   │   ├── python-dateutil v2.9.0.post0   
│   │   │   └── six v1.17.0                 
│   │   ├── pytz v2025.2        
│   │   └── tzdata v2025.2               
│   ├── rapids-dask-dependency v25.10.0
│   │   ├── dask v2025.9.1           
│   │   │   ├── click v8.3.0    
│   │   │   ├── cloudpickle v3.1.2
│   │   │   ├── fsspec v2025.10.0
│   │   │   ├── packaging v25.0      
│   │   │   ├── partd v1.4.2    
│   │   │   │   ├── locket v1.0.0
│   │   │   │   └── toolz v1.1.0
│   │   │   ├── pyyaml v6.0.3  
│   │   │   └── toolz v1.1.0
│   │   └── distributed v2025.9.1          
│   │       ├── click v8.3.0
│   │       ├── cloudpickle v3.1.2
│   │       ├── dask v2025.9.1 (*)
│   │       ├── jinja2 v3.1.6
│   │       │   └── markupsafe v3.0.3
│   │       ├── locket v1.0.0
│   │       ├── msgpack v1.1.2
│   │       ├── packaging v25.0
│   │       ├── psutil v7.1.3
│   │       ├── pyyaml v6.0.3
│   │       ├── sortedcontainers v2.4.0
│   │       ├── tblib v3.2.1
│   │       ├── toolz v1.1.0
│   │       ├── tornado v6.5.2
│   │       ├── urllib3 v2.5.0
│   │       └── zict v3.0.0
│   └── zict v3.0.0
├── dask-cudf-cu12 v25.10.0
│   ├── cudf-cu12 v25.10.0
│   │   ├── cachetools v6.2.1
│   │   ├── cuda-python v12.9.4 (*)
│   │   ├── cuda-toolkit[nvcc, nvrtc] v12.9.1
│   │   │   ├── nvidia-cuda-nvcc-cu12 v12.9.86 (extra: nvcc)
│   │   │   └── nvidia-cuda-nvrtc-cu12 v12.9.86 (extra: nvrtc)
│   │   ├── cupy-cuda12x v13.6.0
│   │   │   ├── fastrlock v0.8.3
│   │   │   └── numpy v2.2.6
│   │   ├── fsspec v2025.10.0
│   │   ├── libcudf-cu12 v25.10.0
│   │   │   ├── libkvikio-cu12 v25.10.0
│   │   │   ├── librmm-cu12 v25.10.0
│   │   │   │   └── rapids-logger v0.1.19
│   │   │   └── rapids-logger v0.1.19
│   │   ├── numba v0.61.2 (*)
│   │   ├── numba-cuda[cu12] v0.19.1 (*)
│   │   ├── numpy v2.2.6
│   │   ├── nvtx v0.2.13
│   │   ├── packaging v25.0
│   │   ├── pandas v2.3.3 (*)
│   │   ├── pyarrow v22.0.0
│   │   ├── pylibcudf-cu12 v25.10.0
│   │   │   ├── cuda-python v12.9.4 (*)
│   │   │   ├── libcudf-cu12 v25.10.0 (*)
│   │   │   ├── nvtx v0.2.13
│   │   │   ├── packaging v25.0
│   │   │   ├── rmm-cu12 v25.10.0
│   │   │   │   ├── cuda-python v12.9.4 (*)
│   │   │   │   ├── librmm-cu12 v25.10.0 (*)
│   │   │   │   └── numpy v2.2.6
│   │   │   └── typing-extensions v4.15.0
│   │   ├── rich v14.2.0
│   │   │   ├── markdown-it-py v4.0.0
│   │   │   │   └── mdurl v0.1.2
│   │   │   └── pygments v2.19.2
│   │   ├── rmm-cu12 v25.10.0 (*)
│   │   └── typing-extensions v4.15.0
│   ├── cupy-cuda12x v13.6.0 (*)
│   ├── fsspec v2025.10.0
│   ├── numpy v2.2.6
│   ├── nvidia-ml-py v13.580.82
│   ├── pandas v2.3.3 (*)
│   └── rapids-dask-dependency v25.10.0 (*)
└── xgboost v3.1.1
    ├── numpy v2.2.6
    ├── nvidia-nccl-cu12 v2.28.7
    └── scipy v1.16.3
        └── numpy v2.2.6

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions