diff --git a/python-package/xgboost/dask/__init__.py b/python-package/xgboost/dask/__init__.py index b1400bdc5..7a565f5f2 100644 --- a/python-package/xgboost/dask/__init__.py +++ b/python-package/xgboost/dask/__init__.py @@ -1634,7 +1634,7 @@ class DaskScikitLearnBase(XGBModel): if isinstance(predts, dd.DataFrame): predts = predts.to_dask_array() else: - test_dmatrix = await DaskDMatrix( + test_dmatrix: DaskDMatrix = await DaskDMatrix( # type: ignore self.client, data=data, base_margin=base_margin, @@ -1675,7 +1675,7 @@ class DaskScikitLearnBase(XGBModel): iteration_range: Optional[IterationRange] = None, ) -> Any: iteration_range = self._get_iteration_range(iteration_range) - test_dmatrix = await DaskDMatrix( + test_dmatrix: DaskDMatrix = await DaskDMatrix( # type: ignore self.client, data=X, missing=self.missing, diff --git a/tests/ci_build/conda_env/python_lint.yml b/tests/ci_build/conda_env/python_lint.yml index 10a313baa..579c377ef 100644 --- a/tests/ci_build/conda_env/python_lint.yml +++ b/tests/ci_build/conda_env/python_lint.yml @@ -6,8 +6,7 @@ dependencies: - pylint - wheel - setuptools -# https://github.com/python/mypy/issues/17617 -- mypy !=1.11.0,!=1.11.1 +- mypy - numpy - scipy - pandas 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 8ab661a24..c729761b8 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 @@ -647,10 +647,10 @@ async def run_from_dask_array_asyncio(scheduler_address: str) -> dxgb.TrainRetur import cupy as cp X, y, _ = generate_array() - X = X.map_blocks(cp.array) # type: ignore - y = y.map_blocks(cp.array) # type: ignore + X = X.to_backend("cupy") + y = y.to_backend("cupy") - m = await xgb.dask.DaskQuantileDMatrix(client, X, y) + m: xgb.dask.DaskDMatrix = await xgb.dask.DaskQuantileDMatrix(client, X, y) # type: ignore output = await xgb.dask.train( client, {"tree_method": "hist", "device": "cuda"}, dtrain=m )