[dask] Use distributed.MultiLock (#6743)

* [dask] Use `distributed.MultiLock`

This enables training multiple models in parallel.

* Conditionally import `MultiLock`.
* Use async train directly in scikit learn interface.
* Use `worker_client` when available.
This commit is contained in:
Jiaming Yuan
2021-03-16 14:19:41 +08:00
committed by GitHub
parent 19a2c54265
commit 325bc93e16
4 changed files with 212 additions and 55 deletions

View File

@@ -293,7 +293,7 @@ class TestDistributedGPU:
fw = fw - fw.min()
m = dxgb.DaskDMatrix(client, X, y, feature_weights=fw)
workers = list(_get_client_workers(client).keys())
workers = _get_client_workers(client)
rabit_args = client.sync(dxgb._get_rabit_args, len(workers), client)
def worker_fn(worker_addr: str, data_ref: Dict) -> None:
@@ -384,7 +384,7 @@ class TestDistributedGPU:
return subprocess.run([str(exe), test], env=env, stdout=subprocess.PIPE)
with Client(local_cuda_cluster) as client:
workers = list(_get_client_workers(client).keys())
workers = _get_client_workers(client)
rabit_args = client.sync(dxgb._get_rabit_args, workers, client)
futures = client.map(runit,
workers,