dask.py

# pylint: disable=too-many-arguments, too-many-locals, no-name-in-module
# pylint: disable=missing-class-docstring, invalid-name
# pylint: disable=too-many-lines
# pylint: disable=import-error
"""Dask extensions for distributed training. See
https://xgboost.readthedocs.io/en/latest/tutorials/dask.html for simple
tutorial.  Also xgboost/demo/dask for some examples.

There are two sets of APIs in this module, one is the functional API including
``train`` and ``predict`` methods.  Another is stateful Scikit-Learner wrapper
inherited from single-node Scikit-Learn interface.

The implementation is heavily influenced by dask_xgboost:
https://github.com/dask/dask-xgboost

"""
import platform
import logging
from collections import defaultdict
from collections.abc import Sequence
from threading import Thread
from typing import TYPE_CHECKING, List, Tuple, Callable, Optional, Any, Union, Dict, Set
from typing import Awaitable, Generator, TypeVar

import numpy

from . import rabit, config

from .callback import TrainingCallback

from .compat import LazyLoader
from .compat import sparse, scipy_sparse
from .compat import PANDAS_INSTALLED, DataFrame, Series, pandas_concat
from .compat import lazy_isinstance

from .core import DMatrix, DeviceQuantileDMatrix, Booster, _expect, DataIter
from .core import Objective, Metric
from .core import _deprecate_positional_args
from .training import train as worker_train
from .tracker import RabitTracker, get_host_ip
from .sklearn import XGBModel, XGBRegressorBase, XGBClassifierBase, _objective_decorator
from .sklearn import xgboost_model_doc


if TYPE_CHECKING:
    from dask import dataframe as dd
    from dask import array as da
    import dask
    import distributed
else:
    dd = LazyLoader('dd', globals(), 'dask.dataframe')
    da = LazyLoader('da', globals(), 'dask.array')
    dask = LazyLoader('dask', globals(), 'dask')
    distributed = LazyLoader('distributed', globals(), 'dask.distributed')

_DaskCollection = Union["da.Array", "dd.DataFrame", "dd.Series"]

try:
    from mypy_extensions import TypedDict
    TrainReturnT = TypedDict('TrainReturnT', {
        'booster': Booster,
        'history': Dict,
    })
except ImportError:
    TrainReturnT = Dict[str, Any]  # type:ignore

# Current status is considered as initial support, many features are not properly
# supported yet.
#
# TODOs:
#   - CV
#   - Ranking
#
# Note for developers:
#
#   As of writing asyncio is still a new feature of Python and in depth documentation is
#   rare.  Best examples of various asyncio tricks are in dask (luckily).  Classes like
#   Client, Worker are awaitable.  Some general rules for the implementation here:
#
#     - Synchronous world is different from asynchronous one, and they don't mix well.
#     - Write everything with async, then use distributed Client sync function to do the
#       switch.
#     - Use Any for type hint when the return value can be union of Awaitable and plain
#       value.  This is caused by Client.sync can return both types depending on context.
#       Right now there's no good way to silent:
#
#         await train(...)
#
#       if train returns an Union type.


LOGGER = logging.getLogger('[xgboost.dask]')


def _start_tracker(n_workers: int) -> Dict[str, Any]:
    """Start Rabit tracker """
    env = {'DMLC_NUM_WORKER': n_workers}
    host = get_host_ip('auto')
    rabit_context = RabitTracker(hostIP=host, nslave=n_workers)
    env.update(rabit_context.slave_envs())

    rabit_context.start(n_workers)
    thread = Thread(target=rabit_context.join)
    thread.daemon = True
    thread.start()
    return env


def _assert_dask_support() -> None:
    try:
        import dask             # pylint: disable=W0621,W0611
    except ImportError as e:
        raise ImportError(
            'Dask needs to be installed in order to use this module') from e

    if platform.system() == 'Windows':
        msg = 'Windows is not officially supported for dask/xgboost,'
        msg += ' contribution are welcomed.'
        LOGGER.warning(msg)


class RabitContext:
    '''A context controling rabit initialization and finalization.'''
    def __init__(self, args: List[bytes]) -> None:
        self.args = args
        worker = distributed.get_worker()
        self.args.append(
            ('DMLC_TASK_ID=[xgboost.dask]:' + str(worker.address)).encode())

    def __enter__(self) -> None:
        rabit.init(self.args)
        LOGGER.debug('-------------- rabit say hello ------------------')

    def __exit__(self, *args: List) -> None:
        rabit.finalize()
        LOGGER.debug('--------------- rabit say bye ------------------')


def concat(value: Any) -> Any:  # pylint: disable=too-many-return-statements
    '''To be replaced with dask builtin.'''
    if isinstance(value[0], numpy.ndarray):
        return numpy.concatenate(value, axis=0)
    if scipy_sparse and isinstance(value[0], scipy_sparse.spmatrix):
        return scipy_sparse.vstack(value, format='csr')
    if sparse and isinstance(value[0], sparse.SparseArray):
        return sparse.concatenate(value, axis=0)
    if PANDAS_INSTALLED and isinstance(value[0], (DataFrame, Series)):
        return pandas_concat(value, axis=0)
    if lazy_isinstance(value[0], 'cudf.core.dataframe', 'DataFrame') or \
       lazy_isinstance(value[0], 'cudf.core.series', 'Series'):
        from cudf import concat as CUDF_concat  # pylint: disable=import-error
        return CUDF_concat(value, axis=0)
    if lazy_isinstance(value[0], 'cupy.core.core', 'ndarray'):
        import cupy
        # pylint: disable=c-extension-no-member,no-member
        d = cupy.cuda.runtime.getDevice()
        for v in value:
            d_v = v.device.id
            assert d_v == d, 'Concatenating arrays on different devices.'
        return cupy.concatenate(value, axis=0)
    return dd.multi.concat(list(value), axis=0)


def _xgb_get_client(client: Optional["distributed.Client"]) -> "distributed.Client":
    '''Simple wrapper around testing None.'''
    if not isinstance(client, (type(distributed.get_client()), type(None))):
        raise TypeError(
            _expect([type(distributed.get_client()), type(None)], type(client)))
    ret = distributed.get_client() if client is None else client
    return ret

# From the implementation point of view, DaskDMatrix complicates a lots of
# things.  A large portion of the code base is about syncing and extracting
# stuffs from DaskDMatrix.  But having an independent data structure gives us a
# chance to perform some specialized optimizations, like building histogram
# index directly.


