[data] use nums_cpus, memory, batch_size, concurrency

Signed-off-by: Xingyu Long <xingyulong97@gmail.com>

Author: xingyu-long

python/ray/data/preprocessor.py

@@ -4,7 +4,17 @@
 import pickle
 import warnings
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, Union, List, Optional, Callable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Union,
+    List,
+    Optional,
+    Callable,
+    Literal,
+    Tuple,
+)
 
 from ray.air.util.data_batch_conversion import BatchFormat
 from ray.util.annotations import DeveloperAPI, PublicAPI
@@ -49,21 +59,22 @@ class Preprocessor(abc.ABC):
 
     def __init__(
         self,
-        ray_remote_args: Optional[Dict[str, Any]] = None,
-        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        num_cpus: Optional[float] = None,
+        memory: Optional[float] = None,
+        batch_size: Union[int, None, Literal["default"]] = None,
+        concurrency: Optional[int] = None,
     ):
         """
         Args:
-            ray_remote_args: Args to provide to :func:`ray.remote`.
-            ray_remote_args_fn: A function that returns a dictionary of remote args
-                passed to each map worker. The purpose of this argument is to generate
-                dynamic arguments for each actor/task, and will be called each time
-                prior to initializing the worker. Args returned from this dict will
-                always override the args in ``ray_remote_args``. Note: this is an
-                advanced, experimental feature.
+            num_cpus: The number of CPUs to reserve for each parallel map worker.
+            memory: The heap memory in bytes to reserve for each parallel map worker.
+            batch_size: The maximum number of rows to return.
+            concurrency: The maximum number of Ray workers to use concurrently.
         """
-        self._ray_remote_args = ray_remote_args
-        self._ray_remote_args_fn = ray_remote_args_fn
+        self._num_cpus = num_cpus
+        self._memory = memory
+        self._batch_size = batch_size
+        self._concurrency = concurrency
 
     class FitStatus(str, Enum):
         """The fit status of preprocessor."""
@@ -243,10 +254,10 @@ def _transform(self, ds: "Dataset") -> "Dataset":
         # Our user-facing batch format should only be pandas or NumPy, other
         # formats {arrow, simple} are internal.
         kwargs = self._get_transform_config()
-        if self._ray_remote_args is not None:
-            kwargs = dict(kwargs, **self._ray_remote_args)
-        if self._ray_remote_args_fn is not None:
-            kwargs["ray_remote_args_fn"] = self._ray_remote_args_fn
+        kwargs["num_cpus"] = self._num_cpus
+        kwargs["memory"] = self._memory
+        kwargs["batch_size"] = self._batch_size
+        kwargs["concurrency"] = self._concurrency
 
         if transform_type == BatchFormat.PANDAS:
             return ds.map_batches(

python/ray/data/preprocessors/tokenizer.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Optional, Dict, Any
+from typing import Callable, List, Optional, Literal, Union
 
 import pandas as pd
 
@@ -59,13 +59,10 @@ class Tokenizer(Preprocessor):
             columns will be the same as the input columns. If not None, the length of
             ``output_columns`` must match the length of ``columns``, othwerwise an error
             will be raised.
-        ray_remote_args: Args to provide to :func:`ray.remote`.
-        ray_remote_args_fn: A function that returns a dictionary of remote args
-            passed to each map worker. The purpose of this argument is to generate
-            dynamic arguments for each actor/task, and will be called each time
-            prior to initializing the worker. Args returned from this dict will
-            always override the args in ``ray_remote_args``. Note: this is an
-            advanced, experimental feature.
+        num_cpus: The number of CPUs to reserve for each parallel map worker.
+        memory: The heap memory in bytes to reserve for each parallel map worker.
+        batch_size: The maximum number of rows to return.
+        concurrency: The maximum number of Ray workers to use concurrently.
     """
 
     _is_fittable = False
@@ -75,12 +72,17 @@ def __init__(
         columns: List[str],
         tokenization_fn: Optional[Callable[[str], List[str]]] = None,
         output_columns: Optional[List[str]] = None,
-        ray_remote_args: Optional[Dict[str, Any]] = None,
-        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        *,
+        num_cpus: Optional[float] = None,
+        memory: Optional[float] = None,
+        batch_size: Union[int, None, Literal["default"]] = None,
+        concurrency: Optional[int] = None,
     ):
         super().__init__(
-            ray_remote_args=ray_remote_args,
-            ray_remote_args_fn=ray_remote_args_fn,
+            num_cpus=num_cpus,
+            memory=memory,
+            batch_size=batch_size,
+            concurrency=concurrency,
         )
         self.columns = columns
         # TODO(matt): Add a more robust default tokenizer.

python/ray/data/tests/preprocessors/test_preprocessors.py

@@ -165,17 +165,9 @@ def test_fit_twice(mocked_warn):
     mocked_warn.assert_called_once_with(msg)
 
 
-def test_ray_remote_args_and_fn():
+def test_initialization_parameters():
     batch_size = 2
 
-    ray_remote_args = {"num_cpus": 2}
-
-    def func(df):
-        import os
-
-        df["value"][:] = int(os.environ["__MY_TEST__"])
-        return df
-
     class DummyPreprocessor(Preprocessor):
         _is_fittable = False
 
@@ -185,23 +177,23 @@ def _get_transform_config(self):
         def _transform_numpy(self, data):
             assert (
                 ray.get_runtime_context().get_assigned_resources()["CPU"]
-                == ray_remote_args["num_cpus"]
+                == self._num_cpus
             )
             assert len(data["value"]) == batch_size
-            func(data)
             return data
 
         def _determine_transform_to_use(self):
             return "numpy"
 
     prep = DummyPreprocessor(
-        ray_remote_args=ray_remote_args,
-        ray_remote_args_fn=lambda: {"runtime_env": {"env_vars": {"__MY_TEST__": "69"}}},
+        num_cpus=2,
+        concurrency=2,
+        batch_size=batch_size,
     )
     ds = ray.data.from_pandas(pd.DataFrame({"value": list(range(10))}))
     ds = prep.transform(ds)
 
-    assert sorted([x["value"] for x in ds.take(5)]) == [69, 69, 69, 69, 69]
+    assert [x["value"] for x in ds.take(5)] == [0, 1, 2, 3, 4]
 
 
 def test_transform_config():