[data][llm] Add per-stage map kwargs for build_llm_processor preprocess/postprocess (#57826)

nrghosh · web-flow · commit 22c755dbef06 · 2025-10-17T19:36:25.000-07:00
Signed-off-by: Nikhil Ghosh &lt;nikhil@anyscale.com&gt;
diff --git a/python/ray/data/llm.py b/python/ray/data/llm.py
@@ -368,6 +368,8 @@ def build_llm_processor(
     config: ProcessorConfig,
     preprocess: Optional[UserDefinedFunction] = None,
     postprocess: Optional[UserDefinedFunction] = None,
+    preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
+    postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
     builder_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Processor:
     """Build a LLM processor using the given config.
@@ -383,6 +385,12 @@ def build_llm_processor(
         postprocess: An optional lambda function that takes a row (dict) as input
             and returns a postprocessed row (dict). To keep all the original columns,
             you can use the `**row` syntax to return all the original columns.
+        preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            preprocess stage. Useful for controlling resources (e.g., num_cpus=0.5)
+            and concurrency independently of the main LLM stage.
+        postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            postprocess stage. Useful for controlling resources (e.g., num_cpus=0.25)
+            and concurrency independently of the main LLM stage.
         builder_kwargs: Optional additional kwargs to pass to the processor builder
             function. These will be passed through to the registered builder and
             should match the signature of the specific builder being used.
@@ -435,6 +443,36 @@ def build_llm_processor(
             for row in ds.take_all():
                 print(row)
 
+        Using map_kwargs to control preprocess/postprocess resources:
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+            from ray.data.llm import vLLMEngineProcessorConfig, build_llm_processor
+
+            config = vLLMEngineProcessorConfig(
+                model_source="meta-llama/Meta-Llama-3.1-8B-Instruct",
+                concurrency=1,
+                batch_size=64,
+            )
+
+            processor = build_llm_processor(
+                config,
+                preprocess=lambda row: dict(
+                    messages=[{"role": "user", "content": row["prompt"]}],
+                    sampling_params=dict(temperature=0.3, max_tokens=20),
+                ),
+                postprocess=lambda row: dict(resp=row["generated_text"]),
+                preprocess_map_kwargs={"num_cpus": 0.5},
+                postprocess_map_kwargs={"num_cpus": 0.25},
+            )
+
+            ds = ray.data.range(300)
+            ds = processor(ds)
+            for row in ds.take_all():
+                print(row)
+
         Using builder_kwargs to pass chat_template_kwargs:
 
         .. testcode::
@@ -474,9 +512,12 @@ def build_llm_processor(
     from ray.llm._internal.batch.processor import ProcessorBuilder
 
     ProcessorBuilder.validate_builder_kwargs(builder_kwargs)
+
     build_kwargs = dict(
         preprocess=preprocess,
         postprocess=postprocess,
+        preprocess_map_kwargs=preprocess_map_kwargs,
+        postprocess_map_kwargs=postprocess_map_kwargs,
     )
 
     # Pass through any additional builder kwargs
diff --git a/python/ray/llm/_internal/batch/processor/base.py b/python/ray/llm/_internal/batch/processor/base.py
@@ -193,6 +193,10 @@ class Processor:
             required fields for the following processing stages.
         postprocess: An optional lambda function that takes a row (dict) as input
             and returns a postprocessed row (dict).
+        preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            preprocess stage (e.g., num_cpus, memory, concurrency).
+        postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            postprocess stage (e.g., num_cpus, memory, concurrency).
     """
 
     # The internal used data column name ("__data"). Your input
@@ -206,10 +210,14 @@ def __init__(
         stages: List[StatefulStage],
         preprocess: Optional[UserDefinedFunction] = None,
         postprocess: Optional[UserDefinedFunction] = None,
+        preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
+        postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
     ):
         self.config = config
         self.preprocess = None
         self.postprocess = None
+        self.preprocess_map_kwargs = preprocess_map_kwargs or {}
+        self.postprocess_map_kwargs = postprocess_map_kwargs or {}
         self.stages: OrderedDict[str, StatefulStage] = OrderedDict()
 
         # FIXES: https://github.com/ray-project/ray/issues/53124
@@ -251,7 +259,7 @@ def __call__(self, dataset: Dataset) -> Dataset:
             The output dataset.
         """
         if self.preprocess is not None:
-            dataset = dataset.map(self.preprocess)
+            dataset = dataset.map(self.preprocess, **self.preprocess_map_kwargs)
 
         # Apply stages.
         for stage in self.stages.values():
@@ -262,7 +270,7 @@ def __call__(self, dataset: Dataset) -> Dataset:
             dataset = dataset.map_batches(stage.fn, **kwargs)
 
         if self.postprocess is not None:
-            dataset = dataset.map(self.postprocess)
+            dataset = dataset.map(self.postprocess, **self.postprocess_map_kwargs)
         return dataset
 
     def _append_stage(self, stage: StatefulStage) -> None:
@@ -360,7 +368,12 @@ def validate_builder_kwargs(cls, builder_kwargs: Optional[Dict[str, Any]]) -> No
         """
         if builder_kwargs is not None:
             # Check for conflicts with explicitly passed arguments
-            reserved_keys = {"preprocess", "postprocess"}
+            reserved_keys = {
+                "preprocess",
+                "postprocess",
+                "preprocess_map_kwargs",
+                "postprocess_map_kwargs",
+            }
             conflicting_keys = reserved_keys & builder_kwargs.keys()
             if conflicting_keys:
                 raise ValueError(
diff --git a/python/ray/llm/_internal/batch/processor/http_request_proc.py b/python/ray/llm/_internal/batch/processor/http_request_proc.py
@@ -59,6 +59,8 @@ def build_http_request_processor(
     config: HttpRequestProcessorConfig,
     preprocess: Optional[UserDefinedFunction] = None,
     postprocess: Optional[UserDefinedFunction] = None,
+    preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
+    postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Processor:
     """Construct a Processor and configure stages.
 
@@ -69,6 +71,10 @@ def build_http_request_processor(
             required fields for the following processing stages.
         postprocess: An optional lambda function that takes a row (dict) as input
             and returns a postprocessed row (dict).
+        preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            preprocess stage (e.g., num_cpus, memory, concurrency).
+        postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            postprocess stage (e.g., num_cpus, memory, concurrency).
 
     Returns:
         The constructed processor.
@@ -100,6 +106,8 @@ def build_http_request_processor(
         stages,
         preprocess=preprocess,
         postprocess=postprocess,
+        preprocess_map_kwargs=preprocess_map_kwargs,
+        postprocess_map_kwargs=postprocess_map_kwargs,
     )
     return processor
 
diff --git a/python/ray/llm/_internal/batch/processor/serve_deployment_proc.py b/python/ray/llm/_internal/batch/processor/serve_deployment_proc.py
@@ -36,9 +36,10 @@ def build_serve_deployment_processor(
     config: ServeDeploymentProcessorConfig,
     preprocess: Optional[UserDefinedFunction] = None,
     postprocess: Optional[UserDefinedFunction] = None,
+    preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
+    postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Processor:
-    """
-    Construct a processor that runs a serve deployment.
+    """Construct a processor that runs a serve deployment.
 
     Args:
         config: The configuration for the processor.
@@ -47,6 +48,10 @@ def build_serve_deployment_processor(
             required fields for the following processing stages.
         postprocess: An optional lambda function that takes a row (dict) as input
             and returns a postprocessed row (dict).
+        preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            preprocess stage (e.g., num_cpus, memory, concurrency).
+        postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            postprocess stage (e.g., num_cpus, memory, concurrency).
 
     Returns:
         The constructed processor.
@@ -69,6 +74,8 @@ def build_serve_deployment_processor(
         stages,
         preprocess=preprocess,
         postprocess=postprocess,
+        preprocess_map_kwargs=preprocess_map_kwargs,
+        postprocess_map_kwargs=postprocess_map_kwargs,
     )
     return processor
 
diff --git a/python/ray/llm/_internal/batch/processor/sglang_engine_proc.py b/python/ray/llm/_internal/batch/processor/sglang_engine_proc.py
@@ -58,9 +58,12 @@ def build_sglang_engine_processor(
     chat_template_kwargs: Optional[Dict[str, Any]] = None,
     preprocess: Optional[UserDefinedFunction] = None,
     postprocess: Optional[UserDefinedFunction] = None,
+    preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
+    postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
     telemetry_agent: Optional[TelemetryAgent] = None,
 ) -> Processor:
     """Construct a Processor and configure stages.
+
     Args:
         config: The configuration for the processor.
         chat_template_kwargs: The optional kwargs to pass to apply_chat_template.
@@ -69,6 +72,10 @@ def build_sglang_engine_processor(
             required fields for the following processing stages.
         postprocess: An optional lambda function that takes a row (dict) as input
             and returns a postprocessed row (dict).
+        preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            preprocess stage (e.g., num_cpus, memory, concurrency).
+        postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            postprocess stage (e.g., num_cpus, memory, concurrency).
         telemetry_agent: An optional telemetry agent for collecting usage telemetry.
 
     Returns:
@@ -179,6 +186,8 @@ def build_sglang_engine_processor(
         stages,
         preprocess=preprocess,
         postprocess=postprocess,
+        preprocess_map_kwargs=preprocess_map_kwargs,
+        postprocess_map_kwargs=postprocess_map_kwargs,
     )
     return processor
 
diff --git a/python/ray/llm/_internal/batch/processor/vllm_engine_proc.py b/python/ray/llm/_internal/batch/processor/vllm_engine_proc.py
@@ -106,9 +106,12 @@ def build_vllm_engine_processor(
     chat_template_kwargs: Optional[Dict[str, Any]] = None,
     preprocess: Optional[UserDefinedFunction] = None,
     postprocess: Optional[UserDefinedFunction] = None,
+    preprocess_map_kwargs: Optional[Dict[str, Any]] = None,
+    postprocess_map_kwargs: Optional[Dict[str, Any]] = None,
     telemetry_agent: Optional[TelemetryAgent] = None,
 ) -> Processor:
     """Construct a Processor and configure stages.
+
     Args:
         config: The configuration for the processor.
         chat_template_kwargs: The optional kwargs to pass to apply_chat_template.
@@ -117,9 +120,12 @@ def build_vllm_engine_processor(
             required fields for the following processing stages.
         postprocess: An optional lambda function that takes a row (dict) as input
             and returns a postprocessed row (dict).
+        preprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            preprocess stage (e.g., num_cpus, memory, concurrency).
+        postprocess_map_kwargs: Optional kwargs to pass to Dataset.map() for the
+            postprocess stage (e.g., num_cpus, memory, concurrency).
         telemetry_agent: An optional telemetry agent for collecting usage telemetry.
 
-
     Returns:
         The constructed processor.
     """
@@ -262,6 +268,8 @@ def build_vllm_engine_processor(
         stages,
         preprocess=preprocess,
         postprocess=postprocess,
+        preprocess_map_kwargs=preprocess_map_kwargs,
+        postprocess_map_kwargs=postprocess_map_kwargs,
     )
     return processor
 
diff --git a/python/ray/llm/tests/batch/cpu/processor/test_processor_base.py b/python/ray/llm/tests/batch/cpu/processor/test_processor_base.py
