Repository: modin-project/modin
Pull request: FEAT-#3243: Add async execution support for Modin xgb.predict
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/CONTRIBUTING.html
if you have questions about contributing.
-->

## What do these changes do?
PR adds to Modin xgb.predict the next:

1. async execution support;
2. change approach from Ray actors using to Ray tasks;
3. exclude force data redistribution;
4. improve performance on a single node/cluster of nodes.

 
<!-- Please give a short brief about these changes. -->

- [x] commit message follows format outlined [here](https://modin.readthedocs.io/en/latest/contributing.html)
- [x] passes `flake8 modin`
- [x] passes `black --check modin`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #3243 ? <!-- issue must be created for each patch -->
- [x] tests added and passing
- [x] module layout described at `docs/developer/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Problem statement:
Add asynchronous execution support in Modin xgboost.predict
### Describe the problem
Current implementation of `predict` in Modin XGBoost can be improved by the next ways:

1. Exclude force partitions redistribution between nodes in Ray cluster.
2. Add asynchronous execution support by excluding blocking calls during execution.

Hints from issue discussion:

Commit message: FEAT-#3243: Apply part of review comments

Signed-off-by: Alexey Prutskov <alexey.prutskov@intel.com>

Patch to review:
diff --git a/docs/flow/modin/experimental/xgboost.rst b/docs/flow/modin/experimental/xgboost.rst
index 658b81e55a4..cb053d19bb2 100644
--- a/docs/flow/modin/experimental/xgboost.rst
+++ b/docs/flow/modin/experimental/xgboost.rst
@@ -27,15 +27,13 @@ row partitions using the function :py:func:`~modin.distributed.dataframe.pandas.
   :members:
 
 :py:class:`~modin.experimental.xgboost.Booster` inherits original class ``xgboost.Booster`` and
-overrides method ``predict``. The main differences from original class interface for ``predict``
-method are: (1) changing the type of the `data` parameter to :py:class:`~modin.experimental.xgboost.DMatrix`, and (2)
-a new parameter `num_actors`, which specifies the number of actors to run for prediction.
+overrides method ``predict``. The difference from original class interface for ``predict``
+method is changing the type of the `data` parameter to :py:class:`~modin.experimental.xgboost.DMatrix`.
 
 .. autoclass:: modin.experimental.xgboost.Booster
     :members:
 
-:py:func:`~modin.experimental.xgboost.train` function (similar to ``predict`` method of
-:py:class:`~modin.experimental.xgboost.Booster`) has 2 differences from the original ``train`` function - (1) the
+:py:func:`~modin.experimental.xgboost.train` function has 2 differences from the original ``train`` function - (1) the
 data type of `dtrain` parameter is :py:class:`~modin.experimental.xgboost.DMatrix` and (2) a new parameter `num_actors`.
 
 .. autofunction:: modin.experimental.xgboost.train
@@ -44,10 +42,13 @@ Internal execution flow on Ray backend
 ''''''''''''''''''''''''''''''''''''''
 
 Internal functions :py:func:`~modin.experimental.xgboost.xgboost_ray._train` and
-:py:func:`~modin.experimental.xgboost.xgboost_ray._predict` work similar to xgboost. Approximate execution flow of
-internal implementation is the following:
+:py:func:`~modin.experimental.xgboost.xgboost_ray._predict` work similar to xgboost.
 
-1. The data is passed to :py:func:`~modin.experimental.xgboost.xgboost_ray._train`/:py:func:`~modin.experimental.xgboost.xgboost_ray._predict`
+
+Training
+********
+
+1. The data is passed to :py:func:`~modin.experimental.xgboost.xgboost_ray._train`
    function as a :py:class:`~modin.experimental.xgboost.DMatrix` object. Using an iterator of
    :py:class:`~modin.experimental.xgboost.DMatrix`, lists of ``ray.ObjectRef`` with row partitions of Modin DataFrame are exctracted. Example:
 
@@ -58,54 +59,46 @@ internal implementation is the following:
    ..
 
 2. On this step, the parameter `num_actors` is processed. The internal function :py:func:`~modin.experimental.xgboost.xgboost_ray._get_num_actors`
-   examines the value provided by the user and checks if it fits in the set of expected values
-   (int, "default_train", "default_predict").
-
-   * int - `num_actors` won't be changed. This value will be used.
-   * "default_train" - `num_actors` will be computed using condition that 1 actor should use maximum 2 CPUs.
-     This condition was chosen for using maximum parallel workers with multithreaded XGBoost training (2 threads
-     per worker will be used in this case).
-   * "default_predict" - `num_actors` will be computed using condition that 1 actor should use maximum 8 CPUs.
-     This condition was chosen to combine parallelization techniques: parallel actors and parallel threads.
-
-.. note:: `num_actors` parameter is made available for public functions :py:func:`~modin.experimental.xgboost.train` and ``predict``
-  method of :py:class:`~modin.experimental.xgboost.Booster` class to allow fine-tuning for obtaining the best
-  performance in specific use cases.
-
-3. ``ray.util.placement_group`` is created to reserve all available Ray resources. After that
-   :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` objects are created using resources of the
-   previously created placement group.
-
-4. Data (`dtrain` for :py:func:`~modin.experimental.xgboost.xgboost_ray._train`, `data` for
-   :py:func:`~modin.experimental.xgboost.xgboost_ray._predict`) is split between actors evenly. The internal function
+   examines the value provided by the user. In case the value wasn't provided, the `num_actors` will be computed using condition that 1 actor should use maximum 2 CPUs.
+   This condition was chosen for using maximum parallel workers with multithreaded XGBoost training (2 threads
+   per worker will be used in this case).
+
+.. note:: `num_actors` parameter is made available for public function :py:func:`~modin.experimental.xgboost.train` to allow
+  fine-tuning for obtaining the best performance in specific use cases.
+
+3. :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` objects are created.
+
+4. Data `dtrain` is split between actors evenly. The internal function
    :py:func:`~modin.experimental.xgboost.xgboost_ray._split_data_across_actors` runs assigning row partitions to actors
    using internal function :py:func:`~modin.experimental.xgboost.xgboost_ray._assign_row_partitions_to_actors`.
-   This function creates a dictionary in the form: `{actor_rank: ([part_i0, part_i3, ..], [0, 3, ..]), ..}` for training,
-   `{actor_rank: [part_i0, part_i1, ..], ..}` for prediction.
+   This function creates a dictionary in the form: `{actor_rank: ([part_i0, part_i3, ..], [0, 3, ..]), ..}`.
 
 .. note:: :py:func:`~modin.experimental.xgboost.xgboost_ray._assign_row_partitions_to_actors` takes into account IP
   addresses of row partitions of `dtrain` data to minimize excess data transfer.
 
-1. For each :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` the object methods ``set_train_data`` or ``set_predict_data`` are
-   called remotely. Those methods run by loading row partitions in actor according to the dictionary with partitions
+5. For each :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` the object method ``set_train_data`` is
+   called remotely. This method runs loading row partitions in actor according to the dictionary with partitions
    distribution from previous step. When data is passed to the actor, the row partitions are automatically materialized
    (``ray.ObjectRef`` -> ``pandas.DataFrame``).
 
-2. Methods ``train`` or ``predict`` of :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` class object are called remotely.
+6. Method ``train`` of :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` class object is called remotely. This method
+   runs XGBoost training on local data of actor, connects to ``Rabit Tracker`` for sharing training state between
+   actors and returns dictionary with `booster` and `evaluation results`.
+
+7. On the final stage results from actors are returned. `booster` and `evals_result` is returned using ``ray.get``
+   function from remote actor.
+
+
+Prediction
+**********
 
-   * ``train``: method runs XGBoost training on local data of actor, connects to ``Rabit Tracker`` for sharing
-     training state between actors and returns dictionary with `booster` and `evaluation results`.
-   * ``predict``: method runs XGBoost prediction on local data of actor and returns IP address of actor and partial
-     prediction (``pandas.DataFrame``).
+1. The data is passed to :py:func:`~modin.experimental.xgboost.xgboost_ray._predict`
+   function as a :py:class:`~modin.experimental.xgboost.DMatrix` object.
 
-3. On the final stage results from actors are returned.
+2. The function :py:func:`~modin.experimental.xgboost.xgboost_ray._map_predict` is applied remotely for each partition
+   of the data to make a partial prediction.
 
-   * ``train``: `booster` and `evals_result` is returned using ``ray.get`` function from remote actor. Placement
-     group which was created on the step 3 is removed to free resources. :py:class:`~modin.experimental.xgboost.Booster`
-     object is created and returned to user. 
-   * ``predict``: using ``ray.wait`` function we wait until all actors finish computing local predictions. Placement group
-     which was created on the step 3 is removed to free resources. ``modin.pandas.DataFrame`` is created from
-     ``ray.ObjectRef`` objects which is returned from actors. Modin DataFrame is returned to user.
+3. Result ``modin.pandas.DataFrame`` is created from ``ray.ObjectRef`` objects, obtained in the previous step.
 
 
 Internal API
@@ -119,3 +112,4 @@ Internal API
 .. autofunction:: modin.experimental.xgboost.xgboost_ray._predict
 .. autofunction:: modin.experimental.xgboost.xgboost_ray._get_num_actors
 .. autofunction:: modin.experimental.xgboost.xgboost_ray._split_data_across_actors
+.. autofunction:: modin.experimental.xgboost.xgboost_ray._map_predict
diff --git a/modin/experimental/xgboost/test/test_xgboost.py b/modin/experimental/xgboost/test/test_xgboost.py
index 430232e655d..a74a9399bdd 100644
--- a/modin/experimental/xgboost/test/test_xgboost.py
+++ b/modin/experimental/xgboost/test/test_xgboost.py
@@ -104,7 +104,7 @@ def test_xgb_with_binary_classification_datasets(data, num_actors, modin_type_y)
             )
 
     predictions = bst.predict(xgb_dmatrix)