class DaskDMatrix:
    # pylint: disable=missing-docstring, too-many-instance-attributes
    '''DMatrix holding on references to Dask DataFrame or Dask Array.  Constructing
    a `DaskDMatrix` forces all lazy computation to be carried out.  Wait for
    the input data explicitly if you want to see actual computation of
    constructing `DaskDMatrix`.

    .. note::

        DaskDMatrix does not repartition or move data between workers.  It's
        the caller's responsibility to balance the data.

    .. versionadded:: 1.0.0

    Parameters
    ----------
    client :
        Specify the dask client used for training.  Use default client returned from dask
        if it's set to None.
    data :
        data source of DMatrix.
    label :
        label used for trainin.
    missing :
        Value in the input data (e.g. `numpy.ndarray`) which needs to be present as a
        missing value. If None, defaults to np.nan.
    weight :
        Weight for each instance.
    base_margin :
        Global bias for each instance.
    label_lower_bound :
        Upper bound for survival training.
    label_upper_bound :
        Lower bound for survival training.
    feature_weights :
        Weight for features used in column sampling.
    feature_names :
        Set names for features.
    feature_types :
        Set types for features

    '''

    def __init__(
        self,
        client: "distributed.Client",
        data: _DaskCollection,
        label: Optional[_DaskCollection] = None,
        missing: float = None,
        weight: Optional[_DaskCollection] = None,
        base_margin: Optional[_DaskCollection] = None,
        label_lower_bound: Optional[_DaskCollection] = None,
        label_upper_bound: Optional[_DaskCollection] = None,
        feature_weights: Optional[_DaskCollection] = None,
        feature_names: Optional[Union[str, List[str]]] = None,
        feature_types: Optional[Union[Any, List[Any]]] = None
    ) -> None:
        _assert_dask_support()
        client = _xgb_get_client(client)

        self.feature_names = feature_names
        self.feature_types = feature_types
        self.missing = missing

        if len(data.shape) != 2:
            raise ValueError(
                'Expecting 2 dimensional input, got: {shape}'.format(
                    shape=data.shape))

        if not isinstance(data, (dd.DataFrame, da.Array)):
            raise TypeError(_expect((dd.DataFrame, da.Array), type(data)))
        if not isinstance(label, (dd.DataFrame, da.Array, dd.Series,
                                  type(None))):
            raise TypeError(
                _expect((dd.DataFrame, da.Array, dd.Series), type(label)))

        self.worker_map: Dict[str, "distributed.Future"] = defaultdict(list)
        self.is_quantile: bool = False

        self._init = client.sync(self.map_local_data,
                                 client, data, label=label, weights=weight,
                                 base_margin=base_margin,
                                 feature_weights=feature_weights,
                                 label_lower_bound=label_lower_bound,
                                 label_upper_bound=label_upper_bound)

    def __await__(self) -> Generator:
        return self._init.__await__()

    async def map_local_data(
        self,
        client: "distributed.Client",
        data: _DaskCollection,
        label: Optional[_DaskCollection] = None,
        weights: Optional[_DaskCollection] = None,
        base_margin: Optional[_DaskCollection] = None,
        feature_weights: Optional[_DaskCollection] = None,
        label_lower_bound: Optional[_DaskCollection] = None,
        label_upper_bound: Optional[_DaskCollection] = None
    ) -> "DaskDMatrix":
        '''Obtain references to local data.'''

        def inconsistent(
            left: List[Any], left_name: str, right: List[Any], right_name: str
        ) -> str:
            msg = 'Partitions between {a_name} and {b_name} are not ' \
                'consistent: {a_len} != {b_len}.  ' \
                'Please try to repartition/rechunk your data.'.format(
                    a_name=left_name, b_name=right_name, a_len=len(left),
                    b_len=len(right)
                )
            return msg

        def check_columns(parts: Any) -> None:
            # x is required to be 2 dim in __init__
            assert parts.ndim == 1 or parts.shape[1], 'Data should be' \
                ' partitioned by row. To avoid this specify the number' \
                ' of columns for your dask Array explicitly. e.g.' \
                ' chunks=(partition_size, X.shape[1])'

        data = data.persist()
        for meta in [label, weights, base_margin, label_lower_bound,
                     label_upper_bound]:
            if meta is not None:
                meta = meta.persist()
        # Breaking data into partitions, a trick borrowed from dask_xgboost.

        # `to_delayed` downgrades high-level objects into numpy or pandas
        # equivalents.
        X_parts = data.to_delayed()
        if isinstance(X_parts, numpy.ndarray):
            check_columns(X_parts)
            X_parts = X_parts.flatten().tolist()

        def flatten_meta(
            meta: Optional[_DaskCollection]
        ) -> "Optional[List[dask.delayed.Delayed]]":
            if meta is not None:
                meta_parts = meta.to_delayed()
                if isinstance(meta_parts, numpy.ndarray):
                    check_columns(meta_parts)
                    meta_parts = meta_parts.flatten().tolist()
                return meta_parts
            return None

        y_parts = flatten_meta(label)
        w_parts = flatten_meta(weights)
        margin_parts = flatten_meta(base_margin)
        ll_parts = flatten_meta(label_lower_bound)
        lu_parts = flatten_meta(label_upper_bound)

        parts = [X_parts]
        meta_names = []

        def append_meta(
            m_parts: Optional[List["dask.delayed.delayed"]], name: str
        ) -> None:
            if m_parts is not None:
                assert len(X_parts) == len(
                    m_parts), inconsistent(X_parts, 'X', m_parts, name)
                parts.append(m_parts)
                meta_names.append(name)

        append_meta(y_parts, 'labels')
        append_meta(w_parts, 'weights')
        append_meta(margin_parts, 'base_margin')
        append_meta(ll_parts, 'label_lower_bound')
        append_meta(lu_parts, 'label_upper_bound')
        # At this point, `parts` looks like:
        # [(x0, x1, ..), (y0, y1, ..), ..] in delayed form

        # delay the zipped result
        parts = list(map(dask.delayed, zip(*parts)))  # pylint: disable=no-member
        # At this point, the mental model should look like:
        # [(x0, y0, ..), (x1, y1, ..), ..] in delayed form

        parts = client.compute(parts)
        await distributed.wait(parts)  # async wait for parts to be computed

        for part in parts:
            assert part.status == 'finished', part.status

        # Preserving the partition order for prediction.
        self.partition_order = {}
        for i, part in enumerate(parts):
            self.partition_order[part.key] = i

        key_to_partition = {part.key: part for part in parts}
        who_has = await client.scheduler.who_has(keys=[part.key for part in parts])

        worker_map: Dict[str, "distributed.Future"] = defaultdict(list)

        for key, workers in who_has.items():
            worker_map[next(iter(workers))].append(key_to_partition[key])

        self.worker_map = worker_map
        self.meta_names = meta_names

        if feature_weights is None:
            self.feature_weights = None
        else:
            self.feature_weights = await client.compute(feature_weights).result()

        return self

    def create_fn_args(self, worker_addr: str) -> Dict[str, Any]:
        '''Create a dictionary of objects that can be pickled for function
        arguments.

        '''
        return {'feature_names': self.feature_names,
                'feature_types': self.feature_types,
                'feature_weights': self.feature_weights,
                'meta_names': self.meta_names,
                'missing': self.missing,
                'parts': self.worker_map.get(worker_addr, None),
                'is_quantile': self.is_quantile}


