From c7339a76c53491e306f7b5c7f005aec75c16f8e1 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Thu, 6 Feb 2025 17:05:34 +0800 Subject: [PATCH 1/2] Fix with the latest mypy. --- python-package/xgboost/dask/__init__.py | 14 +++++--------- .../test_gpu_with_dask/test_gpu_with_dask.py | 4 +--- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/python-package/xgboost/dask/__init__.py b/python-package/xgboost/dask/__init__.py index 635bedc7d2f5..8ae7fe1fe0a2 100644 --- a/python-package/xgboost/dask/__init__.py +++ b/python-package/xgboost/dask/__init__.py @@ -536,19 +536,16 @@ async def map_worker_partitions( else: args.append(ref) - @wraps(func) - def fn( - *args: _P.args, address: str = addr, **kwargs: _P.kwargs - ) -> List[_MapRetT]: - # Turn result into a list for bag construction + def fn(_address: str, *args: _P.args, **kwargs: _P.kwargs) -> List[_MapRetT]: worker = distributed.get_worker() - if worker.address != address: + if worker.address != _address: raise ValueError( - f"Invalid worker address: {worker.address}, expecting {address}. " + f"Invalid worker address: {worker.address}, expecting {_address}. " "This is likely caused by one of the workers died and Dask " "re-scheduled a different one. Resilience is not yet supported." ) + # Turn result into a list for bag construction return [func(*args, **kwargs)] # XGBoost requires all workers running training tasks to be unique. Meaning, we @@ -568,9 +565,8 @@ def fn( # relax the constraint and prevent Dask from choosing an invalid worker, the # task will simply hangs. We prefer a quick error here. # - fut = client.submit( - fn, + update_wrapper(partial(fn, addr), fn), *args, pure=False, workers=[addr], diff --git a/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py b/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py index 0aa3958e5868..0cae314c685e 100644 --- a/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py +++ b/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py @@ -17,9 +17,7 @@ from xgboost.testing.dask import get_rabit_args from xgboost.testing.params import hist_parameter_strategy -from ..test_with_dask.test_with_dask import ( - generate_array, -) +from ..test_with_dask.test_with_dask import generate_array from ..test_with_dask.test_with_dask import kCols as random_cols from ..test_with_dask.test_with_dask import ( make_categorical, From 541991eb34d61ef282b790c5b439a4fff791472f Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Thu, 6 Feb 2025 23:26:15 +0800 Subject: [PATCH 2/2] conflict with isort and black. --- python-package/xgboost/spark/summary.py | 2 +- .../test_distributed/test_gpu_with_dask/test_gpu_with_dask.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/python-package/xgboost/spark/summary.py b/python-package/xgboost/spark/summary.py index eca5f6b128b7..9ee720739611 100644 --- a/python-package/xgboost/spark/summary.py +++ b/python-package/xgboost/spark/summary.py @@ -16,7 +16,7 @@ class XGBoostTrainingSummary: @staticmethod def from_metrics( - metrics: Dict[str, Dict[str, List[float]]] + metrics: Dict[str, Dict[str, List[float]]], ) -> "XGBoostTrainingSummary": """ Create an XGBoostTrainingSummary instance from a nested dictionary of metrics. diff --git a/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py b/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py index 0cae314c685e..0aa3958e5868 100644 --- a/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py +++ b/tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py @@ -17,7 +17,9 @@ from xgboost.testing.dask import get_rabit_args from xgboost.testing.params import hist_parameter_strategy -from ..test_with_dask.test_with_dask import generate_array +from ..test_with_dask.test_with_dask import ( + generate_array, +) from ..test_with_dask.test_with_dask import kCols as random_cols from ..test_with_dask.test_with_dask import ( make_categorical,