-    modin_predictions = modin_bst.predict(mxgb_dmatrix, num_actors=num_actors)
+    modin_predictions = modin_bst.predict(mxgb_dmatrix)
 
     preds = pd.DataFrame(predictions).apply(lambda x: round(x))
     modin_preds = modin_predictions.apply(lambda x: round(x))
@@ -197,7 +197,7 @@ def test_xgb_with_multiclass_classification_datasets(data, num_actors, modin_typ
         )
 
     predictions = bst.predict(xgb_dmatrix)
-    modin_predictions = modin_bst.predict(mxgb_dmatrix, num_actors=num_actors)
+    modin_predictions = modin_bst.predict(mxgb_dmatrix)
 
     array_preds = np.asarray([np.argmax(line) for line in predictions])
     modin_array_preds = np.asarray(
@@ -280,7 +280,7 @@ def test_xgb_with_regression_datasets(data, num_actors, modin_type_y):
             )
 
     predictions = bst.predict(train_xgb_dmatrix)
-    modin_predictions = modin_bst.predict(train_mxgb_dmatrix, num_actors=num_actors)
+    modin_predictions = modin_bst.predict(train_mxgb_dmatrix)
 
     val = mean_squared_error(y_train, predictions)
     modin_val = mean_squared_error(y_train, modin_predictions)