_DataParts = List[Tuple[Any, Optional[Any], Optional[Any], Optional[Any], Optional[Any],
                        Optional[Any]]]


def _get_worker_parts_ordered(
    meta_names: List[str], list_of_parts: _DataParts
) -> _DataParts:
    # List of partitions like: [(x3, y3, w3, m3, ..), ..], order is not preserved.
    assert isinstance(list_of_parts, list)

    result = []

    for i, _ in enumerate(list_of_parts):
        data = list_of_parts[i][0]
        labels = None
        weights = None
        base_margin = None
        label_lower_bound = None
        label_upper_bound = None
        # Iterate through all possible meta info, brings small overhead as in xgboost
        # there are constant number of meta info available.
        for j, blob in enumerate(list_of_parts[i][1:]):
            if meta_names[j] == 'labels':
                labels = blob
            elif meta_names[j] == 'weights':
                weights = blob
            elif meta_names[j] == 'base_margin':
                base_margin = blob
            elif meta_names[j] == 'label_lower_bound':
                label_lower_bound = blob
            elif meta_names[j] == 'label_upper_bound':
                label_upper_bound = blob
            else:
                raise ValueError('Unknown metainfo:', meta_names[j])
        result.append((data, labels, weights, base_margin, label_lower_bound,
                       label_upper_bound))
    return result


def _unzip(list_of_parts: _DataParts) -> List[Tuple[Any, ...]]:
    return list(zip(*list_of_parts))


def _get_worker_parts(
    list_of_parts: _DataParts, meta_names: List[str]
) -> List[Tuple[Any, ...]]:
    partitions = _get_worker_parts_ordered(meta_names, list_of_parts)
    partitions_unzipped = _unzip(partitions)
    return partitions_unzipped


class DaskPartitionIter(DataIter):  # pylint: disable=R0902
    """A data iterator for `DaskDeviceQuantileDMatrix`."""

    def __init__(
        self,
        data: Tuple[Any, ...],
        label: Optional[Tuple[Any, ...]] = None,
        weight: Optional[Tuple[Any, ...]] = None,
        base_margin: Optional[Tuple[Any, ...]] = None,
        label_lower_bound: Optional[Tuple[Any, ...]] = None,
        label_upper_bound: Optional[Tuple[Any, ...]] = None,
        feature_names: Optional[Union[str, List[str]]] = None,
        feature_types: Optional[Union[Any, List[Any]]] = None
    ) -> None:
        self._data = data
        self._labels = label
        self._weights = weight
        self._base_margin = base_margin
        self._label_lower_bound = label_lower_bound
        self._label_upper_bound = label_upper_bound
        self._feature_names = feature_names
        self._feature_types = feature_types

        assert isinstance(self._data, Sequence)

        types = (Sequence, type(None))
        assert isinstance(self._labels, types)
        assert isinstance(self._weights, types)
        assert isinstance(self._base_margin, types)
        assert isinstance(self._label_lower_bound, types)
        assert isinstance(self._label_upper_bound, types)

        self._iter = 0             # set iterator to 0
        super().__init__()

    def data(self) -> Any:
        '''Utility function for obtaining current batch of data.'''
        return self._data[self._iter]

    def labels(self) -> Any:
        '''Utility function for obtaining current batch of label.'''
        if self._labels is not None:
            return self._labels[self._iter]
        return None

    def weights(self) -> Any:
        '''Utility function for obtaining current batch of label.'''
        if self._weights is not None:
            return self._weights[self._iter]
        return None

    def base_margins(self) -> Any:
        '''Utility function for obtaining current batch of base_margin.'''
        if self._base_margin is not None:
            return self._base_margin[self._iter]
        return None

    def label_lower_bounds(self) -> Any:
        '''Utility function for obtaining current batch of label_lower_bound.
        '''
        if self._label_lower_bound is not None:
            return self._label_lower_bound[self._iter]
        return None

    def label_upper_bounds(self) -> Any:
        '''Utility function for obtaining current batch of label_upper_bound.
        '''
        if self._label_upper_bound is not None:
            return self._label_upper_bound[self._iter]
        return None

    def reset(self) -> None:
        '''Reset the iterator'''
        self._iter = 0

    def next(self, input_data: Callable) -> int:
        '''Yield next batch of data'''
        if self._iter == len(self._data):
            # Return 0 when there's no more batch.
            return 0
        feature_names: Optional[Union[List[str], str]] = None
        if self._feature_names:
            feature_names = self._feature_names
        else:
            if hasattr(self.data(), 'columns'):
                feature_names = self.data().columns.format()
            else:
                feature_names = None
        input_data(data=self.data(), label=self.labels(),
                   weight=self.weights(), group=None,
                   label_lower_bound=self.label_lower_bounds(),
                   label_upper_bound=self.label_upper_bounds(),
                   feature_names=feature_names,
                   feature_types=self._feature_types)
        self._iter += 1
        return 1


class DaskDeviceQuantileDMatrix(DaskDMatrix):
    '''Specialized data type for `gpu_hist` tree method.  This class is used to
    reduce the memory usage by eliminating data copies.  Internally the all
    partitions/chunks of data are merged by weighted GK sketching.  So the
    number of partitions from dask may affect training accuracy as GK generates
    bounded error for each merge.

    .. versionadded:: 1.2.0

    Parameters
    ----------
    max_bin : Number of bins for histogram construction.

    '''
    def __init__(
        self,
        client: "distributed.Client",
        data: _DaskCollection,
        label: Optional[_DaskCollection] = None,
        missing: float = None,
        weight: Optional[_DaskCollection] = None,
        base_margin: Optional[_DaskCollection] = None,
        label_lower_bound: Optional[_DaskCollection] = None,
        label_upper_bound: Optional[_DaskCollection] = None,
        feature_weights: Optional[_DaskCollection] = None,
        feature_names: Optional[Union[str, List[str]]] = None,
        feature_types: Optional[Union[Any, List[Any]]] = None,
        max_bin: int = 256
    ) -> None:
        super().__init__(client=client, data=data, label=label,
                         missing=missing,
                         weight=weight, base_margin=base_margin,
                         label_lower_bound=label_lower_bound,
                         label_upper_bound=label_upper_bound,
                         feature_names=feature_names,
                         feature_types=feature_types)
        self.max_bin = max_bin
        self.is_quantile = True

    def create_fn_args(self, worker_addr: str) -> Dict[str, Any]:
        args = super().create_fn_args(worker_addr)
        args['max_bin'] = self.max_bin
        return args


