Repository: modin-project/modin

Issue:
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.




Discussion:

Resolving patch:
diff --git a/modin/experimental/xgboost/xgboost.py b/modin/experimental/xgboost/xgboost.py
--- 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
--- 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 the remote worker.
 
     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 a 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


Test patch:
diff --git a/modin/experimental/xgboost/test/test_xgboost.py b/modin/experimental/xgboost/test/test_xgboost.py
--- 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)