diff --git a/modin/experimental/xgboost/xgboost.py b/modin/experimental/xgboost/xgboost.py
index 2505bad3f67..4d569ce5158 100644
--- a/modin/experimental/xgboost/xgboost.py
+++ b/modin/experimental/xgboost/xgboost.py
@@ -54,6 +54,12 @@ def __init__(self, data, label):
         self.data = unwrap_partitions(data, axis=0, get_ip=True)
         self.label = unwrap_partitions(label, axis=0)
 
+        self.metadata = (
+            data.index,
+            data.columns,
+            data._query_compiler._modin_frame._row_lengths,
+        )
+
     def __iter__(self):
         """
         Return unwrapped `self.data` and `self.label`.
@@ -94,25 +100,20 @@ def __init__(self, params=None, cache=(), model_file=None):  # noqa: MD01
     def predict(
         self,
         data: DMatrix,
-        num_actors: Optional[int] = None,
         **kwargs,
     ):
         """
         Run distributed prediction with a trained booster.
 
-        During work it evenly distributes `data` between workers,
-        runs xgb.predict on each worker for subset of `data` and creates
-        Modin DataFrame with prediction results.
+        During execution it runs ``xgb.predict`` on each worker for subset of `data`
+        and creates Modin DataFrame with prediction results.
 
         Parameters
         ----------
         data : modin.experimental.xgboost.DMatrix
             Input data used for prediction.
-        num_actors : int, optional
-            Number of actors for prediction. If unspecified, this value will be
-            computed automatically.
         **kwargs : dict
-            Other parameters are the same as `xgboost.Booster.predict`.
+            Other parameters are the same as for ``xgboost.Booster.predict``.
 
         Returns
         -------
@@ -130,7 +131,7 @@ def predict(
             data, DMatrix
         ), f"Type of `data` is {type(data)}, but expected {DMatrix}."
 
-        result = _predict(self.copy(), data, num_actors, **kwargs)
+        result = _predict(self.copy(), data, **kwargs)
         LOGGER.info("Prediction finished")
 
         return result