def _create_device_quantile_dmatrix(
    feature_names: Optional[Union[str, List[str]]],
    feature_types: Optional[Union[Any, List[Any]]],
    feature_weights: Optional[Any],
    meta_names: List[str],
    missing: float,
    parts: Optional[_DataParts],
    max_bin: int
) -> DeviceQuantileDMatrix:
    worker = distributed.get_worker()
    if parts is None:
        msg = 'worker {address} has an empty DMatrix.  '.format(
            address=worker.address)
        LOGGER.warning(msg)
        import cupy
        d = DeviceQuantileDMatrix(cupy.zeros((0, 0)),
                                  feature_names=feature_names,
                                  feature_types=feature_types,
                                  max_bin=max_bin)
        return d

    (data, labels, weights, base_margin,
     label_lower_bound, label_upper_bound) = _get_worker_parts(
         parts, meta_names)
    it = DaskPartitionIter(data=data, label=labels, weight=weights,
                           base_margin=base_margin,
                           label_lower_bound=label_lower_bound,
                           label_upper_bound=label_upper_bound)

    dmatrix = DeviceQuantileDMatrix(it,
                                    missing=missing,
                                    feature_names=feature_names,
                                    feature_types=feature_types,
                                    nthread=worker.nthreads,
                                    max_bin=max_bin)
    dmatrix.set_info(feature_weights=feature_weights)
    return dmatrix


def _create_dmatrix(
    feature_names: Optional[Union[str, List[str]]],
    feature_types: Optional[Union[Any, List[Any]]],
    feature_weights: Optional[Any],
    meta_names: List[str],
    missing: float,
    parts: Optional[_DataParts]
) -> DMatrix:
    '''Get data that local to worker from DaskDMatrix.

      Returns
      -------
      A DMatrix object.

    '''
    worker = distributed.get_worker()
    list_of_parts = parts
    if list_of_parts is None:
        msg = 'worker {address} has an empty DMatrix.  '.format(address=worker.address)
        LOGGER.warning(msg)
        d = DMatrix(numpy.empty((0, 0)),
                    feature_names=feature_names,
                    feature_types=feature_types)
        return d

    T = TypeVar('T')

    def concat_or_none(data: Tuple[Optional[T], ...]) -> Optional[T]:
        if any([part is None for part in data]):
            return None
        return concat(data)

    (data, labels, weights, base_margin,
     label_lower_bound, label_upper_bound) = _get_worker_parts(list_of_parts, meta_names)

    _labels = concat_or_none(labels)
    _weights = concat_or_none(weights)
    _base_margin = concat_or_none(base_margin)
    _label_lower_bound = concat_or_none(label_lower_bound)
    _label_upper_bound = concat_or_none(label_upper_bound)

    _data = concat(data)
    dmatrix = DMatrix(_data,
                      _labels,
                      missing=missing,
                      feature_names=feature_names,
                      feature_types=feature_types,
                      nthread=worker.nthreads)
    dmatrix.set_info(base_margin=_base_margin, weight=_weights,
                     label_lower_bound=_label_lower_bound,
                     label_upper_bound=_label_upper_bound,
                     feature_weights=feature_weights)
    return dmatrix


def _dmatrix_from_list_of_parts(
    is_quantile: bool, **kwargs: Any
) -> Union[DMatrix, DeviceQuantileDMatrix]:
    if is_quantile:
        return _create_device_quantile_dmatrix(**kwargs)
    return _create_dmatrix(**kwargs)


async def _get_rabit_args(n_workers: int, client: "distributed.Client") -> List[bytes]:
    '''Get rabit context arguments from data distribution in DaskDMatrix.'''
    env = await client.run_on_scheduler(_start_tracker, n_workers)
    rabit_args = [('%s=%s' % item).encode() for item in env.items()]
    return rabit_args

# train and predict methods are supposed to be "functional", which meets the
# dask paradigm.  But as a side effect, the `evals_result` in single-node API
# is no longer supported since it mutates the input parameter, and it's not
# intuitive to sync the mutation result.  Therefore, a dictionary containing
# evaluation history is instead returned.


def _get_workers_from_data(
    dtrain: DaskDMatrix,
    evals: Optional[List[Tuple[DaskDMatrix, str]]]
) -> Set[str]:
    X_worker_map: Set[str] = set(dtrain.worker_map.keys())
    if evals:
        for e in evals:
            assert len(e) == 2
            assert isinstance(e[0], DaskDMatrix) and isinstance(e[1], str)
            worker_map = set(e[0].worker_map.keys())
            X_worker_map = X_worker_map.union(worker_map)
    return X_worker_map


async def _train_async(
    client: "distributed.Client",
    global_config: Dict[str, Any],
    params: Dict[str, Any],
    dtrain: DaskDMatrix,
    num_boost_round: int,
    evals: Optional[List[Tuple[DaskDMatrix, str]]],
    obj: Optional[Objective],
    feval: Optional[Metric],
    early_stopping_rounds: Optional[int],
    verbose_eval: Union[int, bool],
    xgb_model: Optional[Booster],
    callbacks: Optional[List[TrainingCallback]]
) -> Optional[TrainReturnT]:
    workers = list(_get_workers_from_data(dtrain, evals))
    _rabit_args = await _get_rabit_args(len(workers), client)

    def dispatched_train(
        worker_addr: str,
        rabit_args: List[bytes],
        dtrain_ref: Dict,
        dtrain_idt: int,
        evals_ref: Dict
    ) -> Optional[Dict[str, Union[Booster, Dict]]]:
        '''Perform training on a single worker.  A local function prevents pickling.

        '''
        LOGGER.info('Training on %s', str(worker_addr))
        worker = distributed.get_worker()
        with RabitContext(rabit_args), config.config_context(**global_config):
            local_dtrain = _dmatrix_from_list_of_parts(**dtrain_ref)
            local_evals = []
            if evals_ref:
                for ref, name, idt in evals_ref:
                    if idt == dtrain_idt:
                        local_evals.append((local_dtrain, name))
                        continue
                    local_evals.append((_dmatrix_from_list_of_parts(**ref), name))

            local_history: Dict = {}
            local_param = params.copy()  # just to be consistent
            msg = 'Overriding `nthreads` defined in dask worker.'
            override = ['nthread', 'n_jobs']
            for p in override:
                val = local_param.get(p, None)
                if val is not None and val != worker.nthreads:
                    LOGGER.info(msg)
                else:
                    local_param[p] = worker.nthreads
            bst = worker_train(params=local_param,
                               dtrain=local_dtrain,
                               num_boost_round=num_boost_round,
                               evals_result=local_history,
                               evals=local_evals,
                               obj=obj,
                               feval=feval,
                               early_stopping_rounds=early_stopping_rounds,
                               verbose_eval=verbose_eval,
                               xgb_model=xgb_model,
                               callbacks=callbacks)
            ret: Optional[Dict[str, Union[Booster, Dict]]] = {
                'booster': bst, 'history': local_history}
            if local_dtrain.num_row() == 0:
                ret = None
            return ret

    # Note for function purity:
    # XGBoost is deterministic in most of the cases, which means train function is
    # supposed to be idempotent.  One known exception is gblinear with shotgun updater.
    # We haven't been able to do a full verification so here we keep pure to be False.
    futures = []
    for i, worker_addr in enumerate(workers):
        if evals:
            evals_per_worker = [(e.create_fn_args(worker_addr), name, id(e))
                                for e, name in evals]
        else:
            evals_per_worker = []
        f = client.submit(dispatched_train,
                          worker_addr,
                          _rabit_args,
                          dtrain.create_fn_args(workers[i]),
                          id(dtrain),
                          evals_per_worker,
                          pure=False,
                          workers=[worker_addr])
        futures.append(f)

    results = await client.gather(futures)
    return list(filter(lambda ret: ret is not None, results))[0]


