[data] support ray_remote_args and ray_remote_args_fn for preprocessor

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

Author: xingyu-long

python/ray/data/preprocessor.py

@@ -4,7 +4,7 @@
 import pickle
 import warnings
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, Union, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, Union, List, Optional, Callable
 
 from ray.air.util.data_batch_conversion import BatchFormat
 from ray.util.annotations import DeveloperAPI, PublicAPI
@@ -47,6 +47,24 @@ class Preprocessor(abc.ABC):
       implemented method.
     """
 
+    def __init__(
+        self,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = 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.
+        """
+        self._ray_remote_args = ray_remote_args
+        self._ray_remote_args_fn = ray_remote_args_fn
+
     class FitStatus(str, Enum):
         """The fit status of preprocessor."""
 
@@ -225,6 +243,11 @@ 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
+
         if transform_type == BatchFormat.PANDAS:
             return ds.map_batches(
                 self._transform_pandas, batch_format=BatchFormat.PANDAS, **kwargs

python/ray/data/preprocessors/tokenizer.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Dict, Any
 
 import pandas as pd
 
@@ -59,6 +59,13 @@ 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.
     """
 
     _is_fittable = False
@@ -68,7 +75,13 @@ 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,
     ):
+        super().__init__(
+            ray_remote_args=ray_remote_args,
+            ray_remote_args_fn=ray_remote_args_fn,
+        )
         self.columns = columns
         # TODO(matt): Add a more robust default tokenizer.
         self.tokenization_fn = tokenization_fn or simple_split_tokenizer