Fix mypy errors. (#11389)

Author: trivialfis

python-package/xgboost/dask/__init__.py

@@ -55,7 +55,7 @@
 import logging
 from collections import defaultdict
 from contextlib import contextmanager
-from functools import partial, update_wrapper, wraps
+from functools import partial, update_wrapper
 from threading import Thread
 from typing import (
     Any,
@@ -354,7 +354,7 @@ def __init__(
             label_upper_bound=label_upper_bound,
         )
 
-    def __await__(self) -> Generator:
+    def __await__(self) -> Generator[None, None, "DaskDMatrix"]:
         return self._init.__await__()
 
     async def _map_local_data(
@@ -1490,7 +1490,7 @@ async def _predict_async(
             if isinstance(predts, dd.DataFrame):
                 predts = predts.to_dask_array()
         else:
-            test_dmatrix: DaskDMatrix = await DaskDMatrix(  # type: ignore
+            test_dmatrix: DaskDMatrix = await DaskDMatrix(
                 self.client,
                 data=data,
                 base_margin=base_margin,
@@ -1532,7 +1532,7 @@ async def _apply_async(
         iteration_range: Optional[IterationRange] = None,
     ) -> Any:
         iteration_range = self._get_iteration_range(iteration_range)
-        test_dmatrix: DaskDMatrix = await DaskDMatrix(  # type: ignore
+        test_dmatrix: DaskDMatrix = await DaskDMatrix(
             self.client,
             data=X,
             missing=self.missing,

tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py

@@ -670,9 +670,7 @@ async def run_from_dask_array_asyncio(scheduler_address: str) -> dxgb.TrainRetur
         X = X.to_backend("cupy")
         y = y.to_backend("cupy")
 
-        m: dxgb.DaskDMatrix = await dxgb.DaskQuantileDMatrix(
-            client, X, y
-        )  # type:ignore
+        m: dxgb.DaskDMatrix = await dxgb.DaskQuantileDMatrix(client, X, y)
         output = await dxgb.train(
             client, {"tree_method": "hist", "device": "cuda"}, dtrain=m
         )

tests/test_distributed/test_with_dask/test_with_dask.py

@@ -953,7 +953,7 @@ def test_empty_dmatrix(tree_method) -> None:
 async def run_from_dask_array_asyncio(scheduler_address: str) -> dxgb.TrainReturnT:
     async with Client(scheduler_address, asynchronous=True) as client:
         X, y, _ = generate_array()
-        m = await DaskDMatrix(client, X, y)  # type: ignore
+        m = await DaskDMatrix(client, X, y)
         output = await dxgb.train(client, {}, dtrain=m)
 
         with_m = await dxgb.predict(client, output, m)
@@ -1058,8 +1058,8 @@ async def train() -> None:
         ) as cluster:
             async with Client(cluster, asynchronous=True) as client:
                 X, y, w = generate_array(with_weights=True)
-                dtrain = await DaskDMatrix(client, X, y, weight=w)  # type: ignore
-                dvalid = await DaskDMatrix(client, X, y, weight=w)  # type: ignore
+                dtrain = await DaskDMatrix(client, X, y, weight=w)
+                dvalid = await DaskDMatrix(client, X, y, weight=w)
                 output = await dxgb.train(client, {}, dtrain=dtrain)
                 await dxgb.predict(client, output, data=dvalid)
 
@@ -2195,7 +2195,7 @@ async def test_worker_left(c: Client, s: Scheduler, a: Worker, b: Worker):
     async with Worker(s.address):
         dx = da.random.random((1000, 10)).rechunk(chunks=(10, None))
         dy = da.random.random((1000,)).rechunk(chunks=(10,))
-        d_train = await dxgb.DaskDMatrix(  # type: ignore
+        d_train = await dxgb.DaskDMatrix(
             c,
             dx,
             dy,