def train(
    client: "distributed.Client",
    params: Dict[str, Any],
    dtrain: DaskDMatrix,
    num_boost_round: int = 10,
    evals: Optional[List[Tuple[DaskDMatrix, str]]] = None,
    obj: Optional[Objective] = None,
    feval: Optional[Metric] = None,
    early_stopping_rounds: Optional[int] = None,
    xgb_model: Optional[Booster] = None,
    verbose_eval: Union[int, bool] = True,
    callbacks: Optional[List[TrainingCallback]] = None
) -> Any:
    '''Train XGBoost model.

    .. versionadded:: 1.0.0

    .. note::

        Other parameters are the same as `xgboost.train` except for `evals_result`, which
        is returned as part of function return value instead of argument.

    Parameters
    ----------
    client :
        Specify the dask client used for training.  Use default client returned from dask
        if it's set to None.

    Returns
    -------
    results: dict
        A dictionary containing trained booster and evaluation history.  `history` field
        is the same as `eval_result` from `xgboost.train`.

        .. code-block:: python

            {'booster': xgboost.Booster,
             'history': {'train': {'logloss': ['0.48253', '0.35953']},
                         'eval': {'logloss': ['0.480385', '0.357756']}}}
    '''
    _assert_dask_support()
    client = _xgb_get_client(client)
    # Get global configuration before transferring computation to another thread or
    # process.
    global_config = config.get_config()
    return client.sync(_train_async,
                       client=client,
                       global_config=global_config,
                       num_boost_round=num_boost_round,
                       obj=obj,
                       feval=feval,
                       params=params,
                       dtrain=dtrain,
                       evals=evals,
                       early_stopping_rounds=early_stopping_rounds,
                       verbose_eval=verbose_eval,
                       xgb_model=xgb_model,
                       callbacks=callbacks)


async def _direct_predict_impl(
    client: "distributed.Client",
    data: _DaskCollection,
    predict_fn: Callable
) -> _DaskCollection:
    if isinstance(data, da.Array):
        predictions = await client.submit(
            da.map_blocks,
            predict_fn, data, False, drop_axis=1,
            dtype=numpy.float32
        ).result()
        return predictions
    if isinstance(data, dd.DataFrame):
        predictions = await client.submit(
            dd.map_partitions,
            predict_fn, data, True,
            meta=dd.utils.make_meta({'prediction': 'f4'})
        ).result()
        return predictions.iloc[:, 0]
    raise TypeError('data of type: ' + str(type(data)) +
                    ' is not supported by direct prediction')


# pylint: disable=too-many-statements
async def _predict_async(
    client: "distributed.Client",
    global_config: Dict[str, Any],
    model: Union[Booster, Dict],
    data: _DaskCollection,
    output_margin: bool,
    missing: float,
    pred_leaf: bool,
    pred_contribs: bool,
    approx_contribs: bool,
    pred_interactions: bool,
    validate_features: bool
) -> _DaskCollection:
    if isinstance(model, Booster):
        booster = model
    elif isinstance(model, dict):
        booster = model['booster']
    else:
        raise TypeError(_expect([Booster, dict], type(model)))
    if not isinstance(data, (DaskDMatrix, da.Array, dd.DataFrame)):
        raise TypeError(_expect([DaskDMatrix, da.Array, dd.DataFrame],
                                type(data)))

    def mapped_predict(partition: Any, is_df: bool) -> Any:
        worker = distributed.get_worker()
        with config.config_context(**global_config):
            booster.set_param({'nthread': worker.nthreads})
            m = DMatrix(partition, missing=missing, nthread=worker.nthreads)
            predt = booster.predict(
                data=m,
                output_margin=output_margin,
                pred_leaf=pred_leaf,
                pred_contribs=pred_contribs,
                approx_contribs=approx_contribs,
                pred_interactions=pred_interactions,
                validate_features=validate_features
            )
            if is_df:
                if lazy_isinstance(partition, 'cudf', 'core.dataframe.DataFrame'):
                    import cudf
                    predt = cudf.DataFrame(predt, columns=['prediction'])
                else:
                    predt = DataFrame(predt, columns=['prediction'])
            return predt
    # Predict on dask collection directly.
    if isinstance(data, (da.Array, dd.DataFrame)):
        return await _direct_predict_impl(client, data, mapped_predict)

    # Prediction on dask DMatrix.
    worker_map = data.worker_map
    partition_order = data.partition_order
    feature_names = data.feature_names
    feature_types = data.feature_types
    missing = data.missing
    meta_names = data.meta_names

    def dispatched_predict(
            worker_id: int, list_of_orders: List[int], list_of_parts: _DataParts
    ) -> List[Tuple[Tuple["dask.delayed.Delayed", int], int]]:
        '''Perform prediction on each worker.'''
        LOGGER.info('Predicting on %d', worker_id)
        with config.config_context(**global_config):
            worker = distributed.get_worker()
            list_of_parts = _get_worker_parts_ordered(meta_names, list_of_parts)
            predictions = []

            booster.set_param({'nthread': worker.nthreads})
            for i, parts in enumerate(list_of_parts):
                (data, _, _, base_margin, _, _) = parts
                order = list_of_orders[i]
                local_part = DMatrix(
                    data,
                    base_margin=base_margin,
                    feature_names=feature_names,
                    feature_types=feature_types,
                    missing=missing,
                    nthread=worker.nthreads
                )
                predt = booster.predict(
                    data=local_part,
                    output_margin=output_margin,
                    pred_leaf=pred_leaf,
                    pred_contribs=pred_contribs,
                    approx_contribs=approx_contribs,
                    pred_interactions=pred_interactions,
                    validate_features=validate_features
                )
                columns = 1 if len(predt.shape) == 1 else predt.shape[1]
                ret = ((dask.delayed(predt), columns), order)  # pylint: disable=no-member
                predictions.append(ret)

            return predictions

    def dispatched_get_shape(
        worker_id: int, list_of_orders: List[int], list_of_parts: _DataParts
    ) -> List[Tuple[int, int]]:
        '''Get shape of data in each worker.'''
        LOGGER.info('Get shape on %d', worker_id)
        list_of_parts = _get_worker_parts_ordered(meta_names, list_of_parts)
        shapes = []
        for i, parts in enumerate(list_of_parts):
            (data, _, _, _, _, _) = parts
            shapes.append((data.shape, list_of_orders[i]))
        return shapes

    async def map_function(
        func: Callable[[int, List[int], _DataParts], Any]
    ) -> List[Any]:
        '''Run function for each part of the data.'''
        futures = []
        workers_address = list(worker_map.keys())
        for wid, worker_addr in enumerate(workers_address):
            worker_addr = workers_address[wid]
            list_of_parts = worker_map[worker_addr]
            list_of_orders = [partition_order[part.key] for part in list_of_parts]

            f = client.submit(func, worker_id=wid,
                              list_of_orders=list_of_orders,
                              list_of_parts=list_of_parts,
                              pure=True, workers=[worker_addr])
            assert isinstance(f, distributed.client.Future)
            futures.append(f)
        # Get delayed objects
        results = await client.gather(futures)
        # flatten into 1 dim list
        results = [t for list_per_worker in results for t in list_per_worker]
        # sort by order, l[0] is the delayed object, l[1] is its order
        results = sorted(results, key=lambda l: l[1])
        results = [predt for predt, order in results]  # remove order
        return results

    results = await map_function(dispatched_predict)
    shapes = await map_function(dispatched_get_shape)

    # Constructing a dask array from list of numpy arrays
    # See https://docs.dask.org/en/latest/array-creation.html
    arrays = []
    for i, shape in enumerate(shapes):
        arrays.append(da.from_delayed(
            results[i][0], shape=(shape[0],)
            if results[i][1] == 1 else (shape[0], results[i][1]),
            dtype=numpy.float32))
    predictions = await da.concatenate(arrays, axis=0)
    return predictions


