[dask] Supoort running on GKE. (#6343)

* Avoid accessing `scheduler_info()['workers']`.
* Avoid calling `client.gather` inside task.
* Avoid using `client.scheduler_address`.
This commit is contained in:
Jiaming Yuan
2020-11-11 18:04:34 +08:00
committed by GitHub
parent 8a17610666
commit 6e12c2a6f8
4 changed files with 150 additions and 128 deletions

View File

@@ -23,11 +23,12 @@ def test_rabit_tracker():
def run_rabit_ops(client, n_workers):
from xgboost.dask import RabitContext, _get_rabit_args, _get_client_workers
from test_with_dask import _get_client_workers
from xgboost.dask import RabitContext, _get_rabit_args
from xgboost import rabit
workers = list(_get_client_workers(client).keys())
rabit_args = client.sync(_get_rabit_args, workers, client)
rabit_args = client.sync(_get_rabit_args, len(workers), client)
assert not rabit.is_distributed()
n_workers_from_dask = len(workers)
assert n_workers == n_workers_from_dask