resolve comments

Signed-off-by: Yu Chen <yuchen.ecnu@gmail.com>

Author: yuchen-ecnu

python/ray/data/llm.py

@@ -14,6 +14,7 @@
 from ray.llm._internal.batch.stages.configs import (
     ChatTemplateStageConfig as _ChatTemplateStageConfig,
     DetokenizeStageConfig as _DetokenizeStageConfig,
+    HttpRequestStageConfig as _HttpRequestStageConfig,
     PrepareImageStageConfig as _PrepareImageStageConfig,
     PrepareMultimodalStageConfig as _PrepareMultimodalStageConfig,
     TokenizerStageConfig as _TokenizerStageConfig,
@@ -512,6 +513,29 @@ class TokenizerStageConfig(_TokenizerStageConfig):
     pass
 
 
+@PublicAPI(stability="alpha")
+class HttpRequestStageConfig(_HttpRequestStageConfig):
+    """The configuration for the http request stage.
+
+    Args:
+        enabled: Whether this stage is enabled. Defaults to True.
+        batch_size: Rows per batch. If not specified, will use the processor-level
+            batch_size.
+        concurrency: Actor pool size or range for this stage. If not specified,
+            will use the processor-level concurrency. If ``concurrency`` is a
+            tuple ``(m, n)``, Ray creates an autoscaling actor pool that scales
+            between ``m`` and ``n`` workers (``1 <= m <= n``). If ``concurrency``
+            is an ``int`` ``n``, CPU stages use an autoscaling pool from ``(1, n)``.
+        runtime_env: Optional runtime environment for this stage. If not specified,
+            will use the processor-level runtime_env. See
+            :ref:`this doc <handling_dependencies>` for more details.
+        num_cpus: Number of CPUs to reserve for each map worker in this stage.
+        memory: Heap memory in bytes to reserve for each map worker in this stage.
+    """
+
+    pass
+
+
 @PublicAPI(stability="alpha")
 class PrepareImageStageConfig(_PrepareImageStageConfig):
     """The configuration for the prepare image stage.
@@ -737,6 +761,7 @@ def build_processor(
     "DetokenizeStageConfig",
     "PrepareMultimodalStageConfig",
     "TokenizerStageConfig",
+    "HttpRequestStageConfig",
     "PrepareImageStageConfig",
     "build_llm_processor",
     "build_processor",

python/ray/llm/_internal/batch/processor/http_request_proc.py

@@ -4,7 +4,6 @@
 
 from pydantic import Field
 
-from ray.data import ActorPoolStrategy
 from ray.data.block import UserDefinedFunction
 from ray.llm._internal.batch.observability.usage_telemetry.usage import (
     BatchModelTelemetry,
@@ -15,8 +14,12 @@
     ProcessorBuilder,
     ProcessorConfig,
 )
-from ray.llm._internal.batch.processor.utils import extract_resource_kwargs
+from ray.llm._internal.batch.processor.utils import build_cpu_stage_map_kwargs
 from ray.llm._internal.batch.stages import HttpRequestStage
+from ray.llm._internal.batch.stages.configs import (
+    HttpRequestStageConfig,
+    resolve_stage_config,
+)
 
 
 class HttpRequestProcessorConfig(ProcessorConfig):
@@ -55,14 +58,9 @@ class HttpRequestProcessorConfig(ProcessorConfig):
         # exclude from JSON serialization since `session_factory` is a callable
         exclude=True,
     )
-    num_cpus: Optional[float] = Field(
-        default=None,
-        description="Number of CPUs per HttpRequestUDF worker. Defaults to 1 if None. "
-        "For I/O-bound workloads, use fractional values (e.g., 0.1).",
-    )
-    memory: Optional[float] = Field(
-        default=None,
-        description="Heap memory in bytes to reserve for each HttpRequestUDF worker.",
+    http_request_stage: Any = Field(
+        default=True,
+        description="Chat templating stage config (bool | dict | HttpRequestStageConfig).",
     )
 
 
@@ -90,6 +88,25 @@ def build_http_request_processor(
     Returns:
         The constructed processor.
     """
+
+    # Prepare processor defaults for merging into stage configs
+    processor_defaults = {
+        "batch_size": config.batch_size,
+        "concurrency": config.concurrency,
+    }
+
+    # Resolve and build HttpRequestStage if enabled
+    http_request_stage_cfg = resolve_stage_config(
+        config.http_request_stage,
+        HttpRequestStageConfig,
+        processor_defaults,
+    )
+
+    if not http_request_stage_cfg.enabled:
+        raise ValueError(
+            "The HTTP request stage is required and cannot be disabled in HttpRequestProcessorConfig."
+        )
+
     stages = [
         HttpRequestStage(
             fn_constructor_kwargs=dict(
@@ -100,16 +117,7 @@ def build_http_request_processor(
                 base_retry_wait_time_in_s=config.base_retry_wait_time_in_s,
                 session_factory=config.session_factory,
             ),
-            map_batches_kwargs=dict(
-                compute=ActorPoolStrategy(
-                    **config.get_concurrency(autoscaling_enabled=False),
-                ),
-                **extract_resource_kwargs(
-                    None,
-                    config.num_cpus,
-                    config.memory,
-                ),
-            ),
+            map_batches_kwargs=build_cpu_stage_map_kwargs(http_request_stage_cfg),
         )
     ]
     telemetry_agent = get_or_create_telemetry_agent()

python/ray/llm/_internal/batch/stages/configs.py

@@ -68,6 +68,10 @@ class PrepareMultimodalStageConfig(_StageConfigBase):
     )
 
 
+class HttpRequestStageConfig(_StageConfigBase):
+    pass
+
+
 def resolve_stage_config(
     stage_cfg_value: Union[bool, Dict[str, Any], _StageConfigBase],
     stage_config_cls: Type[T],

python/ray/llm/tests/batch/cpu/processor/test_http_request_proc.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from ray.data.llm import HttpRequestStageConfig
 from ray.llm._internal.batch.processor import ProcessorBuilder
 from ray.llm._internal.batch.processor.http_request_proc import (
     HttpRequestProcessorConfig,
@@ -15,15 +16,17 @@ def test_http_request_processor():
         qps=2,
         concurrency=4,
         batch_size=64,
-        num_cpus=0.5,
-        memory=100000,
+        http_request_stage=HttpRequestStageConfig(
+            num_cpus=0.5,
+            memory=100000,
+        ),
     )
     processor = ProcessorBuilder.build(config)
     assert processor.list_stage_names() == ["HttpRequestStage"]
     stage = processor.get_stage_by_name("HttpRequestStage")
     assert stage.map_batches_kwargs["num_cpus"] == 0.5
     assert stage.map_batches_kwargs["memory"] == 100000
-    assert stage.map_batches_kwargs["compute"].min_size == 4
+    assert stage.map_batches_kwargs["compute"].min_size == 1
     assert stage.map_batches_kwargs["compute"].max_size == 4
     assert stage.fn_constructor_kwargs["url"] == "http://localhost:8000"
     assert stage.fn_constructor_kwargs["additional_header"] == {