Add typehint to tracker. (#7338)

dmlc · Oct 20, 2021 · f53da41 · f53da41
1 parent 5ff210e
commit f53da41
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 94 deletions.
diff --git a/Makefile b/Makefile
@@ -93,11 +93,14 @@ mypy:
 	cd python-package; \
 	mypy ./xgboost/dask.py && \
 	mypy ./xgboost/rabit.py && \
+	mypy ./xgboost/tracker.py && \
+	mypy ./xgboost/sklearn.py && \
 	mypy ../demo/guide-python/external_memory.py && \
+	mypy ../demo/guide-python/categorical.py && \
+	mypy ../demo/guide-python/cat_in_the_dat.py && \
 	mypy ../tests/python-gpu/test_gpu_with_dask.py && \
 	mypy ../tests/python/test_data_iterator.py && \
-	mypy ../tests/python-gpu/test_gpu_data_iterator.py && \
-	mypy ./xgboost/sklearn.py || exit 1; \
+	mypy ../tests/python-gpu/test_gpu_data_iterator.py  || exit 1; \
 	mypy . || true ;
 
 clean:

diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
@@ -138,10 +138,10 @@ async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
 
 def _start_tracker(n_workers: int) -> Dict[str, Any]:
     """Start Rabit tracker """
-    env = {'DMLC_NUM_WORKER': n_workers}
+    env: Dict[str, Union[int, str]] = {'DMLC_NUM_WORKER': n_workers}
     host = get_host_ip('auto')
-    rabit_context = RabitTracker(hostIP=host, nslave=n_workers, use_logger=False)
-    env.update(rabit_context.slave_envs())
+    rabit_context = RabitTracker(hostIP=host, n_workers=n_workers, use_logger=False)
+    env.update(rabit_context.worker_envs())
 
     rabit_context.start(n_workers)
     thread = Thread(target=rabit_context.join)