Properly await async method client.wait_for_workers (dmlc#8558) (dmlc#8567)

trivialfis · mrocklin · web-flow · commit a980e10744e5 · 2022-12-07T23:25:05.000+08:00
* Properly await async method client.wait_for_workers

* ignore mypy error.

Co-authored-by: jiamingy &lt;jm.yuan@outlook.com&gt;

Co-authored-by: Matthew Rocklin &lt;mrocklin@gmail.com&gt;
diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
@@ -853,7 +853,7 @@ async def _get_rabit_args(
         sched_addr = None
 
     # make sure all workers are online so that we can obtain reliable scheduler_info
-    client.wait_for_workers(n_workers)
+    await client.wait_for_workers(n_workers)  # type: ignore
     env = await client.run_on_scheduler(
         _start_tracker, n_workers, sched_addr, user_addr
     )

Original file line number	Diff line number	Diff line change
`@@ -853,7 +853,7 @@ async def _get_rabit_args(`
`853`	`853`	`sched_addr = None`
`854`	`854`
`855`	`855`	`# make sure all workers are online so that we can obtain reliable scheduler_info`
`856`		`- client.wait_for_workers(n_workers)`
	`856`	`+ await client.wait_for_workers(n_workers) # type: ignore`
`857`	`857`	`env = await client.run_on_scheduler(`
`858`	`858`	`_start_tracker, n_workers, sched_addr, user_addr`
`859`	`859`	`)`