def predict(
    client: "distributed.Client",
    model: Union[TrainReturnT, Booster],
    data: Union[DaskDMatrix, _DaskCollection],
    output_margin: bool = False,
    missing: float = numpy.nan,
    pred_leaf: bool = False,
    pred_contribs: bool = False,
    approx_contribs: bool = False,
    pred_interactions: bool = False,
    validate_features: bool = True
) -> Any:
    '''Run prediction with a trained booster.

    .. note::

        Only default prediction mode is supported right now.

    .. versionadded:: 1.0.0

    Parameters
    ----------
    client:
        Specify the dask client used for training.  Use default client
        returned from dask if it's set to None.
    model:
        The trained model.
    data:
        Input data used for prediction.  When input is a dataframe object,
        prediction output is a series.
    missing:
        Used when input data is not DaskDMatrix.  Specify the value
        considered as missing.

    Returns
    -------
    prediction: dask.array.Array/dask.dataframe.Series

    '''
    _assert_dask_support()
    client = _xgb_get_client(client)
    global_config = config.get_config()
    return client.sync(
        _predict_async, client, global_config, model, data,
        output_margin=output_margin,
        missing=missing,
        pred_leaf=pred_leaf,
        pred_contribs=pred_contribs,
        approx_contribs=approx_contribs,
        pred_interactions=pred_interactions,
        validate_features=validate_features
    )


async def _inplace_predict_async(
    client: "distributed.Client",
    global_config: Dict[str, Any],
    model: Union[Booster, Dict],
    data: _DaskCollection,
    iteration_range: Tuple[int, int] = (0, 0),
    predict_type: str = 'value',
    missing: float = numpy.nan
) -> _DaskCollection:
    client = _xgb_get_client(client)
    if isinstance(model, Booster):
        booster = model
    elif isinstance(model, dict):
        booster = model['booster']
    else:
        raise TypeError(_expect([Booster, dict], type(model)))
    if not isinstance(data, (da.Array, dd.DataFrame)):
        raise TypeError(_expect([da.Array, dd.DataFrame], type(data)))

    def mapped_predict(data: Any, is_df: bool) -> Any:
        worker = distributed.get_worker()
        config.set_config(**global_config)
        booster.set_param({'nthread': worker.nthreads})
        prediction = booster.inplace_predict(
            data,
            iteration_range=iteration_range,
            predict_type=predict_type,
            missing=missing)
        if is_df:
            if lazy_isinstance(data, 'cudf.core.dataframe', 'DataFrame'):
                import cudf
                prediction = cudf.DataFrame({'prediction': prediction},
                                            dtype=numpy.float32)
            else:
                # If it's  from pandas, the partition is a numpy array
                prediction = DataFrame(prediction, columns=['prediction'],
                                       dtype=numpy.float32)
        return prediction

    return await _direct_predict_impl(client, data, mapped_predict)


def inplace_predict(
    client: "distributed.Client",
    model: Union[TrainReturnT, Booster],
    data: _DaskCollection,
    iteration_range: Tuple[int, int] = (0, 0),
    predict_type: str = 'value',
    missing: float = numpy.nan
) -> Any:
    '''Inplace prediction.

    .. versionadded:: 1.1.0

    Parameters
    ----------
    client:
        Specify the dask client used for training.  Use default client
        returned from dask if it's set to None.
    model:
        The trained model.
    iteration_range:
        Specify the range of trees used for prediction.
    predict_type:
        * 'value': Normal prediction result.
        * 'margin': Output the raw untransformed margin value.
    missing:
        Value in the input data which needs to be present as a missing
        value. If None, defaults to np.nan.

    Returns
    -------
    prediction
    '''
    _assert_dask_support()
    client = _xgb_get_client(client)
    global_config = config.get_config()
    return client.sync(_inplace_predict_async, client, global_config, model=model,
                       data=data,
                       iteration_range=iteration_range,
                       predict_type=predict_type,
                       missing=missing)


async def _evaluation_matrices(
    client: "distributed.Client",
    validation_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]],
    sample_weight: Optional[List[_DaskCollection]],
    missing: float
) -> Optional[List[Tuple[DaskDMatrix, str]]]:
    '''
    Parameters
    ----------
    validation_set: list of tuples
        Each tuple contains a validation dataset including input X and label y.
        E.g.:

        .. code-block:: python

          [(X_0, y_0), (X_1, y_1), ... ]

    sample_weights: list of arrays
        The weight vector for validation data.

    Returns
    -------
    evals: list of validation DMatrix
    '''
    evals: Optional[List[Tuple[DaskDMatrix, str]]] = []
    if validation_set is not None:
        assert isinstance(validation_set, list)
        for i, e in enumerate(validation_set):
            w = (sample_weight[i] if sample_weight is not None else None)
            dmat = await DaskDMatrix(client=client, data=e[0], label=e[1],
                                     weight=w, missing=missing)
            assert isinstance(evals, list)
            evals.append((dmat, 'validation_{}'.format(i)))
    else:
        evals = None
    return evals