@@ -188,7 +189,7 @@ def train(
     assert isinstance(
         dtrain, DMatrix
     ), f"Type of `dtrain` is {type(dtrain)}, but expected {DMatrix}."
-    result = _train(dtrain, num_actors, params, *args, evals=evals, **kwargs)
+    result = _train(dtrain, params, *args, num_actors=num_actors, evals=evals, **kwargs)
     if isinstance(evals_result, dict):
         evals_result.update(result["history"])
 
diff --git a/modin/experimental/xgboost/xgboost_ray.py b/modin/experimental/xgboost/xgboost_ray.py
index 59edb293ef2..482b18e08a5 100644
--- a/modin/experimental/xgboost/xgboost_ray.py
+++ b/modin/experimental/xgboost/xgboost_ray.py
@@ -24,6 +24,7 @@
 import math
 from collections import defaultdict
 
+import numpy as np
 import xgboost as xgb
 import ray
 from ray.services import get_node_ip_address
@@ -38,7 +39,7 @@
 @ray.remote(num_cpus=0)
 class ModinXGBoostActor:
     """
-    Ray actor-class runs training/prediction on remote workers.
+    Ray actor-class runs training on remote workers.
 
     Parameters
     ----------
@@ -87,17 +88,6 @@ def _get_dmatrix(self, X_y):
 
         return xgb.DMatrix(X, y, nthread=self._nthreads)
 
-    def get_ip(self):
-        """
-        Get IP address of actor.
-
-        Returns
-        -------
-        str
-            An IP address of node used by actor.
-        """
-        return get_node_ip_address()
-
     def set_train_data(self, *X_y, add_as_eval_method=None):
         """
         Set train data for actor.
@@ -116,25 +106,6 @@ def set_train_data(self, *X_y, add_as_eval_method=None):
         if add_as_eval_method is not None:
             self._evals.append((self._dtrain, add_as_eval_method))
 
-    def set_predict_data(
-        self,
-        *X,
-    ):
-        """
-        Set prediction data for actor.
-
-        Parameters
-        ----------
-        *X : iterable
-            Sequence of ray.ObjectRef objects. When it is passed in actor, auto-materialization
-            of ray.ObjectRef -> pandas.DataFrame happens.
-        """
-        X = pandas.concat(X, axis=0)
-        self._dpredict = {
-            "dmatrix": xgb.DMatrix(X, nthread=self._nthreads),
-            "index": X.index,
-        }
-
     def add_eval_data(self, *X_y, eval_method):
         """
         Add evaluation data for actor.
@@ -196,36 +167,6 @@ def train(self, rabit_args, params, *args, **kwargs):
             LOGGER.info(f"Local training time: {time.time() - s} s")
             return {"booster": bst, "history": evals_result}
 
-    def predict(self, booster: xgb.Booster, **kwargs):
-        """
-        Run local XGBoost prediction.
-
-        Parameters
-        ----------
-        booster : xgboost.Booster
-            A trained booster.
-        **kwargs : dict
-            Other parameters for `xgboost.Booster.predict`.
-
-        Returns
-        -------
-        tuple
-            Pair of IP address of caller and pandas.DataFrame
-            with partial prediction result.
-        """
-        local_dpredict = self._dpredict
-        booster.set_param({"nthread": self._nthreads})
-
-        s = time.time()
-
-        predictions = pandas.DataFrame(
-            booster.predict(local_dpredict["dmatrix"], **kwargs),
-            index=local_dpredict["index"],
-        )
-        LOGGER.info(f"Local prediction time: {time.time() - s} s")
-
-        return get_node_ip_address(), predictions
-
 
 def _get_cluster_cpus():
     """
@@ -276,18 +217,15 @@ def _get_cpus_per_actor(num_actors):
     return cpus_per_actor
 
 
-def _get_num_actors(num_actors):
+def _get_num_actors(num_actors=None):
     """
-    Get number on actors to create.
-
-    In case `num_actors` is 'default_train' string,
-    integer number of actors will be computed  by condition
-    2 CPUs per 1 actor, in case 'default_predict' 8 CPUs per 1 actor.
+    Get number of actors to create.
 
     Parameters
     ----------
-    num_actors : int or {'default_train', 'default_predict'}
-        Number of actors to create or str to set defaults.
+    num_actors : int, optional
+        Desired number of actors. If is None, integer number of actors
+        will be computed by condition 2 CPUs per 1 actor.
 
     Returns
     -------
@@ -295,12 +233,9 @@ def _get_num_actors(num_actors):
         Number of actors to create.
     """
     min_cpus_per_node = _get_min_cpus_per_node()
-    if num_actors == "default_train":
+    if num_actors is None:
         num_actors_per_node = max(1, int(min_cpus_per_node // 2))
         return num_actors_per_node * len(ray.nodes())
-    elif num_actors == "default_predict":
-        num_actors_per_node = max(1, int(min_cpus_per_node // 8))
-        return num_actors_per_node * len(ray.nodes())
     elif isinstance(num_actors, int):
         assert (
             num_actors % len(ray.nodes()) == 0
@@ -344,7 +279,10 @@ def create_actors(num_actors):
 
 
 def _split_data_across_actors(
-    actors: List, set_func, X_parts, y_parts=None, is_predict=False
+    actors: List,
+    set_func,
+    X_parts,
+    y_parts,
 ):
     """
     Split row partitions of data between actors.
@@ -357,49 +295,35 @@ def _split_data_across_actors(
         The function for setting data in actor.
     X_parts : list
         Row partitions of X data.
-    y_parts : list, optional
+    y_parts : list
         Row partitions of y data.
-    is_predict : bool, default: False
-        Is split data for predict or not.
     """
     X_parts_by_actors = _assign_row_partitions_to_actors(
         actors,
         X_parts,
-        is_predict=is_predict,
     )
 
-    if y_parts is not None:
-        y_parts_by_actors = _assign_row_partitions_to_actors(
-            actors,
-            y_parts,
-            data_for_aligning=X_parts_by_actors,
-        )
+    y_parts_by_actors = _assign_row_partitions_to_actors(
+        actors,
+        y_parts,
+        data_for_aligning=X_parts_by_actors,
+    )
 
     for rank, (_, actor) in enumerate(actors):
-        X_parts = X_parts_by_actors[rank] if is_predict else X_parts_by_actors[rank][0]
-        if y_parts is None:
-            set_func(actor, *X_parts)
-        else:
-            y_parts = y_parts_by_actors[rank][0]
-            set_func(actor, *(X_parts + y_parts))
+        set_func(actor, *(X_parts_by_actors[rank][0] + y_parts_by_actors[rank][0]))
 
 
 def _assign_row_partitions_to_actors(
     actors: List,
     row_partitions,
     data_for_aligning=None,
-    is_predict=False,
 ):
     """
     Assign row_partitions to actors.
 
-    In case of `is_predict` == False, `row_partitions` will be
-    assigned to actors according to their IPs. If distribution isn't
-    even, partitions will be moved from actor with excess partitions
-    to actor with lack of them.
-
-    In case of `is_predict` == True, `row_partitions` will be
-    assigned evenly to actors in order.
+    `row_partitions` will be assigned to actors according to their IPs.
+    If distribution isn't even, partitions will be moved from actor
+    with excess partitions to actor with lack of them.
 
     Parameters
     ----------
@@ -410,8 +334,6 @@ def _assign_row_partitions_to_actors(
     data_for_aligning : dict, optional
         Data according to the order of which should be
         distributed `row_partitions`. Used to align y with X.
-    is_predict : bool, default: False
-        Is split data for predict or not.
 
     Returns
     -------
@@ -420,138 +342,115 @@ def _assign_row_partitions_to_actors(
         as {actor_rank: (partitions, order)}.
     """
     num_actors = len(actors)
-    if not is_predict:
-        if data_for_aligning is None:
-            parts_ips_ref, parts_ref = zip(*row_partitions)
-
-            # Group actors which are one the same ip
-            actor_ips = defaultdict(list)
-            for rank, (ip, _) in enumerate(actors):
-                actor_ips[ip].append(rank)
-
-            # Get distribution of parts between nodes ({ip:[(part, position),..],..})
-            init_parts_distribution = defaultdict(list)
-            for idx, (ip, part_ref) in enumerate(
-                zip(ray.get(list(parts_ips_ref)), parts_ref)
-            ):
-                init_parts_distribution[ip].append((part_ref, idx))
-
-            num_parts = len(parts_ref)
-            min_parts_per_actor = math.floor(num_parts / num_actors)
-            max_parts_per_actor = math.ceil(num_parts / num_actors)
-            num_actors_with_max_parts = num_parts % num_actors
-
-            row_partitions_by_actors = defaultdict(list)
-            # Fill actors without movement parts between ips
-            for actor_ip, ranks in actor_ips.items():
-                # Loop across actors which are placed on actor_ip
-                for rank in ranks:
-                    num_parts_on_ip = len(init_parts_distribution[actor_ip])
-
-                    # Check that have something to distribute on this ip
-                    if num_parts_on_ip == 0:
-                        break
-                    # Check that node with `actor_ip` has enough parts for minimal
-                    # filling actor with `rank`
-                    if num_parts_on_ip >= min_parts_per_actor:
-                        # Check that node has enough parts for max filling
-                        # actor with `rank`
-                        if (
-                            num_parts_on_ip >= max_parts_per_actor
-                            and num_actors_with_max_parts > 0
-                        ):
-                            pop_slice = slice(0, max_parts_per_actor)
-                            num_actors_with_max_parts -= 1
-                        else:
-                            pop_slice = slice(0, min_parts_per_actor)
-
-                        row_partitions_by_actors[rank].extend(
-                            init_parts_distribution[actor_ip][pop_slice]
-                        )
-                        # Delete parts which we already assign
-                        del init_parts_distribution[actor_ip][pop_slice]
-                    else:
-                        row_partitions_by_actors[rank].extend(
-                            init_parts_distribution[actor_ip]
-                        )
-                        init_parts_distribution[actor_ip] = []
-
-            # Remove empty IPs
-            for ip in list(init_parts_distribution):
-                if len(init_parts_distribution[ip]) == 0:
-                    init_parts_distribution.pop(ip)
-
-            # IP's aren't necessary now
-            init_parts_distribution = [
-                pair for pairs in init_parts_distribution.values() for pair in pairs
-            ]
-
-            # Fill the actors with extra parts (movements data between nodes)
-            for rank in range(len(actors)):
-                num_parts_on_rank = len(row_partitions_by_actors[rank])
-
-                if num_parts_on_rank == max_parts_per_actor or (
-                    num_parts_on_rank == min_parts_per_actor
-                    and num_actors_with_max_parts == 0
-                ):
-                    continue
-
-                if num_actors_with_max_parts > 0:
-                    pop_slice = slice(0, max_parts_per_actor - num_parts_on_rank)
-                    num_actors_with_max_parts -= 1
-                else:
-                    pop_slice = slice(0, min_parts_per_actor - num_parts_on_rank)
-
-                row_partitions_by_actors[rank].extend(
-                    init_parts_distribution[pop_slice]
-                )
-                del init_parts_distribution[pop_slice]
-
-            if len(init_parts_distribution) != 0:
-                raise RuntimeError(
-                    f"Not all partitions were ditributed between actors: {len(init_parts_distribution)} left."
-                )
-
-            row_parts_by_ranks = dict()
-            for rank, pairs_part_pos in dict(row_partitions_by_actors).items():
-                parts, order = zip(*pairs_part_pos)
-                row_parts_by_ranks[rank] = (list(parts), list(order))
-        else:
-            row_parts_by_ranks = {rank: ([], []) for rank in range(len(actors))}
-
-            for rank, (_, order_of_indexes) in data_for_aligning.items():
-                row_parts_by_ranks[rank][1].extend(order_of_indexes)
-                for row_idx in order_of_indexes:
-                    row_parts_by_ranks[rank][0].append(row_partitions[row_idx])
-    else:
-        row_parts_by_ranks = defaultdict(list)
-        _, parts_ref = zip(*row_partitions)
+    if data_for_aligning is None:
+        parts_ips_ref, parts_ref = zip(*row_partitions)
+
+        # Group actors which are one the same ip
+        actor_ips = defaultdict(list)
+        for rank, (ip, _) in enumerate(actors):
+            actor_ips[ip].append(rank)
+
+        # Get distribution of parts between nodes ({ip:[(part, position),..],..})
+        init_parts_distribution = defaultdict(list)
+        for idx, (ip, part_ref) in enumerate(
+            zip(ray.get(list(parts_ips_ref)), parts_ref)
+        ):
+            init_parts_distribution[ip].append((part_ref, idx))
 
         num_parts = len(parts_ref)
         min_parts_per_actor = math.floor(num_parts / num_actors)
         max_parts_per_actor = math.ceil(num_parts / num_actors)
         num_actors_with_max_parts = num_parts % num_actors
 
-        start_idx = 0
-        for rank, actor in enumerate(actors):
+        row_partitions_by_actors = defaultdict(list)
+        # Fill actors without movement parts between ips
+        for actor_ip, ranks in actor_ips.items():
+            # Loop across actors which are placed on actor_ip
+            for rank in ranks:
+                num_parts_on_ip = len(init_parts_distribution[actor_ip])
+
+                # Check that have something to distribute on this ip
+                if num_parts_on_ip == 0:
+                    break
+                # Check that node with `actor_ip` has enough parts for minimal
+                # filling actor with `rank`
+                if num_parts_on_ip >= min_parts_per_actor:
+                    # Check that node has enough parts for max filling
+                    # actor with `rank`
+                    if (
+                        num_parts_on_ip >= max_parts_per_actor
+                        and num_actors_with_max_parts > 0
+                    ):
+                        pop_slice = slice(0, max_parts_per_actor)
+                        num_actors_with_max_parts -= 1
+                    else:
+                        pop_slice = slice(0, min_parts_per_actor)
+
+                    row_partitions_by_actors[rank].extend(
+                        init_parts_distribution[actor_ip][pop_slice]
+                    )
+                    # Delete parts which we already assign
+                    del init_parts_distribution[actor_ip][pop_slice]
+                else:
+                    row_partitions_by_actors[rank].extend(
+                        init_parts_distribution[actor_ip]
+                    )
+                    init_parts_distribution[actor_ip] = []
+
+        # Remove empty IPs
+        for ip in list(init_parts_distribution):
+            if len(init_parts_distribution[ip]) == 0:
+                init_parts_distribution.pop(ip)
+
+        # IP's aren't necessary now
+        init_parts_distribution = [
+            pair for pairs in init_parts_distribution.values() for pair in pairs
+        ]
+
+        # Fill the actors with extra parts (movements data between nodes)
+        for rank in range(len(actors)):
+            num_parts_on_rank = len(row_partitions_by_actors[rank])
+
+            if num_parts_on_rank == max_parts_per_actor or (
+                num_parts_on_rank == min_parts_per_actor
+                and num_actors_with_max_parts == 0
+            ):
+                continue
+
             if num_actors_with_max_parts > 0:
-                num_actor_parts = max_parts_per_actor
+                pop_slice = slice(0, max_parts_per_actor - num_parts_on_rank)
                 num_actors_with_max_parts -= 1
             else:
-                num_actor_parts = min_parts_per_actor
+                pop_slice = slice(0, min_parts_per_actor - num_parts_on_rank)
+
+            row_partitions_by_actors[rank].extend(init_parts_distribution[pop_slice])
+            del init_parts_distribution[pop_slice]
+
+        if len(init_parts_distribution) != 0:
+            raise RuntimeError(
+                f"Not all partitions were ditributed between actors: {len(init_parts_distribution)} left."
+            )
+
+        row_parts_by_ranks = dict()
+        for rank, pairs_part_pos in dict(row_partitions_by_actors).items():
+            parts, order = zip(*pairs_part_pos)
+            row_parts_by_ranks[rank] = (list(parts), list(order))
+    else:
+        row_parts_by_ranks = {rank: ([], []) for rank in range(len(actors))}
 
-            idx_slice = slice(start_idx, start_idx + num_actor_parts)
-            row_parts_by_ranks[rank].extend(parts_ref[idx_slice])
-            start_idx += num_actor_parts
+        for rank, (_, order_of_indexes) in data_for_aligning.items():
+            row_parts_by_ranks[rank][1].extend(order_of_indexes)
+            for row_idx in order_of_indexes:
+                row_parts_by_ranks[rank][0].append(row_partitions[row_idx])
 
     return row_parts_by_ranks
 
 
 def _train(
     dtrain,
-    num_actors,
     params: Dict,
     *args,
+    num_actors=None,
     evals=(),
     **kwargs,
 ):
@@ -568,13 +467,13 @@ def _train(
     ----------
     dtrain : modin.experimental.DMatrix
         Data to be trained against.
-    num_actors : int, optional
-        Number of actors for training. If unspecified, this value will be
-        computed automatically.
     params : dict
         Booster params.
     *args : iterable
         Other parameters for `xgboost.train`.
+    num_actors : int, optional
+        Number of actors for training. If unspecified, this value will be
+        computed automatically.
     evals : list of pairs (modin.experimental.xgboost.DMatrix, str), default: empty
         List of validation sets for which metrics will be evaluated during training.
         Validation metrics will help us track the performance of the model.
@@ -594,9 +493,7 @@ def _train(
 
     assert len(X_row_parts) == len(y_row_parts), "Unaligned train data"
 
-    num_actors = _get_num_actors(
-        num_actors if isinstance(num_actors, int) else "default_train"
-    )
+    num_actors = _get_num_actors(num_actors)
 
     if num_actors > len(X_row_parts):
         num_actors = len(X_row_parts)
@@ -618,7 +515,7 @@ def _train(
                     *X_y, eval_method=eval_method
                 ),
                 eval_X,
-                y_parts=eval_y,
+                eval_y,
             )
 
     # Split data across workers
@@ -628,7 +525,7 @@ def _train(
             *X_y, add_as_eval_method=add_as_eval_method
         ),
         X_row_parts,
-        y_parts=y_row_parts,
+        y_row_parts,
     )
     LOGGER.info(f"Data preparation time: {time.time() - s} s")
 
@@ -648,18 +545,46 @@ def _train(
         return result
 
 
+@ray.remote
+def _map_predict(booster, part, columns, **kwargs):
+    """
+    Run prediction on remote worker.
+
+    Parameters
+    ----------
+    booster : xgboost.Booster or ray.ObjectRef
+        A trained booster.
+    part : pandas.DataFrame or ray.ObjectRef
+        Partition of full data used for local prediction.
+    columns : list or ray.ObjectRef
+        Columns for the result.
+    **kwargs : dict
+        Other parameters are the same as for ``xgboost.Booster.predict``.
+
+    Returns
+    -------
+    ray.ObjectRef
+        ``ray.ObjectRef`` with partial prediction.
+    """
+    dmatrix = xgb.DMatrix(part)
+    prediction = pandas.DataFrame(
+        booster.predict(dmatrix, **kwargs),
+        index=part.index,
+        columns=columns,
+    )
+    return prediction
+
+
 def _predict(
     booster,
     data,
-    num_actors,
     **kwargs,
 ):
     """
     Run distributed prediction with a trained booster on Ray backend.
 
-    During work it evenly distributes `data` between workers,
-    runs xgb.predict on each worker for subset of `data` and creates
-    Modin DataFrame with prediction results.
+    During execution it runs ``xgb.predict`` on each worker for subset of `data`
+    and creates Modin DataFrame with prediction results.
 
     Parameters
     ----------
@@ -667,11 +592,8 @@ def _predict(
         A trained booster.
     data : modin.experimental.xgboost.DMatrix
         Input data used for prediction.
-    num_actors : int, optional
-        Number of actors for prediction. If unspecified, this value will be
-        computed automatically.
-    **kwargs : dist
-        Other parameters are the same as `xgboost.Booster.predict`.
+    **kwargs : dict
+        Other parameters are the same as for ``xgboost.Booster.predict``.
 
     Returns
     -------
@@ -680,37 +602,39 @@ def _predict(
     """
     s = time.time()
 
-    X_row_parts, _ = data
-
-    num_actors = _get_num_actors(
-        num_actors if isinstance(num_actors, int) else "default_predict"
-    )
-
-    if num_actors > len(X_row_parts):
-        num_actors = len(X_row_parts)
-
-    # Create remote actors
-    actors = create_actors(num_actors)
+    # Get metadata from DMatrix
+    input_index, input_columns, row_lengths = data.metadata
 
-    # Split data across workers
-    _split_data_across_actors(
-        actors,
-        lambda actor, *X: actor.set_predict_data.remote(*X),
-        X_row_parts,
-        is_predict=True,
-    )
+    # Infer columns of result
+    def _get_num_columns(booster, n_features, **kwargs):
+        rng = np.random.RandomState(777)
+        test_data = rng.randn(1, n_features)
+        test_predictions = booster.predict(
+            xgb.DMatrix(test_data), validate_features=False, **kwargs
+        )
+        num_columns = (
+            test_predictions.shape[1] if len(test_predictions.shape) > 1 else 1
+        )
+        return num_columns
 
-    LOGGER.info(f"Data preparation time: {time.time() - s} s")
-    s = time.time()
+    result_num_columns = _get_num_columns(booster, len(input_columns), **kwargs)
+    new_columns = list(range(result_num_columns))
 
+    # Put common data in object store
     booster = ray.put(booster)
+    new_columns_ref = ray.put(new_columns)
 
-    predictions = [
-        tuple(actor.predict.options(num_returns=2).remote(booster, **kwargs))
-        for _, actor in actors
+    prediction_refs = [
+        _map_predict.remote(booster, part, new_columns_ref, **kwargs)
+        for _, part in data.data
     ]
-
-    result = from_partitions(predictions, 0)
+    predictions = from_partitions(
+        prediction_refs,
+        0,
+        index=input_index,
+        columns=new_columns,
+        row_lengths=row_lengths,
+        column_widths=[len(new_columns)],
+    )
     LOGGER.info(f"Prediction time: {time.time() - s} s")
-
-    return result
+    return predictions


Code review comments:
diff_hunk: @@ -58,54 +59,46 @@ internal implementation is the following:
    ..
 
 2. On this step, the parameter `num_actors` is processed. The internal function :py:func:`~modin.experimental.xgboost.xgboost_ray._get_num_actors`
-   examines the value provided by the user and checks if it fits in the set of expected values
-   (int, "default_train", "default_predict").
-
-   * int - `num_actors` won't be changed. This value will be used.
-   * "default_train" - `num_actors` will be computed using condition that 1 actor should use maximum 2 CPUs.
-     This condition was chosen for using maximum parallel workers with multithreaded XGBoost training (2 threads
-     per worker will be used in this case).
-   * "default_predict" - `num_actors` will be computed using condition that 1 actor should use maximum 8 CPUs.
-     This condition was chosen to combine parallelization techniques: parallel actors and parallel threads.
-
-.. note:: `num_actors` parameter is made available for public functions :py:func:`~modin.experimental.xgboost.train` and ``predict``
-  method of :py:class:`~modin.experimental.xgboost.Booster` class to allow fine-tuning for obtaining the best
-  performance in specific use cases.
-
-3. ``ray.util.placement_group`` is created to reserve all available Ray resources. After that
-   :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` objects are created using resources of the
-   previously created placement group.
-
-4. Data (`dtrain` for :py:func:`~modin.experimental.xgboost.xgboost_ray._train`, `data` for
-   :py:func:`~modin.experimental.xgboost.xgboost_ray._predict`) is split between actors evenly. The internal function
+   examines the value provided by the user. In case the value wasn't provided, the `num_actors` will be computed using condition that 1 actor should use maximum 2 CPUs.
+   This condition was chosen for using maximum parallel workers with multithreaded XGBoost training (2 threads
+   per worker will be used in this case).
+
+.. note:: `num_actors` parameter is made available for public function :py:func:`~modin.experimental.xgboost.train` to allow
+  fine-tuning for obtaining the best performance in specific use cases.
+
+3. :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` objects are created.
+
+4. Data `dtrain` is split between actors evenly. The internal function
    :py:func:`~modin.experimental.xgboost.xgboost_ray._split_data_across_actors` runs assigning row partitions to actors
    using internal function :py:func:`~modin.experimental.xgboost.xgboost_ray._assign_row_partitions_to_actors`.
-   This function creates a dictionary in the form: `{actor_rank: ([part_i0, part_i3, ..], [0, 3, ..]), ..}` for training,
-   `{actor_rank: [part_i0, part_i1, ..], ..}` for prediction.
+   This function creates a dictionary in the form: `{actor_rank: ([part_i0, part_i3, ..], [0, 3, ..]), ..}`.
 
 .. note:: :py:func:`~modin.experimental.xgboost.xgboost_ray._assign_row_partitions_to_actors` takes into account IP
   addresses of row partitions of `dtrain` data to minimize excess data transfer.
 
-1. For each :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` the object methods ``set_train_data`` or ``set_predict_data`` are
-   called remotely. Those methods run by loading row partitions in actor according to the dictionary with partitions
+5. For each :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` the object method ``set_train_data`` is
+   called remotely. This method runs loading row partitions in actor according to the dictionary with partitions
    distribution from previous step. When data is passed to the actor, the row partitions are automatically materialized
    (``ray.ObjectRef`` -> ``pandas.DataFrame``).
 
-2. Methods ``train`` or ``predict`` of :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` class object are called remotely.
+6. Method ``train`` of :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` class object is called remotely. This method
+   runs XGBoost training on local data of actor, connects to ``Rabit Tracker`` for sharing training state between
+   actors and returns dictionary with `booster` and `evaluation results`.
+
+7. On the final stage results from actors are returned. `booster` and `evals_result` is returned using ``ray.get``
+   function from remote actor.
+
+
+Prediction
+**********
 
-   * ``train``: method runs XGBoost training on local data of actor, connects to ``Rabit Tracker`` for sharing
-     training state between actors and returns dictionary with `booster` and `evaluation results`.
-   * ``predict``: method runs XGBoost prediction on local data of actor and returns IP address of actor and partial
-     prediction (``pandas.DataFrame``).
+1. The data is passed to :py:func:`~modin.experimental.xgboost.xgboost_ray._predict`
+   function as a :py:class:`~modin.experimental.xgboost.DMatrix` object.
 
-3. On the final stage results from actors are returned.
+2. The function :py:func:`~modin.experimental.xgboost.xgboost_ray._map_predict` is applied remotely for each partition
path: docs/flow/modin/experimental/xgboost.rst
text: @user1:
```suggestion
2. :py:func:`~modin.experimental.xgboost.xgboost_ray._map_predict` function is applied remotely for each partition
```

@author:
fixed

diff_hunk: @@ -58,54 +59,46 @@ internal implementation is the following:
    ..
 
 2. On this step, the parameter `num_actors` is processed. The internal function :py:func:`~modin.experimental.xgboost.xgboost_ray._get_num_actors`
-   examines the value provided by the user and checks if it fits in the set of expected values
-   (int, "default_train", "default_predict").
-
-   * int - `num_actors` won't be changed. This value will be used.
-   * "default_train" - `num_actors` will be computed using condition that 1 actor should use maximum 2 CPUs.
-     This condition was chosen for using maximum parallel workers with multithreaded XGBoost training (2 threads
-     per worker will be used in this case).
-   * "default_predict" - `num_actors` will be computed using condition that 1 actor should use maximum 8 CPUs.
-     This condition was chosen to combine parallelization techniques: parallel actors and parallel threads.
-
-.. note:: `num_actors` parameter is made available for public functions :py:func:`~modin.experimental.xgboost.train` and ``predict``
-  method of :py:class:`~modin.experimental.xgboost.Booster` class to allow fine-tuning for obtaining the best
-  performance in specific use cases.
-
-3. ``ray.util.placement_group`` is created to reserve all available Ray resources. After that
-   :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` objects are created using resources of the
-   previously created placement group.
-
-4. Data (`dtrain` for :py:func:`~modin.experimental.xgboost.xgboost_ray._train`, `data` for
-   :py:func:`~modin.experimental.xgboost.xgboost_ray._predict`) is split between actors evenly. The internal function
+   examines the value provided by the user. In case the value wasn't provided, the `num_actors` will be computed using condition that 1 actor should use maximum 2 CPUs.
+   This condition was chosen for using maximum parallel workers with multithreaded XGBoost training (2 threads
+   per worker will be used in this case).
+
+.. note:: `num_actors` parameter is made available for public function :py:func:`~modin.experimental.xgboost.train` to allow
+  fine-tuning for obtaining the best performance in specific use cases.
+
+3. :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` objects are created.
+
+4. Data `dtrain` is split between actors evenly. The internal function
    :py:func:`~modin.experimental.xgboost.xgboost_ray._split_data_across_actors` runs assigning row partitions to actors
    using internal function :py:func:`~modin.experimental.xgboost.xgboost_ray._assign_row_partitions_to_actors`.
-   This function creates a dictionary in the form: `{actor_rank: ([part_i0, part_i3, ..], [0, 3, ..]), ..}` for training,
-   `{actor_rank: [part_i0, part_i1, ..], ..}` for prediction.
+   This function creates a dictionary in the form: `{actor_rank: ([part_i0, part_i3, ..], [0, 3, ..]), ..}`.
 
 .. note:: :py:func:`~modin.experimental.xgboost.xgboost_ray._assign_row_partitions_to_actors` takes into account IP
   addresses of row partitions of `dtrain` data to minimize excess data transfer.
 
-1. For each :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` the object methods ``set_train_data`` or ``set_predict_data`` are
-   called remotely. Those methods run by loading row partitions in actor according to the dictionary with partitions
+5. For each :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` the object method ``set_train_data`` is
+   called remotely. This method runs loading row partitions in actor according to the dictionary with partitions
    distribution from previous step. When data is passed to the actor, the row partitions are automatically materialized
    (``ray.ObjectRef`` -> ``pandas.DataFrame``).
 
-2. Methods ``train`` or ``predict`` of :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` class object are called remotely.
+6. Method ``train`` of :py:class:`~modin.experimental.xgboost.xgboost_ray.ModinXGBoostActor` class object is called remotely. This method
+   runs XGBoost training on local data of actor, connects to ``Rabit Tracker`` for sharing training state between
+   actors and returns dictionary with `booster` and `evaluation results`.
+
+7. On the final stage results from actors are returned. `booster` and `evals_result` is returned using ``ray.get``
path: docs/flow/modin/experimental/xgboost.rst
text: @user1:
```suggestion
7. At the final stage results from actors are returned. `booster` and `evals_result` are returned using ``ray.get``
```

@author:
fixed