class DaskScikitLearnBase(XGBModel):
    '''Base class for implementing scikit-learn interface with Dask'''

    _client = None

    # pylint: disable=arguments-differ
    @_deprecate_positional_args
    def fit(
        self,
        X: _DaskCollection,
        y: _DaskCollection,
        *,
        sample_weight: Optional[_DaskCollection] = None,
        base_margin: Optional[_DaskCollection] = None,
        eval_set: List[Tuple[_DaskCollection, _DaskCollection]] = None,
        eval_metric: Optional[Callable] = None,
        sample_weight_eval_set: Optional[List[_DaskCollection]] = None,
        early_stopping_rounds: Optional[int] = None,
        verbose: bool = True,
        feature_weights: Optional[_DaskCollection] = None,
        callbacks: List[TrainingCallback] = None
    ) -> "DaskScikitLearnBase":
        '''Fit gradient boosting model

        Parameters
        ----------
        X : array_like
            Feature matrix
        y : array_like
            Labels
        sample_weight : array_like
            instance weights
        eval_set : list, optional
            A list of (X, y) tuple pairs to use as validation sets, for which
            metrics will be computed.
            Validation metrics will help us track the performance of the model.
        eval_metric : str, list of str, or callable, optional
        sample_weight_eval_set : list, optional
            A list of the form [L_1, L_2, ..., L_n], where each L_i is a list
            of group weights on the i-th validation set.
        early_stopping_rounds : int
            Activates early stopping.
        verbose : bool
            If `verbose` and an evaluation set is used, writes the evaluation
            metric measured on the validation set to stderr.
        feature_weights: array_like
            Weight for each feature, defines the probability of each feature being
            selected when colsample is being used.  All values must be greater than 0,
            otherwise a `ValueError` is thrown.  Only available for `hist`, `gpu_hist` and
            `exact` tree methods.
        callbacks : list of callback functions
            List of callback functions that are applied at end of each iteration.
            It is possible to use predefined callbacks by using :ref:`callback_api`.
            Example:

            .. code-block:: python

                callbacks = [xgb.callback.EarlyStopping(rounds=early_stopping_rounds,
                                                        save_best=True)]

        '''
        raise NotImplementedError

    def predict(
        self,
        data: _DaskCollection,
        output_margin: bool = False,
        ntree_limit: Optional[int] = None,
        validate_features: bool = True,
        base_margin: Optional[_DaskCollection] = None
    ) -> Any:
        '''Predict with `data`.
        Parameters
        ----------
        data: data that can be used to construct a DaskDMatrix
        output_margin : Whether to output the raw untransformed margin value.
        ntree_limit : NOT supported on dask interface.
        validate_features :
            When this is True, validate that the Booster's and data's feature_names are
            identical.  Otherwise, it is assumed that the feature_names are the same.
        Returns
        -------
        prediction:

        '''
        raise NotImplementedError

    def __await__(self) -> Awaitable[Any]:
        # Generate a coroutine wrapper to make this class awaitable.
        async def _() -> Awaitable[Any]:
            return self
        return self.client.sync(_).__await__()

    @property
    def client(self) -> "distributed.Client":
        '''The dask client used in this model.'''
        client = _xgb_get_client(self._client)
        return client

    @client.setter
    def client(self, clt: "distributed.Client") -> None:
        self._client = clt


@xgboost_model_doc("""Implementation of the Scikit-Learn API for XGBoost.""",
                   ['estimators', 'model'])
class DaskXGBRegressor(DaskScikitLearnBase, XGBRegressorBase):
    # pylint: disable=missing-class-docstring
    async def _fit_async(
        self, X: _DaskCollection,
        y: _DaskCollection,
        sample_weight: Optional[_DaskCollection],
        base_margin: Optional[_DaskCollection],
        eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]],
        eval_metric: Optional[Union[str, List[str], Callable]],
        sample_weight_eval_set: Optional[List[_DaskCollection]],
        early_stopping_rounds: int,
        verbose: bool,
        feature_weights: Optional[_DaskCollection],
        callbacks: Optional[List[TrainingCallback]]
    ) -> _DaskCollection:
        dtrain = await DaskDMatrix(client=self.client,
                                   data=X,
                                   label=y,
                                   weight=sample_weight,
                                   base_margin=base_margin,
                                   feature_weights=feature_weights,
                                   missing=self.missing)
        params = self.get_xgb_params()
        evals = await _evaluation_matrices(self.client, eval_set,
                                           sample_weight_eval_set,
                                           self.missing)

        if callable(self.objective):
            obj = _objective_decorator(self.objective)
        else:
            obj = None
        metric = eval_metric if callable(eval_metric) else None
        if eval_metric is not None:
            if callable(eval_metric):
                eval_metric = None
            else:
                params.update({"eval_metric": eval_metric})

        results = await train(client=self.client,
                              params=params,
                              dtrain=dtrain,
                              num_boost_round=self.get_num_boosting_rounds(),
                              evals=evals,
                              feval=metric,
                              obj=obj,
                              verbose_eval=verbose,
                              early_stopping_rounds=early_stopping_rounds,
                              callbacks=callbacks)
        self._Booster = results['booster']
        # pylint: disable=attribute-defined-outside-init
        self.evals_result_ = results['history']
        return self

    # pylint: disable=missing-docstring
    @_deprecate_positional_args
    def fit(
        self,
        X: _DaskCollection,
        y: _DaskCollection,
        *,
        sample_weight: Optional[_DaskCollection] = None,
        base_margin: Optional[_DaskCollection] = None,
        eval_set: List[Tuple[_DaskCollection, _DaskCollection]] = None,
        eval_metric: Optional[Callable] = None,
        sample_weight_eval_set: Optional[List[_DaskCollection]] = None,
        early_stopping_rounds: Optional[int] = None,
        verbose: bool = True,
        feature_weights: Optional[_DaskCollection] = None,
        callbacks: List[TrainingCallback] = None
    ) -> "DaskXGBRegressor":
        _assert_dask_support()
        return self.client.sync(self._fit_async,
                                X=X,
                                y=y,
                                sample_weight=sample_weight,
                                base_margin=base_margin,
                                eval_set=eval_set,
                                eval_metric=eval_metric,
                                sample_weight_eval_set=sample_weight_eval_set,
                                early_stopping_rounds=early_stopping_rounds,
                                verbose=verbose,
                                feature_weights=feature_weights,
                                callbacks=callbacks)

    async def _predict_async(
        self, data: _DaskCollection,
        output_margin: bool = False,
        validate_features: bool = True,
        base_margin: Optional[_DaskCollection] = None
    ) -> _DaskCollection:
        test_dmatrix = await DaskDMatrix(
            client=self.client, data=data, base_margin=base_margin,
            missing=self.missing
        )
        pred_probs = await predict(client=self.client,
                                   model=self.get_booster(), data=test_dmatrix,
                                   output_margin=output_margin,
                                   validate_features=validate_features)
        return pred_probs

    # pylint: disable=arguments-differ
    def predict(
        self,
        data: _DaskCollection,
        output_margin: bool = False,
        ntree_limit: Optional[int] = None,
        validate_features: bool = True,
        base_margin: Optional[_DaskCollection] = None
    ) -> Any:
        _assert_dask_support()
        msg = '`ntree_limit` is not supported on dask, use model slicing instead.'
        assert ntree_limit is None, msg
        return self.client.sync(self._predict_async, data,
                                output_margin=output_margin,
                                validate_features=validate_features,
                                base_margin=base_margin)


@xgboost_model_doc(
    'Implementation of the scikit-learn API for XGBoost classification.',
    ['estimators', 'model'])
class DaskXGBClassifier(DaskScikitLearnBase, XGBClassifierBase):
    # pylint: disable=missing-class-docstring
    async def _fit_async(
        self, X: _DaskCollection, y: _DaskCollection,
        sample_weight: Optional[_DaskCollection],
        base_margin: Optional[_DaskCollection],
        eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]],
        eval_metric: Optional[Union[str, List[str], Callable]],
        sample_weight_eval_set: Optional[List[_DaskCollection]],
        early_stopping_rounds: int,
        verbose: bool,
        feature_weights: Optional[_DaskCollection],
        callbacks: Optional[List[TrainingCallback]]
    ) -> "DaskXGBClassifier":
        dtrain = await DaskDMatrix(client=self.client,
                                   data=X,
                                   label=y,
                                   weight=sample_weight,
                                   base_margin=base_margin,
                                   feature_weights=feature_weights,
                                   missing=self.missing)
        params = self.get_xgb_params()

        # pylint: disable=attribute-defined-outside-init
        if isinstance(y, (da.Array)):
            self.classes_ = await self.client.compute(da.unique(y))
        else:
            self.classes_ = await self.client.compute(y.drop_duplicates())
        self.n_classes_ = len(self.classes_)

        if self.n_classes_ > 2:
            params["objective"] = "multi:softprob"
            params['num_class'] = self.n_classes_
        else:
            params["objective"] = "binary:logistic"

        evals = await _evaluation_matrices(self.client, eval_set,
                                           sample_weight_eval_set,
                                           self.missing)

        if callable(self.objective):
            obj = _objective_decorator(self.objective)
        else:
            obj = None
        metric = eval_metric if callable(eval_metric) else None
        if eval_metric is not None:
            if callable(eval_metric):
                eval_metric = None
            else:
                params.update({"eval_metric": eval_metric})

        results = await train(client=self.client,
                              params=params,
                              dtrain=dtrain,
                              num_boost_round=self.get_num_boosting_rounds(),
                              evals=evals,
                              obj=obj,
                              feval=metric,
                              verbose_eval=verbose,
                              early_stopping_rounds=early_stopping_rounds,
                              callbacks=callbacks)
        self._Booster = results['booster']
        # pylint: disable=attribute-defined-outside-init
        self.evals_result_ = results['history']
        return self

    @_deprecate_positional_args
    def fit(
        self,
        X: _DaskCollection,
        y: _DaskCollection,
        *,
        sample_weight: Optional[_DaskCollection] = None,
        base_margin: Optional[_DaskCollection] = None,
        eval_set: Optional[List[Tuple[_DaskCollection, _DaskCollection]]] = None,
        eval_metric: Optional[Union[str, List[str], Callable]] = None,
        sample_weight_eval_set: Optional[List[_DaskCollection]] = None,
        early_stopping_rounds: int = None,
        verbose: bool = True,
        feature_weights: _DaskCollection = None,
        callbacks: Optional[List[TrainingCallback]] = None
    ) -> "DaskXGBClassifier":
        _assert_dask_support()
        return self.client.sync(self._fit_async,
                                X=X,
                                y=y,
                                sample_weight=sample_weight,
                                base_margin=base_margin,
                                eval_set=eval_set,
                                eval_metric=eval_metric,
                                sample_weight_eval_set=sample_weight_eval_set,
                                early_stopping_rounds=early_stopping_rounds,
                                verbose=verbose,
                                feature_weights=feature_weights,
                                callbacks=callbacks)

    async def _predict_proba_async(
        self,
        X: _DaskCollection,
        validate_features: bool,
        output_margin: bool,
        base_margin: Optional[_DaskCollection]
    ) -> _DaskCollection:
        test_dmatrix = await DaskDMatrix(
            client=self.client, data=X, base_margin=base_margin,
            missing=self.missing
        )
        pred_probs = await predict(client=self.client,
                                   model=self.get_booster(),
                                   data=test_dmatrix,
                                   validate_features=validate_features,
                                   output_margin=output_margin)
        return pred_probs

    # pylint: disable=arguments-differ,missing-docstring
    def predict_proba(
        self,
        X: _DaskCollection,
        ntree_limit: Optional[int] = None,
        validate_features: bool = True,
        output_margin: bool = False,
        base_margin: Optional[_DaskCollection] = None
    ) -> Any:
        _assert_dask_support()
        msg = '`ntree_limit` is not supported on dask, use model slicing instead.'
        assert ntree_limit is None, msg
        return self.client.sync(
            self._predict_proba_async,
            X=X,
            validate_features=validate_features,
            output_margin=output_margin,
            base_margin=base_margin
        )

    async def _predict_async(
        self, data: _DaskCollection,
        output_margin: bool = False,
        validate_features: bool = True,
        base_margin: Optional[_DaskCollection] = None
    ) -> _DaskCollection:
        test_dmatrix = await DaskDMatrix(
            client=self.client, data=data, base_margin=base_margin,
            missing=self.missing
        )
        pred_probs = await predict(
            client=self.client,
            model=self.get_booster(),
            data=test_dmatrix,
            output_margin=output_margin,
            validate_features=validate_features
        )

        if self.n_classes_ == 2:
            preds = (pred_probs > 0.5).astype(int)
        else:
            preds = da.argmax(pred_probs, axis=1)

        return preds

    # pylint: disable=arguments-differ
    def predict(
        self,
        data: _DaskCollection,
        output_margin: bool = False,
        ntree_limit: Optional[int] = None,
        validate_features: bool = True,
        base_margin: Optional[_DaskCollection] = None
    ) -> Any:
        _assert_dask_support()
        msg = '`ntree_limit` is not supported on dask, use model slicing instead.'
        assert ntree_limit is None, msg
        return self.client.sync(
            self._predict_async,
            data,
            output_margin=output_margin,
            validate_features=validate_features,
            base_margin=base_margin
        )