chore: import upstream snapshot with attribution

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

@@ -0,0 +1,99 @@
+import sys
+import warnings
+
+import pytest
+
+from ray.llm._internal.batch.processor.vllm_engine_proc import vLLMEngineProcessorConfig
+from ray.llm._internal.batch.stages.configs import (
+    ChatTemplateStageConfig,
+    DetokenizeStageConfig,
+    PrepareMultimodalStageConfig,
+    TokenizerStageConfig,
+)
+
+
+def test_legacy_booleans_coerced_to_stage_configs():
+    """Legacy flags → stage configs (dict form)."""
+    config = vLLMEngineProcessorConfig(
+        model_source="test-model",
+        apply_chat_template=True,
+        tokenize=False,
+        detokenize=True,
+    )
+
+    # Legacy flags should be coerced to stage configs
+    assert isinstance(config.chat_template_stage, dict)
+    assert config.chat_template_stage["enabled"] is True
+
+    assert isinstance(config.tokenize_stage, dict)
+    assert config.tokenize_stage["enabled"] is False
+
+    assert isinstance(config.detokenize_stage, dict)
+    assert config.detokenize_stage["enabled"] is True
+
+
+def test_explicit_stage_configs_preserved():
+    """Explicit stage configs not overwritten by legacy flags."""
+    explicit_chat_template = ChatTemplateStageConfig(enabled=False, batch_size=64)
+    config = vLLMEngineProcessorConfig(
+        model_source="test-model",
+        chat_template_stage=explicit_chat_template,
+        apply_chat_template=True,  # Legacy flag should be ignored
+    )
+
+    # Explicit stage config should be preserved
+    assert config.chat_template_stage is explicit_chat_template
+    assert config.chat_template_stage.enabled is False
+    assert config.chat_template_stage.batch_size == 64
+
+
+def test_chat_template_fields_merged():
+    """apply_chat_template + chat_template → merged into stage config."""
+    config = vLLMEngineProcessorConfig(
+        model_source="test-model",
+        apply_chat_template=True,
+        chat_template="custom_template",
+    )
+
+    assert isinstance(config.chat_template_stage, dict)
+    assert config.chat_template_stage["enabled"] is True
+    assert config.chat_template_stage["chat_template"] == "custom_template"
+
+
+def test_no_warnings_when_using_new_api():
+    """No warnings when only new API used."""
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        vLLMEngineProcessorConfig(
+            model_source="test-model",
+            chat_template_stage=ChatTemplateStageConfig(enabled=True),
+            tokenize_stage=TokenizerStageConfig(enabled=True),
+            detokenize_stage=DetokenizeStageConfig(enabled=True),
+            prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=False),
+        )
+        # Filter out any non-UserWarning warnings
+        deprecation_warnings = [
+            warning for warning in w if issubclass(warning.category, UserWarning)
+        ]
+        assert len(deprecation_warnings) == 0
+
+
+def test_legacy_dict_stage_config():
+    """Dict form stage configs work correctly."""
+    config = vLLMEngineProcessorConfig(
+        model_source="test-model",
+        chat_template_stage={"enabled": False, "batch_size": 128},
+        tokenize_stage={"enabled": True, "concurrency": 4},
+    )
+
+    assert isinstance(config.chat_template_stage, dict)
+    assert config.chat_template_stage["enabled"] is False
+    assert config.chat_template_stage["batch_size"] == 128
+
+    assert isinstance(config.tokenize_stage, dict)
+    assert config.tokenize_stage["enabled"] is True
+    assert config.tokenize_stage["concurrency"] == 4
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,39 @@
+import sys
+
+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,
+)
+
+
+def test_http_request_processor():
+    config = HttpRequestProcessorConfig(
+        url="http://localhost:8000",
+        headers={"Authorization": "Bearer 1234567890"},
+        qps=2,
+        concurrency=4,
+        batch_size=64,
+        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 == 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"] == {
+        "Authorization": "Bearer 1234567890"
+    }
+    assert stage.fn_constructor_kwargs["qps"] == 2
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,177 @@
+"""Regression tests for lazy imports in ``ray.llm._internal.batch``.
+
+The ``stages`` and ``processor`` packages re-export classes whose defining
+modules pull in heavy optional dependencies (``transformers``, ``vllm``,
+``sglang``, ``mistral_common``). They are wired up via PEP 562
+``__getattr__`` so that, e.g., importing ``HttpRequestProcessorConfig`` does
+not drag the entire ML stack into ``sys.modules``.
+
+These tests run in a fresh Python subprocess so that ``sys.modules`` is
+guaranteed to be clean -- otherwise the modules-under-test could already be
+loaded by an earlier test.
+"""
+
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+
+def _run_in_subprocess(script: str) -> str:
+    """Run ``script`` in a clean Python subprocess and return stdout."""
+    result = subprocess.run(
+        [sys.executable, "-c", textwrap.dedent(script)],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return result.stdout
+
+
+# Module names that the lightweight HTTP processor must NOT pull in.
+# Each one is loaded by a different stage / processor module, so seeing any
+# of them after importing the HTTP processor means the lazy wiring broke.
+_HEAVY_MODULES = (
+    "transformers",
+    "tokenizers",
+    "huggingface_hub",
+    "mistral_common",
+    "vllm.transformers_utils",
+    # Stage submodules that pull in the heavy deps above.
+    "ray.llm._internal.batch.stages.tokenize_stage",
+    "ray.llm._internal.batch.stages.chat_template_stage",
+    "ray.llm._internal.batch.stages.vllm_engine_stage",
+    "ray.llm._internal.batch.stages.sglang_engine_stage",
+    "ray.llm._internal.batch.stages.prepare_multimodal_stage",
+    "ray.llm._internal.batch.stages.serve_deployment_stage",
+    # Processor submodules whose top-level statements import heavy deps
+    # directly (e.g. ``import transformers`` in sglang_engine_proc.py).
+    "ray.llm._internal.batch.processor.sglang_engine_proc",
+    "ray.llm._internal.batch.processor.vllm_engine_proc",
+    "ray.llm._internal.batch.processor.serve_deployment_proc",
+)
+
+
+def test_http_request_processor_does_not_import_heavy_deps():
+    """HTTP-only processor imports must not load transformers/vllm/etc."""
+    out = _run_in_subprocess(
+        f"""
+        import sys
+        from ray.llm._internal.batch import HttpRequestProcessorConfig  # noqa: F401
+
+        heavy = {_HEAVY_MODULES!r}
+        loaded = [m for m in heavy if m in sys.modules]
+        print(','.join(loaded))
+        """
+    )
+    loaded = [m for m in out.strip().split(",") if m]
+    assert loaded == [], (
+        "Importing HttpRequestProcessorConfig must not pull in heavy ML "
+        f"dependencies, but the following modules ended up loaded: {loaded}"
+    )
+
+
+def test_http_request_stage_only_loads_its_own_submodule():
+    """``from ...stages import HttpRequestStage`` must only load that stage."""
+    out = _run_in_subprocess(
+        """
+        import sys
+        from ray.llm._internal.batch.stages import HttpRequestStage  # noqa: F401
+
+        stage_modules = sorted(
+            m for m in sys.modules
+            if m.startswith('ray.llm._internal.batch.stages.')
+            and m != 'ray.llm._internal.batch.stages.base'
+            and m != 'ray.llm._internal.batch.stages.configs'
+            and m != 'ray.llm._internal.batch.stages.common'
+        )
+        print(','.join(stage_modules))
+        """
+    )
+    loaded = [m for m in out.strip().split(",") if m]
+    assert loaded == [
+        "ray.llm._internal.batch.stages.http_request_stage"
+    ], f"Expected only http_request_stage to load, got: {loaded}"
+
+
+@pytest.mark.parametrize(
+    "name,submodule",
+    [
+        ("HttpRequestStage", "http_request_stage"),
+        ("TokenizeStage", "tokenize_stage"),
+        ("DetokenizeStage", "tokenize_stage"),
+        ("ChatTemplateStage", "chat_template_stage"),
+        ("PrepareMultimodalStage", "prepare_multimodal_stage"),
+        ("ServeDeploymentStage", "serve_deployment_stage"),
+        ("SGLangEngineStage", "sglang_engine_stage"),
+        ("vLLMEngineStage", "vllm_engine_stage"),
+    ],
+)
+def test_stage_lazy_attr_resolves(name, submodule):
+    """Each lazy stage attr resolves to the class from the right submodule."""
+    import ray.llm._internal.batch.stages as stages
+
+    cls = getattr(stages, name)
+    assert cls.__name__ == name
+    assert cls.__module__ == f"ray.llm._internal.batch.stages.{submodule}"
+
+
+@pytest.mark.parametrize(
+    "name,submodule",
+    [
+        ("HttpRequestProcessorConfig", "http_request_proc"),
+        ("ServeDeploymentProcessorConfig", "serve_deployment_proc"),
+        ("SGLangEngineProcessorConfig", "sglang_engine_proc"),
+        ("vLLMEngineProcessorConfig", "vllm_engine_proc"),
+    ],
+)
+def test_processor_lazy_attr_resolves(name, submodule):
+    """Each lazy processor-config attr resolves to the right class."""
+    import ray.llm._internal.batch.processor as processor
+
+    cls = getattr(processor, name)
+    assert cls.__name__ == name
+    assert cls.__module__ == f"ray.llm._internal.batch.processor.{submodule}"
+
+
+def test_unknown_attr_raises_attribute_error():
+    """``__getattr__`` must raise ``AttributeError`` for unknown names so
+    that ``hasattr`` and other attribute-introspection paths behave correctly.
+    """
+    import ray.llm._internal.batch.processor as processor
+    import ray.llm._internal.batch.stages as stages
+
+    with pytest.raises(AttributeError):
+        processor.DefinitelyNotAProcessor  # noqa: B018
+    with pytest.raises(AttributeError):
+        stages.DefinitelyNotAStage  # noqa: B018
+
+
+def test_dir_lists_lazy_attrs():
+    """``dir(pkg)`` must list the lazy attributes (for IDE completion etc.)."""
+    import ray.llm._internal.batch.processor as processor
+    import ray.llm._internal.batch.stages as stages
+
+    for name in (
+        "HttpRequestProcessorConfig",
+        "ServeDeploymentProcessorConfig",
+        "SGLangEngineProcessorConfig",
+        "vLLMEngineProcessorConfig",
+    ):
+        assert name in dir(processor)
+    for name in (
+        "HttpRequestStage",
+        "TokenizeStage",
+        "DetokenizeStage",
+        "ChatTemplateStage",
+        "PrepareMultimodalStage",
+        "ServeDeploymentStage",
+        "SGLangEngineStage",
+        "vLLMEngineStage",
+    ):
+        assert name in dir(stages)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-vv", __file__]))

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

@@ -0,0 +1,595 @@
+import sys
+from typing import Any, AsyncIterator, Dict, List, Type
+from unittest.mock import patch
+
+import pydantic
+import pytest
+
+import ray
+from ray.data.llm import build_processor
+from ray.llm._internal.batch.processor import (
+    base as processor_base,
+    vLLMEngineProcessorConfig,
+)
+from ray.llm._internal.batch.processor.base import (
+    Processor,
+    ProcessorBuilder,
+    ProcessorConfig,
+)
+from ray.llm._internal.batch.stages.base import StatefulStage, StatefulStageUDF
+
+
+def test_empty_processor():
+    """Test processor with only preprocess and postprocess."""
+
+    processor = Processor(
+        config=ProcessorConfig(
+            batch_size=64,
+            accelerator_type=None,
+            concurrency=1,
+        ),
+        stages=[],
+        # {id} -> {__data: {id, val}}
+        preprocess=lambda row: {"val": row["id"] + 5},
+        # {__data: {id, val}} -> {id, result}
+        postprocess=lambda row: {"result": row["val"], "id": row["id"]},
+    )
+
+    ds = ray.data.range(5)
+    ds = processor(ds).take_all()
+    for row in ds:
+        assert "val" not in row
+        assert "id" in row
+        assert "result" in row
+
+
+def test_processor_with_no_preprocess_or_postprocess():
+    """Test processor with no preprocess or postprocess."""
+
+    processor = Processor(
+        config=ProcessorConfig(
+            batch_size=64,
+            accelerator_type=None,
+            concurrency=1,
+        ),
+        stages=[],
+    )
+
+    ds = ray.data.range(5)
+    ds = processor(ds).take_all()
+    for row in ds:
+        assert "id" in row
+
+
+@pytest.mark.parametrize("has_extra", [True, False])
+def test_processor_with_stages(has_extra: bool):
+    """Test processor with multiple stages."""
+
+    class DummyStatefulStageUDF(StatefulStageUDF):
+        def __init__(
+            self,
+            data_column: str,
+            expected_input_keys: List[str],
+            factor: int,
+        ):
+            super().__init__(data_column, expected_input_keys)
+            self.factor = factor
+
+        async def udf(
+            self, batch: List[Dict[str, Any]]
+        ) -> AsyncIterator[Dict[str, Any]]:
+            for row in batch:
+                answer = row["val"] * self.factor
+                if "extra" in row:  # Optional input column.
+                    answer += row["extra"]
+                yield {
+                    # Use the same name to chain multiple dummy stages.
+                    "val": answer,
+                    self.IDX_IN_BATCH_COLUMN: row[self.IDX_IN_BATCH_COLUMN],
+                }
+
+    class DummyStage(StatefulStage):
+        fn: Type[StatefulStageUDF] = DummyStatefulStageUDF
+        fn_constructor_kwargs: Dict[str, Any] = {}
+        map_batches_kwargs: Dict[str, Any] = dict(concurrency=1)
+
+        def get_required_input_keys(self) -> Dict[str, str]:
+            return {"val": "The value to multiply."}
+
+    stages = [
+        DummyStage(fn_constructor_kwargs=dict(factor=2)),
+        DummyStage(fn_constructor_kwargs=dict(factor=3)),
+    ]
+
+    processor = Processor(
+        config=ProcessorConfig(
+            accelerator_type=None,
+            concurrency=1,
+            batch_size=64,
+        ),
+        stages=stages,
+        preprocess=lambda row: {"val": row["id"]},
+        postprocess=lambda row: {"result": row["val"], "id": row["id"]},
+    )
+
+    # Check the stage names.
+    stage_names = processor.list_stage_names()
+    assert stage_names == [
+        "DummyStage",
+        "DummyStage_1",
+    ]
+
+    # Check the stages.
+    for stage_name, stage in zip(stage_names, stages):
+        assert processor.get_stage_by_name(stage_name) == stage
+
+    # Run the processor twice with different datasets to test
+    # whether the processor is reusable.
+    for _ in range(2):
+        ds = ray.data.range(5)
+        ds = ds.map(
+            lambda row: {
+                "id": row["id"],
+                **({"extra": 1} if has_extra else {}),
+            }
+        )
+
+        ds = processor(ds).take_all()
+        extra = 1 if has_extra else 0
+        for row in ds:
+            assert "id" in row
+            assert "result" in row
+
+            # The final output should be the result of the last stage.
+            assert row["result"] == (row["id"] * 2 + extra) * 3 + extra
+
+
+# Common dummy classes for testing
+class DummyStatefulStageUDF(StatefulStageUDF):
+    async def udf(self, batch: List[Dict[str, Any]]) -> AsyncIterator[Dict[str, Any]]:
+        for row in batch:
+            yield row
+
+
+class DummyStage(StatefulStage):
+    fn: Type[StatefulStageUDF] = DummyStatefulStageUDF
+    fn_constructor_kwargs: Dict[str, Any] = {}
+    map_batches_kwargs: Dict[str, Any] = {}
+
+
+class DummyProcessorConfig(ProcessorConfig):
+    pass
+
+
+def test_builder():
+    def build_processor(config: ProcessorConfig) -> Processor:
+        stages = [
+            DummyStage(
+                fn_constructor_kwargs=dict(),
+                map_batches_kwargs=dict(concurrency=1),
+            )
+        ]
+        processor = Processor(config, stages)
+        return processor
+
+    ProcessorBuilder.register(DummyProcessorConfig, build_processor)
+
+    processor = ProcessorBuilder.build(DummyProcessorConfig(batch_size=64))
+    assert isinstance(processor.config, DummyProcessorConfig)
+    assert processor.list_stage_names() == ["DummyStage"]
+    assert (
+        processor.get_stage_by_name("DummyStage").map_batches_kwargs["concurrency"] == 1
+    )
+
+    def overrider(name: str, stage: StatefulStage):
+        if name.startswith("DummyStage"):
+            stage.map_batches_kwargs["concurrency"] = 2
+
+    processor = ProcessorBuilder.build(
+        DummyProcessorConfig(batch_size=64),
+        override_stage_config_fn=overrider,
+    )
+    assert processor.list_stage_names() == ["DummyStage"]
+    assert (
+        processor.get_stage_by_name("DummyStage").map_batches_kwargs["concurrency"] == 2
+    )
+
+
+class TestBuilderKwargsValidation:
+    @pytest.fixture
+    def build_processor_with_kwargs(self):
+        def build_processor_with_kwargs(
+            config: ProcessorConfig,
+            preprocess=None,
+            postprocess=None,
+            preprocess_map_kwargs=None,
+            postprocess_map_kwargs=None,
+            custom_kwarg=None,
+            another_kwarg=None,
+        ) -> Processor:
+            stages = [
+                DummyStage(
+                    fn_constructor_kwargs=dict(
+                        custom_kwarg=custom_kwarg,
+                        another_kwarg=another_kwarg,
+                    ),
+                    map_batches_kwargs=dict(concurrency=1),
+                )
+            ]
+            processor = Processor(
+                config,
+                stages,
+                preprocess=preprocess,
+                postprocess=postprocess,
+                preprocess_map_kwargs=preprocess_map_kwargs,
+                postprocess_map_kwargs=postprocess_map_kwargs,
+            )
+            return processor
+
+        return build_processor_with_kwargs
+
+    @pytest.fixture(autouse=True)
+    def clear_registry(self):
+        ProcessorBuilder.clear_registry()
+
+    def test_builder_kwargs_passthrough(self, build_processor_with_kwargs):
+        ProcessorBuilder.register(DummyProcessorConfig, build_processor_with_kwargs)
+
+        config = DummyProcessorConfig(batch_size=64)
+        processor = build_processor(
+            config,
+            preprocess=lambda row: {"val": row["id"]},
+            postprocess=lambda row: {"result": row["val"]},
+            builder_kwargs=dict(
+                custom_kwarg="test_value",
+                another_kwarg=42,
+            ),
+        )
+        assert processor.list_stage_names() == ["DummyStage"]
+        stage = processor.get_stage_by_name("DummyStage")
+        assert stage.fn_constructor_kwargs["custom_kwarg"] == "test_value"
+        assert stage.fn_constructor_kwargs["another_kwarg"] == 42
+
+    def test_unsupported_kwargs(self):
+        def build_processor_no_kwargs(
+            config: ProcessorConfig,
+            preprocess=None,
+            postprocess=None,
+        ) -> Processor:
+            stages = []
+            processor = Processor(
+                config, stages, preprocess=preprocess, postprocess=postprocess
+            )
+            return processor
+
+        ProcessorBuilder.register(DummyProcessorConfig, build_processor_no_kwargs)
+
+        config = DummyProcessorConfig(batch_size=64)
+        with pytest.raises(TypeError, match="unsupported_kwarg"):
+            build_processor(
+                config,
+                builder_kwargs=dict(unsupported_kwarg="value"),
+            )
+
+    @pytest.mark.parametrize("conflicting_key", ["preprocess", "postprocess"])
+    def test_error_builder_kwargs_conflict(
+        self, conflicting_key, build_processor_with_kwargs
+    ):
+        ProcessorBuilder.register(DummyProcessorConfig, build_processor_with_kwargs)
+
+        config = DummyProcessorConfig(batch_size=64)
+        with pytest.raises(ValueError, match="builder_kwargs cannot contain"):
+            build_processor(
+                config,
+                preprocess=lambda row: {"val": row["id"]},
+                builder_kwargs={conflicting_key: lambda row: {"other": row["id"]}},
+            )
+
+
+class TestProcessorConfig:
+    def test_valid_concurrency(self):
+        config = vLLMEngineProcessorConfig(
+            model_source="unsloth/Llama-3.2-1B-Instruct",
+            concurrency=(1, 2),
+        )
+        assert config.concurrency == (1, 2)
+
+        config = vLLMEngineProcessorConfig(
+            model_source="unsloth/Llama-3.2-1B-Instruct",
+        )
+        assert config.concurrency == 1
+
+    def test_invalid_concurrency(self):
+        with pytest.raises(pydantic.ValidationError):
+            vLLMEngineProcessorConfig(
+                model_source="unsloth/Llama-3.2-1B-Instruct",
+                concurrency=1.1,
+            )
+
+        with pytest.raises(pydantic.ValidationError):
+            vLLMEngineProcessorConfig(
+                model_source="unsloth/Llama-3.2-1B-Instruct",
+                concurrency=[1, 2, 3],
+            )
+
+    @pytest.mark.parametrize("n", [1, 2, 10])
+    def test_positive_int_not_fail(self, n):
+        conf = ProcessorConfig(concurrency=n)
+        assert conf.concurrency == n
+
+    def test_positive_int_unusual_not_fail(self):
+        assert ProcessorConfig(concurrency="1").concurrency == 1
+        assert ProcessorConfig(concurrency=1.0).concurrency == 1
+        assert ProcessorConfig(concurrency="1.0").concurrency == 1
+
+    @pytest.mark.parametrize("pair", [(1, 1), (1, 2), (2, 8)])
+    def test_valid_tuple_not_fail(self, pair):
+        conf = ProcessorConfig(concurrency=pair)
+        assert conf.concurrency == pair
+
+    def test_valid_tuple_unusual_not_fail(self):
+        assert ProcessorConfig(concurrency=("1", 2)).concurrency == (1, 2)
+        assert ProcessorConfig(concurrency=(1, "2")).concurrency == (1, 2)
+        assert ProcessorConfig(concurrency=[1, "2"]).concurrency == (1, 2)
+
+    @pytest.mark.parametrize(
+        "bad,msg_part",
+        [
+            (0, "positive integer"),
+            (-5, "positive integer"),
+            ((1, 2, 3), "at most 2 items"),
+            ((0, 1), "positive integers"),
+            ((1, 0), "positive integers"),
+            ((-1, 2), "positive integers"),
+            ((1, -2), "positive integers"),
+            ((1, 2.5), "a number with a fractional part"),
+            ("2.1", "unable to parse string"),
+            ((5, 2), "min > max"),
+        ],
+    )
+    def test_invalid_inputs_raise(self, bad, msg_part):
+        with pytest.raises(pydantic.ValidationError) as e:
+            ProcessorConfig(concurrency=bad)
+        assert msg_part in str(e.value)
+
+    @pytest.mark.parametrize(
+        "n,expected",
+        [
+            (1, {"min_size": 1, "max_size": 1}),
+            (4, {"min_size": 1, "max_size": 4}),
+            (10, {"min_size": 1, "max_size": 10}),
+            ("10", {"min_size": 1, "max_size": 10}),
+        ],
+    )
+    def test_with_int_concurrency_scaling(self, n, expected):
+        conf = ProcessorConfig(concurrency=n)
+        assert conf.get_concurrency() == expected
+
+    @pytest.mark.parametrize(
+        "n,expected",
+        [
+            (1, {"size": 1}),
+            (4, {"size": 4}),
+            (10, {"size": 10}),
+        ],
+    )
+    def test_with_int_concurrency_fixed(self, n, expected):
+        conf = ProcessorConfig(concurrency=n)
+        assert conf.get_concurrency(autoscaling_enabled=False) == expected
+
+    @pytest.mark.parametrize(
+        "pair,expected",
+        [
+            ((1, 1), {"min_size": 1, "max_size": 1}),
+            ((1, 3), {"min_size": 1, "max_size": 3}),
+            ((2, 8), {"min_size": 2, "max_size": 8}),
+        ],
+    )
+    def test_with_tuple_concurrency(self, pair, expected):
+        conf = ProcessorConfig(concurrency=pair)
+        assert conf.get_concurrency() == expected
+
+
+class TestOfflineProcessorConfig:
+    @pytest.mark.parametrize(
+        "kwargs, expected",
+        [
+            ({"max_tasks_in_flight_per_actor": 10}, 10),
+            ({}, None),
+            # Field stays None; the formula runs in Ray Data, not here.
+            ({"max_concurrent_batches": 4}, None),
+        ],
+    )
+    def test_max_tasks_in_flight_per_actor_passthrough(self, kwargs, expected):
+        """Field passes through to ActorPoolStrategy; None defers resolution."""
+        config = vLLMEngineProcessorConfig(
+            model_source="unsloth/Llama-3.2-1B-Instruct",
+            **kwargs,
+        )
+        assert config.max_tasks_in_flight_per_actor == expected
+        assert config.max_concurrent_batches == kwargs.get("max_concurrent_batches", 8)
+
+    def test_experimental_max_tasks_in_flight_per_actor_deprecated(self):
+        """Setting `experimental['max_tasks_in_flight_per_actor']` migrates to
+        the top-level field with a deprecation log; the explicit top-level
+        field overrides it but the warning still fires."""
+
+        def has_deprecation_log(warning_mock):
+            return any(
+                "max_tasks_in_flight_per_actor" in call.args[0]
+                and "deprecated" in call.args[0]
+                for call in warning_mock.call_args_list
+            )
+
+        # Migration: experimental → top-level field.
+        with patch.object(processor_base.logger, "warning") as warning_mock:
+            cfg = vLLMEngineProcessorConfig(
+                model_source="unsloth/Llama-3.2-1B-Instruct",
+                experimental={"max_tasks_in_flight_per_actor": 10},
+            )
+        assert cfg.max_tasks_in_flight_per_actor == 10
+        assert has_deprecation_log(warning_mock)
+
+        # Explicit top-level beats experimental, but warning still fires.
+        with patch.object(processor_base.logger, "warning") as warning_mock:
+            cfg = vLLMEngineProcessorConfig(
+                model_source="unsloth/Llama-3.2-1B-Instruct",
+                max_tasks_in_flight_per_actor=20,
+                experimental={"max_tasks_in_flight_per_actor": 10},
+            )
+        assert cfg.max_tasks_in_flight_per_actor == 20
+        assert has_deprecation_log(warning_mock)
+
+    def test_max_tasks_in_flight_under_max_concurrent_batches_warns(self):
+        with patch.object(processor_base.logger, "warning") as warning_mock:
+            cfg = vLLMEngineProcessorConfig(
+                model_source="unsloth/Llama-3.2-1B-Instruct",
+                max_tasks_in_flight_per_actor=1,
+                max_concurrent_batches=8,
+            )
+
+        assert cfg.max_tasks_in_flight_per_actor == 1
+        assert cfg.max_concurrent_batches == 8
+        warning_messages = [call.args[0] for call in warning_mock.call_args_list]
+        assert any(
+            "max_tasks_in_flight_per_actor" in message
+            and "max_concurrent_batches" in message
+            and "underutilize" in message
+            for message in warning_messages
+        )
+
+    @pytest.mark.parametrize(
+        "kwargs",
+        [
+            {},
+            {"max_tasks_in_flight_per_actor": 8, "max_concurrent_batches": 8},
+            {"max_tasks_in_flight_per_actor": 16, "max_concurrent_batches": 8},
+        ],
+    )
+    def test_max_tasks_in_flight_does_not_warn_when_not_underutilized(self, kwargs):
+        with patch.object(processor_base.logger, "warning") as warning_mock:
+            vLLMEngineProcessorConfig(
+                model_source="unsloth/Llama-3.2-1B-Instruct",
+                **kwargs,
+            )
+
+        warning_messages = [call.args[0] for call in warning_mock.call_args_list]
+        assert not any("underutilize" in message for message in warning_messages)
+
+
+class TestMapKwargs:
+    """Tests for preprocess_map_kwargs and postprocess_map_kwargs."""
+
+    def test_map_kwargs_stored_in_processor(self):
+        """Test that map kwargs are correctly stored in Processor."""
+        preprocess_kwargs = {"num_cpus": 0.5}
+        postprocess_kwargs = {"num_cpus": 0.25, "memory": 1024}
+
+        processor = Processor(
+            config=ProcessorConfig(batch_size=64),
+            stages=[],
+            preprocess=lambda row: {"val": row["id"]},
+            postprocess=lambda row: {"result": row["val"]},
+            preprocess_map_kwargs=preprocess_kwargs,
+            postprocess_map_kwargs=postprocess_kwargs,
+        )
+
+        assert processor.preprocess_map_kwargs == preprocess_kwargs
+        assert processor.postprocess_map_kwargs == postprocess_kwargs
+
+    def test_map_kwargs_defaults_to_empty_dict(self):
+        """Test that map kwargs default to empty dict when None."""
+        processor = Processor(
+            config=ProcessorConfig(batch_size=64),
+            stages=[],
+        )
+
+        assert processor.preprocess_map_kwargs == {}
+        assert processor.postprocess_map_kwargs == {}
+
+    def test_map_kwargs_passthrough_via_builder(self):
+        """Test that map kwargs are passed through ProcessorBuilder."""
+
+        def build_processor_simple(
+            config: ProcessorConfig,
+            preprocess=None,
+            postprocess=None,
+            preprocess_map_kwargs=None,
+            postprocess_map_kwargs=None,
+        ) -> Processor:
+            return Processor(
+                config,
+                [],
+                preprocess=preprocess,
+                postprocess=postprocess,
+                preprocess_map_kwargs=preprocess_map_kwargs,
+                postprocess_map_kwargs=postprocess_map_kwargs,
+            )
+
+        ProcessorBuilder.clear_registry()
+        ProcessorBuilder.register(DummyProcessorConfig, build_processor_simple)
+
+        config = DummyProcessorConfig(batch_size=64)
+        # Test through ProcessorBuilder which is called by build_processor
+        processor = ProcessorBuilder.build(
+            config,
+            preprocess=lambda row: {"val": row["id"]},
+            postprocess=lambda row: {"result": row["val"]},
+            preprocess_map_kwargs={"num_cpus": 0.5},
+            postprocess_map_kwargs={"num_cpus": 0.25},
+        )
+
+        assert processor.preprocess_map_kwargs == {"num_cpus": 0.5}
+        assert processor.postprocess_map_kwargs == {"num_cpus": 0.25}
+
+    def test_builder_kwargs_conflict_with_map_kwargs(self):
+        """Test that builder_kwargs validation rejects map kwargs."""
+        # Test the validation that build_processor calls
+        with pytest.raises(ValueError, match="builder_kwargs cannot contain"):
+            ProcessorBuilder.validate_builder_kwargs(
+                {"preprocess_map_kwargs": {"num_cpus": 0.5}}
+            )
+
+        with pytest.raises(ValueError, match="builder_kwargs cannot contain"):
+            ProcessorBuilder.validate_builder_kwargs(
+                {"postprocess_map_kwargs": {"num_cpus": 0.5}}
+            )
+
+    def test_end_to_end_with_map_kwargs(self):
+        """Test end-to-end execution with map kwargs."""
+        processor = Processor(
+            config=ProcessorConfig(batch_size=64),
+            stages=[],
+            preprocess=lambda row: {"val": row["id"] * 2},
+            postprocess=lambda row: {"result": row["val"] + 1, "id": row["id"]},
+            preprocess_map_kwargs={"num_cpus": 0.5},
+            postprocess_map_kwargs={"num_cpus": 0.25},
+        )
+
+        ds = ray.data.range(5)
+        result = processor(ds).take_all()
+
+        for row in result:
+            # Verify the computation: val = id * 2, result = val + 1
+            assert row["result"] == row["id"] * 2 + 1
+
+    def test_backward_compatibility_without_map_kwargs(self):
+        """Test that existing code without map kwargs still works."""
+        processor = Processor(
+            config=ProcessorConfig(batch_size=64),
+            stages=[],
+            preprocess=lambda row: {"val": row["id"]},
+            postprocess=lambda row: {"result": row["val"]},
+        )
+
+        ds = ray.data.range(5)
+        result = processor(ds).take_all()
+        # Sort results by result value since order is not guaranteed
+        result = sorted(result, key=lambda x: x["result"])
+
+        for i, row in enumerate(result):
+            assert row["result"] == i
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

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

@@ -0,0 +1,98 @@
+"""
+Test that Ray Data LLM does not override wait_for_min_actors_s.
+
+With default settings (wait_for_min_actors_s <= 0), processing starts
+as soon as any actor is ready, regardless of concurrency config.
+"""
+import sys
+
+import pytest
+
+from ray.data import DataContext
+from ray.llm._internal.batch.processor import ProcessorBuilder
+from ray.llm._internal.batch.processor.vllm_engine_proc import vLLMEngineProcessorConfig
+
+
+@pytest.fixture(autouse=True)
+def reset_data_context():
+    """Reset DataContext before and after each test."""
+    ctx = DataContext.get_current()
+    original_value = ctx.wait_for_min_actors_s
+    ctx.wait_for_min_actors_s = -1
+    yield
+    ctx.wait_for_min_actors_s = original_value
+
+
+class TestWaitForMinActorsNotOverridden:
+    """Test that Processor does not override wait_for_min_actors_s."""
+
+    def test_processor_does_not_override_default(self):
+        """Processor should not change wait_for_min_actors_s from default."""
+        ctx = DataContext.get_current()
+        ctx.wait_for_min_actors_s = -1
+
+        config = vLLMEngineProcessorConfig(
+            model_source="facebook/opt-125m",
+            concurrency=4,
+        )
+        ProcessorBuilder.build(config)
+
+        assert ctx.wait_for_min_actors_s == -1
+
+    @pytest.mark.parametrize("user_value", [60, 600, 1800])
+    def test_processor_preserves_user_setting(self, user_value):
+        """Processor should preserve user-set wait_for_min_actors_s."""
+        ctx = DataContext.get_current()
+        ctx.wait_for_min_actors_s = user_value
+
+        config = vLLMEngineProcessorConfig(
+            model_source="facebook/opt-125m",
+            concurrency=4,
+        )
+        ProcessorBuilder.build(config)
+
+        assert ctx.wait_for_min_actors_s == user_value
+
+
+class TestConcurrencyConfigPassthrough:
+    """
+    Test that concurrency config correctly sets ActorPoolStrategy.
+
+    This determines blocking behavior when wait_for_min_actors_s > 0:
+    - concurrency=N → min_size=1, max_size=N → blocks for 1 actor (autoscaling)
+    - concurrency=(m, N) → min_size=m, max_size=N → blocks for m actors
+    """
+
+    @pytest.mark.parametrize(
+        "concurrency,expected_min_size,expected_max_size",
+        [
+            (4, 1, 4),  # int: autoscaling pool with min_size=1
+            ((1, 4), 1, 4),  # tuple: autoscaling pool
+            ((2, 8), 2, 8),  # tuple: custom min
+        ],
+        ids=["int_concurrency", "tuple_1_to_n", "tuple_custom_min"],
+    )
+    def test_concurrency_to_actor_pool_strategy(
+        self, concurrency, expected_min_size, expected_max_size
+    ):
+        """Verify concurrency config maps to correct ActorPoolStrategy."""
+        config = vLLMEngineProcessorConfig(
+            model_source="facebook/opt-125m",
+            concurrency=concurrency,
+        )
+        processor = ProcessorBuilder.build(config)
+
+        # Get the vLLM stage and check its compute strategy
+        stage = processor.get_stage_by_name("vLLMEngineStage")
+        compute = stage.map_batches_kwargs.get("compute")
+
+        assert (
+            compute.min_size == expected_min_size
+        ), f"Expected min_size={expected_min_size}, got {compute.min_size}"
+        assert (
+            compute.max_size == expected_max_size
+        ), f"Expected max_size={expected_max_size}, got {compute.max_size}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/cpu/stages/test_chat_template_stage.py

@@ -0,0 +1,229 @@
+import shutil
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray.llm._internal.batch.stages.chat_template_stage import ChatTemplateUDF
+from ray.llm._internal.common.utils.download_utils import get_model_entrypoint
+
+
+@pytest.fixture
+def mock_tokenizer_setup():
+    # As of writing, the test environment does not have TYPE_CHECKING, which
+    # triggers lazy module import in transformers/__init__.py. This means the
+    # mocking may not work without explicitly importing AutoProcessor, hence
+    # the following import.
+    from transformers import AutoProcessor  # noqa: F401
+
+    with patch("transformers.AutoProcessor") as mock_auto_processor:
+        mock_processor = MagicMock()
+        mock_auto_processor.from_pretrained.return_value = mock_processor
+        yield mock_processor
+
+
+@pytest.mark.asyncio
+async def test_chat_template_udf_basic(mock_tokenizer_setup):
+    mock_tokenizer = mock_tokenizer_setup
+    mock_tokenizer.apply_chat_template.return_value = "<chat>Hello AI</chat>"
+
+    udf = ChatTemplateUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model="test-model",
+    )
+
+    batch = {
+        "__data": [
+            {
+                "messages": MagicMock(
+                    tolist=lambda: [{"role": "user", "content": "Hello AI"}]
+                )
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 1
+    assert results[0]["prompt"] == "<chat>Hello AI</chat>"
+    mock_tokenizer.apply_chat_template.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_chat_template_udf_multiple_messages(mock_tokenizer_setup):
+    mock_tokenizer = mock_tokenizer_setup
+    mock_tokenizer.apply_chat_template.side_effect = [
+        "<chat>Hello AI</chat>",
+        "<chat>How are you?</chat>",
+    ]
+
+    udf = ChatTemplateUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model="test-model",
+    )
+
+    batch = {
+        "__data": [
+            {
+                "messages": MagicMock(
+                    tolist=lambda: [{"role": "user", "content": "Hello AI"}]
+                )
+            },
+            {
+                "messages": MagicMock(
+                    tolist=lambda: [{"role": "user", "content": "How are you?"}],
+                )
+            },
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.append(result)
+
+    assert len(results) == 1
+    assert results[0]["__data"][0]["prompt"] == "<chat>Hello AI</chat>"
+    assert results[0]["__data"][1]["prompt"] == "<chat>How are you?</chat>"
+    assert mock_tokenizer.apply_chat_template.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "chat_template_kwargs, expected_prompt",
+    [
+        ({"enable_thinking": False}, "Answer without thinking"),
+        ({"enable_thinking": True}, "<think>thinking</think>"),
+        ({}, "<think>thinking</think>"),
+        (
+            {"enable_thinking": True, "custom_param": "test_value", "temperature": 0.7},
+            "<think>thinking</think>",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_chat_template_udf_chat_template_kwargs(
+    mock_tokenizer_setup, chat_template_kwargs, expected_prompt
+):
+    mock_tokenizer = mock_tokenizer_setup
+
+    # Store captured kwargs for verification
+    captured_kwargs = {}
+
+    def side_effect_func(conversation, **kwargs):
+        # Capture all kwargs for later verification
+        captured_kwargs.update(kwargs)
+
+        enable_thinking = kwargs.get("enable_thinking", True)
+        if enable_thinking is False:
+            return "Answer without thinking"
+        else:
+            return "<think>thinking</think>"
+
+    mock_tokenizer.apply_chat_template.side_effect = side_effect_func
+
+    udf = ChatTemplateUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model="test-model",
+        chat_template_kwargs=chat_template_kwargs,
+    )
+
+    # Assert that the chat_template_kwargs were properly stored
+    assert udf.chat_template_kwargs == chat_template_kwargs
+
+    batch = {
+        "__data": [
+            {
+                "messages": MagicMock(
+                    tolist=lambda: [{"role": "user", "content": "Hello AI"}]
+                )
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 1
+    assert results[0]["prompt"] == expected_prompt
+
+    # Verify that all chat_template_kwargs were passed through to apply_chat_template
+    for key, value in chat_template_kwargs.items():
+        assert (
+            key in captured_kwargs
+        ), f"Expected kwargs key '{key}' not found in captured kwargs"
+        assert (
+            captured_kwargs[key] == value
+        ), f"Expected '{key}': {value}, but got '{key}': {captured_kwargs[key]}"
+
+
+@pytest.mark.asyncio
+async def test_chat_template_udf_assistant_prefill(mock_tokenizer_setup):
+    mock_tokenizer = mock_tokenizer_setup
+    mock_tokenizer.apply_chat_template.side_effect = [
+        "<chat>Hello AI<assistant><think>\n</chat>",
+        "<chat>Hello AI</chat>",
+    ]
+
+    udf = ChatTemplateUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model="test-model",
+    )
+
+    batch = {
+        "__data": [
+            {
+                "messages": MagicMock(
+                    tolist=lambda: [
+                        {"role": "user", "content": "Hello AI"},
+                        {"role": "assistant", "content": "<think>\n"},
+                    ]
+                ),
+            },
+            {
+                "messages": MagicMock(
+                    tolist=lambda: [{"role": "user", "content": "How are you?"}],
+                )
+            },
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 2
+    assert mock_tokenizer.apply_chat_template.call_count == 2
+    assert results[0]["prompt"] == "<chat>Hello AI<assistant><think>\n</chat>"
+    assert results[1]["prompt"] == "<chat>Hello AI</chat>"
+    # check if kwargs were set properly
+    call_args_list = mock_tokenizer.apply_chat_template.call_args_list
+    args1, kwargs1 = call_args_list[0]
+    assert not kwargs1.get("add_generation_prompt")
+    assert kwargs1.get("continue_final_message")
+    _, kwargs2 = call_args_list[1]
+    assert kwargs2.get("add_generation_prompt")
+    assert not kwargs2.get("continue_final_message")
+
+
+def test_trust_remote_code(model_internlm2_1_8b):
+    model_entry = get_model_entrypoint(model_internlm2_1_8b)
+    if model_entry != model_internlm2_1_8b:
+        shutil.rmtree(model_entry, ignore_errors=True)
+
+    udf = ChatTemplateUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model=model_internlm2_1_8b,
+        trust_remote_code=True,
+    )
+    assert udf.processor is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/cpu/stages/test_configs.py

@@ -0,0 +1,159 @@
+import sys
+
+import pytest
+
+from ray.llm._internal.batch.stages.configs import (
+    ChatTemplateStageConfig,
+    PrepareMultimodalStageConfig,
+    TokenizerStageConfig,
+    resolve_stage_config,
+)
+
+
+def test_resolve_dict_to_config():
+    """Dict → parsed StageConfig with all fields."""
+    stage_cfg = resolve_stage_config(
+        {"enabled": True, "batch_size": 128, "concurrency": 4},
+        PrepareMultimodalStageConfig,
+    )
+    assert isinstance(stage_cfg, PrepareMultimodalStageConfig)
+    assert stage_cfg.enabled is True
+    assert stage_cfg.batch_size == 128
+    assert stage_cfg.concurrency == 4
+
+
+def test_resolve_typed_config_copied():
+    """Typed config creates copy (doesn't mutate input)."""
+    original = ChatTemplateStageConfig(
+        enabled=True, batch_size=64, model_source="model1"
+    )
+    processor_defaults = {"batch_size": 128, "model_source": "model2"}
+
+    # Resolve with first processor defaults
+    resolved1 = resolve_stage_config(
+        original, ChatTemplateStageConfig, processor_defaults
+    )
+    assert resolved1.batch_size == 64  # Explicit value preserved
+    assert resolved1.model_source == "model1"  # Explicit value preserved
+
+    # Resolve same original with different processor defaults
+    processor_defaults2 = {"batch_size": 256, "model_source": "model3"}
+    resolved2 = resolve_stage_config(
+        original, ChatTemplateStageConfig, processor_defaults2
+    )
+    assert resolved2.batch_size == 64  # Still preserved
+    assert resolved2.model_source == "model1"  # Still preserved
+
+    # Original unchanged
+    assert original.batch_size == 64
+    assert original.model_source == "model1"
+
+
+def test_resolve_merges_processor_defaults():
+    """Defaults merged when None."""
+    stage_cfg = resolve_stage_config(
+        {"enabled": True},
+        TokenizerStageConfig,
+        processor_defaults={
+            "batch_size": 128,
+            "concurrency": 4,
+            "model_source": "test-model",
+        },
+    )
+    assert stage_cfg.batch_size == 128
+    assert stage_cfg.concurrency == 4
+    assert stage_cfg.model_source == "test-model"
+
+
+def test_resolve_preserves_explicit_overrides():
+    """Explicit values not overridden by defaults."""
+    stage_cfg = resolve_stage_config(
+        {"enabled": True, "batch_size": 64, "concurrency": 2},
+        PrepareMultimodalStageConfig,
+        processor_defaults={"batch_size": 128, "concurrency": 4},
+    )
+    assert stage_cfg.batch_size == 64  # Explicit override preserved
+    assert stage_cfg.concurrency == 2  # Explicit override preserved
+
+
+def test_resolve_model_source_fallback():
+    """model_source field uses processor default when None."""
+    stage_cfg = resolve_stage_config(
+        {"enabled": True},
+        TokenizerStageConfig,
+        processor_defaults={"model_source": "default-model"},
+    )
+    assert stage_cfg.model_source == "default-model"
+
+
+def test_resolve_bool_true():
+    """Bool True → enabled StageConfig."""
+    stage_cfg = resolve_stage_config(True, PrepareMultimodalStageConfig)
+    assert isinstance(stage_cfg, PrepareMultimodalStageConfig)
+    assert stage_cfg.enabled is True
+
+
+def test_resolve_bool_false():
+    """Bool False → disabled StageConfig."""
+    stage_cfg = resolve_stage_config(False, PrepareMultimodalStageConfig)
+    assert isinstance(stage_cfg, PrepareMultimodalStageConfig)
+    assert stage_cfg.enabled is False
+
+
+def test_resolve_runtime_env_replacement():
+    """Stage runtime_env replaces processor (not merged)."""
+    stage_cfg = resolve_stage_config(
+        {"enabled": True, "runtime_env": {"env_vars": {"STAGE_VAR": "stage_value"}}},
+        PrepareMultimodalStageConfig,
+        processor_defaults={"runtime_env": {"env_vars": {"PROC_VAR": "proc_value"}}},
+    )
+    # Stage runtime_env completely replaces processor runtime_env
+    assert stage_cfg.runtime_env == {"env_vars": {"STAGE_VAR": "stage_value"}}
+
+
+def test_resolve_same_config_reusable():
+    """Same StageConfig instance can be resolved with different processor defaults without mutation."""
+    original = ChatTemplateStageConfig(enabled=True, batch_size=None, model_source=None)
+    processor_defaults1 = {"batch_size": 128, "model_source": "model1"}
+    processor_defaults2 = {"batch_size": 256, "model_source": "model2"}
+
+    resolved1 = resolve_stage_config(
+        original, ChatTemplateStageConfig, processor_defaults1
+    )
+    resolved2 = resolve_stage_config(
+        original, ChatTemplateStageConfig, processor_defaults2
+    )
+
+    # Each resolution gets its own defaults
+    assert resolved1.batch_size == 128
+    assert resolved1.model_source == "model1"
+    assert resolved2.batch_size == 256
+    assert resolved2.model_source == "model2"
+
+    # Original unchanged (still None)
+    assert original.batch_size is None
+    assert original.model_source is None
+
+
+def test_resolve_unsupported_type():
+    """Unsupported types raise TypeError."""
+    with pytest.raises(TypeError, match="Unsupported type for stage config"):
+        resolve_stage_config(None, PrepareMultimodalStageConfig)
+
+    with pytest.raises(TypeError, match="Unsupported type for stage config"):
+        resolve_stage_config(123, PrepareMultimodalStageConfig)
+
+
+def test_resolve_stage_without_model_source():
+    """Stages without model_source field don't get it from defaults."""
+    stage_cfg = resolve_stage_config(
+        {"enabled": True},
+        PrepareMultimodalStageConfig,
+        processor_defaults={"model_source": "test-model"},
+    )
+    # PrepareMultimodalStageConfig doesn't have a model_source field.
+    assert not hasattr(stage_cfg, "model_source")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/cpu/stages/test_http_request_stage.py

@@ -0,0 +1,226 @@
+import asyncio
+import json
+import sys
+from unittest.mock import AsyncMock, call, patch
+
+import aiohttp
+import aiohttp.web
+import aiohttp.web_exceptions
+import numpy as np
+import pytest
+from aiohttp.test_utils import TestServer
+
+import ray.data
+from ray.llm._internal.batch.processor import ProcessorBuilder
+from ray.llm._internal.batch.processor.http_request_proc import (
+    HttpRequestProcessorConfig,
+)
+from ray.llm._internal.batch.stages.http_request_stage import (
+    HttpRequestUDF,
+    NumpyEncoder,
+)
+
+
+@pytest.fixture
+def mock_response():
+    mock = AsyncMock()
+    mock.json = AsyncMock(return_value={"response": "test"})
+    return mock
+
+
+@pytest.fixture
+def mock_session(mock_response):
+    session = AsyncMock()
+    session.post.return_value.__aenter__.return_value = mock_response
+    session_cm = AsyncMock()
+    session_cm.__aenter__.return_value = session
+    return session_cm
+
+
+@pytest.mark.asyncio
+async def test_http_request_udf_basic(mock_session):
+    udf = HttpRequestUDF(
+        data_column="__data",
+        expected_input_keys=["payload"],
+        url="http://test.com/api",
+        additional_header={"Authorization": "Bearer 1234567890"},
+        qps=None,
+        session_factory=lambda: mock_session,  # noqa: E731
+    )
+
+    batch = {"__data": [{"payload": {"text": "hello", "metadata": "test"}}]}
+
+    async for result in udf(batch):
+        assert result["__data"][0]["http_response"]["response"] == "test"
+
+    mock_session.__aenter__.return_value.post.assert_called_once_with(
+        "http://test.com/api",
+        headers={
+            "Content-Type": "application/json",
+            "Authorization": "Bearer 1234567890",
+        },
+        data=json.dumps({"text": "hello", "metadata": "test"}),
+    )
+
+
+@pytest.mark.asyncio
+async def test_http_request_udf_with_qps(mock_session):
+    udf = HttpRequestUDF(
+        data_column="__data",
+        expected_input_keys=["payload"],
+        url="http://test.com/api",
+        qps=2,
+        session_factory=lambda: mock_session,  # noqa: E731
+    )
+
+    batch = {
+        "__data": [{"payload": {"text": "hello1"}}, {"payload": {"text": "hello2"}}]
+    }
+
+    with patch("time.time") as mock_time, patch("asyncio.sleep") as mock_sleep:
+        # Mock time to test QPS limiting. Req2 cannot be sent until 0.5s,
+        # so the asyncio.sleep should be called once.
+        # [start_time, req1_time, req2_time]
+        mock_time.side_effect = [0, 0.1, 0.2]
+
+        results = []
+        async for result in udf(batch):
+            results.extend(result["__data"])
+
+        assert len(results) == 2
+        assert mock_sleep.called  # Should have called sleep for QPS limiting
+
+
+@pytest.mark.asyncio
+async def test_http_request_udf_with_retry(mock_response):
+    batch = {
+        "__data": [{"payload": {"text": "hello1"}}, {"payload": {"text": "hello2"}}]
+    }
+    # Create a fake session
+    # create another response
+    retry_resp = AsyncMock(status=429)
+    retry_resp.json = AsyncMock(return_value={"detail": "Too Many Requests"})
+
+    session = AsyncMock()
+    session.post.return_value.__aenter__.side_effect = [
+        mock_response,  # First request: success
+        asyncio.TimeoutError(),  # Second request, initial attempt: timeout
+        aiohttp.ClientConnectionError(),  # Second request, first retry: connection error
+        retry_resp,  # Second request, second retry: HTTP 429 error
+        mock_response,  # Final retry: success
+    ]
+    session_cm = AsyncMock()
+    session_cm.__aenter__.return_value = session
+    fake_session_factory = lambda: session_cm  # noqa: E731
+
+    udf = HttpRequestUDF(
+        data_column="__data",
+        expected_input_keys=["payload"],
+        url="http://test.com/api",
+        max_retries=3,
+        base_retry_wait_time_in_s=1,
+        session_factory=fake_session_factory,
+    )
+
+    with patch("asyncio.sleep") as mock_sleep:
+        results = []
+        async for result in udf(batch):
+            results.extend(result["__data"])
+
+        assert len(results) == 2
+        mock_sleep.assert_called()
+        mock_sleep.assert_has_calls(
+            [
+                call(udf.base_retry_wait_time_in_s),
+                call(udf.base_retry_wait_time_in_s * 2),
+                call(udf.base_retry_wait_time_in_s * 2 * 2),
+            ]
+        )
+
+
+def test_numpy_encoder():
+    """Test NumpyEncoder correctly serializes numpy data types."""
+    data = {
+        "ndarray": np.array([1, 2, 3]),
+        "integer": np.int64(10),
+        "float": np.float64(3.14),
+        "bool": np.bool_(True),
+        "list": [np.int32(1), np.float32(2.0)],
+    }
+
+    json_str = json.dumps(data, cls=NumpyEncoder)
+    decoded = json.loads(json_str)
+
+    assert decoded["ndarray"] == [1, 2, 3]
+    assert decoded["integer"] == 10
+    assert decoded["float"] == 3.14
+    assert decoded["bool"] is True
+    assert decoded["list"] == [1, 2.0]
+
+
+@pytest.fixture
+async def numpy_payload_server():
+    # Handler that verifies numpy array was correctly serialized
+    async def handler(request):
+        data = await request.json()
+        assert data["model"] == "test-model"
+        assert data["embedding"] == [1.0, 2.0, 3.0]
+        assert data["flags"] == [True, False]
+        assert len(data["messages"]) == 1
+        assert data["messages"][0]["role"] == "user"
+        assert len(data["messages"][0]["content"]) == 2
+        return aiohttp.web.json_response({"response": "success"})
+
+    # Create test app and server
+    app = aiohttp.web.Application()
+    app.router.add_post("/", handler)
+    server = TestServer(app)
+    await server.start_server()
+    yield server
+    await server.close()
+
+
+@pytest.mark.asyncio
+async def test_http_request_udf_with_numpy_payload_server(numpy_payload_server):
+    """Test HttpRequestUDF with numpy arrays using a real aiohttp server."""
+    data = [
+        {
+            "payload": {
+                "model": "test-model",
+                "embedding": np.array([1.0, 2.0, 3.0]),
+                "flags": np.array([True, False]),
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": "http://example.com/image.png"},
+                            },
+                            {
+                                "type": "text",
+                                "text": "hello",
+                            },
+                        ],
+                    }
+                ],
+            }
+        }
+    ]
+    config = HttpRequestProcessorConfig(
+        url=str(numpy_payload_server.make_url("/")),
+        headers={"Content-Type": "application/json"},
+        qps=None,
+    )
+
+    processor = ProcessorBuilder.build(config)
+    ds = processor(ray.data.from_items(data * 10))
+    results = await asyncio.to_thread(ds.take_all)
+
+    assert len(results) == 10
+    for result in results:
+        assert result["http_response"]["response"] == "success"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/cpu/stages/test_prepare_multimodal_stage.py

@@ -0,0 +1,148 @@
+import sys
+
+import pytest
+
+from ray.llm._internal.batch.stages.prepare_multimodal_stage import (
+    PrepareMultimodalUDF,
+)
+
+
+@pytest.mark.asyncio
+async def test_prepare_multimodal_udf_image_url(image_asset):
+    image_url, _ = image_asset
+    udf = PrepareMultimodalUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model_config_kwargs={"model": "Qwen/Qwen2.5-VL-3B-Instruct"},
+        chat_template_content_format="string",
+    )
+
+    batch = {
+        "__data": [
+            {
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "Describe this image in 10 words.",
+                            },
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": image_url},
+                                "uuid": "image-1-id",
+                            },
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": image_url},
+                                "uuid": "image-2-id",
+                            },
+                        ],
+                    },
+                ]
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.append(result["__data"][0])
+
+    assert len(results) == 1
+    assert "multimodal_data" in results[0]
+    assert len(results[0]["multimodal_data"]["image"]) == 2
+    assert "multimodal_uuids" in results[0]
+    assert results[0]["multimodal_uuids"] == {"image": ["image-1-id", "image-2-id"]}
+    assert "messages" in results[0]
+
+
+@pytest.mark.asyncio
+async def test_prepare_multimodal_udf_pil_image(image_asset):
+    _, image_pil = image_asset
+
+    udf = PrepareMultimodalUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model_config_kwargs={"model": "Qwen/Qwen2.5-VL-3B-Instruct"},
+        chat_template_content_format="string",
+    )
+
+    batch = {
+        "__data": [
+            {
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": "Describe this image in 10 words.",
+                            },
+                            {
+                                "type": "image_pil",
+                                "image_pil": image_pil,
+                            },
+                        ],
+                    },
+                ]
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.append(result["__data"][0])
+
+    assert len(results) == 1
+    assert "multimodal_data" in results[0]
+    assert "messages" in results[0]
+
+
+@pytest.mark.asyncio
+async def test_prepare_multimodal_udf_no_multimodal_content():
+    """
+    Multimodal stage should proceed as normal if there is no multimodal content provided in messages.
+    """
+    udf = PrepareMultimodalUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model_config_kwargs={"model": "Qwen/Qwen2.5-VL-3B-Instruct"},
+        chat_template_content_format="string",
+    )
+
+    batch = {
+        "__data": [
+            {
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "Hello, how are you?"},
+                ]
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.append(result["__data"][0])
+
+    assert len(results) == 1
+    assert "multimodal_data" in results[0]
+    assert results[0]["multimodal_data"] is None
+    assert "messages" in results[0]
+
+
+def test_prepare_multimodal_udf_expected_keys():
+    udf = PrepareMultimodalUDF(
+        data_column="__data",
+        expected_input_keys=["messages"],
+        model_config_kwargs={"model": "Qwen/Qwen2.5-VL-3B-Instruct"},
+        chat_template_content_format="string",
+    )
+    assert udf.expected_input_keys == {"messages"}
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/cpu/stages/test_stage_base.py

@@ -0,0 +1,263 @@
+import sys
+from typing import Any, AsyncIterator, Dict, List, Optional
+
+import pytest
+
+from ray.llm._internal.batch.stages.base import (
+    StatefulStageUDF,
+    wrap_postprocess,
+    wrap_preprocess,
+)
+
+
+def test_wrap_preprocess():
+    # Test function that doubles a number
+    def double(x: dict) -> dict:
+        return {"value": x["id"] * 2}
+
+    # Test with carry_over=True
+    wrapped = wrap_preprocess(double, "__data")
+    result = wrapped({"id": 5, "extra": "memo"})
+    assert result == {"__data": {"id": 5, "extra": "memo", "value": 10}}
+
+
+def test_wrap_postprocess():
+    # Test function that converts number to string
+    def to_string(x: dict) -> dict:
+        return {
+            "result": str(x["value"]),
+            "extra": x["extra"],
+        }
+
+    # Test with carry_over=True
+    wrapped = wrap_postprocess(to_string, "__data")
+    result = wrapped({"__data": {"id": 5, "extra": "memo", "value": 10}})
+    assert result == {"extra": "memo", "result": "10"}
+
+    # Test missing input column
+    with pytest.raises(ValueError):
+        wrapped({"wrong_key": 42})
+
+
+def test_wrap_postprocess_bypasses_error_rows():
+    """Error rows with __inference_error__ set bypass user postprocess."""
+
+    def user_fn(data: dict) -> dict:
+        # Would crash if called with error row (missing generated_text)
+        return {"response": data["generated_text"].upper()}
+
+    wrapped = wrap_postprocess(user_fn, "__data")
+
+    error_row = {
+        "__data": {
+            "__inference_error__": "ValueError: prompt too long",
+            "prompt": "This is a long prompt",
+        }
+    }
+    result = wrapped(error_row)
+    # Error rows return entire data dict to preserve debugging info
+    assert result == {
+        "__inference_error__": "ValueError: prompt too long",
+        "prompt": "This is a long prompt",
+    }
+
+
+def test_wrap_postprocess_success_rows_run_postprocess():
+    """Success rows (__inference_error__ is None) run user postprocess."""
+
+    def user_fn(data: dict) -> dict:
+        return {"response": data["generated_text"], "tokens": data["num_tokens"]}
+
+    wrapped = wrap_postprocess(user_fn, "__data")
+
+    success_row = {
+        "__data": {
+            "generated_text": "Hello world",
+            "num_tokens": 10,
+            "__inference_error__": "",
+        }
+    }
+    result = wrapped(success_row)
+    assert result == {"response": "Hello world", "tokens": 10}
+
+
+def test_wrap_postprocess_include_error_column():
+    """With include_error_column=True, success rows include __inference_error__: None."""
+
+    def user_fn(data: dict) -> dict:
+        return {"response": data["generated_text"]}
+
+    wrapped = wrap_postprocess(user_fn, "__data", include_error_column=True)
+
+    success_row = {
+        "__data": {
+            "generated_text": "Hello world",
+            "__inference_error__": "",
+        }
+    }
+    result = wrapped(success_row)
+    assert result == {"response": "Hello world", "__inference_error__": ""}
+
+    # Error rows return entire data dict to preserve debugging info
+    error_row = {
+        "__data": {
+            "__inference_error__": "ValueError: prompt too long",
+            "prompt": "a long prompt",
+        }
+    }
+    result = wrapped(error_row)
+    assert result == {
+        "__inference_error__": "ValueError: prompt too long",
+        "prompt": "a long prompt",
+    }
+
+
+class TestStatefulStageUDF:
+    class SimpleUDF(StatefulStageUDF):
+        def __init__(
+            self,
+            data_column: str,
+            expected_input_keys: Optional[List[str]] = None,
+            udf_output_missing_idx_in_batch_column: bool = False,
+        ):
+            super().__init__(data_column, expected_input_keys)
+            self.udf_output_missing_idx_in_batch_column = (
+                udf_output_missing_idx_in_batch_column
+            )
+
+        async def udf(
+            self, rows: list[Dict[str, Any]]
+        ) -> AsyncIterator[Dict[str, Any]]:
+
+            # Intentionally output in a reversed order to test OOO.
+            for row in rows[::-1]:
+                ret = {"processed": row["value"] * 2}
+                if not self.udf_output_missing_idx_in_batch_column:
+                    ret[self.IDX_IN_BATCH_COLUMN] = row[self.IDX_IN_BATCH_COLUMN]
+                yield ret
+
+    @pytest.mark.asyncio
+    async def test_basic_processing(self):
+        udf = self.SimpleUDF(data_column="__data", expected_input_keys=["value"])
+
+        batch = {
+            "__data": [{"value": 1, "extra": 10}, {"value": 2, "extra": 20}],
+        }
+
+        results = []
+        async for result in udf(batch):
+            results.extend(result["__data"])
+
+        assert len(results) == 2
+        for data in results:
+            val = data["value"]
+            assert data["processed"] == val * 2
+            assert data["extra"] == 10 * val
+            assert data["value"] == val
+
+    @pytest.mark.asyncio
+    async def test_missing_data_column(self):
+        udf = self.SimpleUDF(data_column="__data", expected_input_keys=["value"])
+
+        batch = {"extra": ["a"]}
+
+        with pytest.raises(ValueError):
+            async for _ in udf(batch):
+                pass
+
+    @pytest.mark.asyncio
+    async def test_missing_required_key(self):
+        udf = self.SimpleUDF(data_column="__data", expected_input_keys=["value"])
+
+        batch = {"__data": [{"wrong_key": 1}]}
+
+        with pytest.raises(ValueError):
+            async for _ in udf(batch):
+                pass
+
+    @pytest.mark.asyncio
+    async def test_missing_idx_in_batch_column(self):
+        udf = self.SimpleUDF(
+            data_column="__data",
+            expected_input_keys=["value"],
+            udf_output_missing_idx_in_batch_column=True,
+        )
+
+        batch = {"__data": [{"value": 1}]}
+
+        with pytest.raises(ValueError):
+            async for _ in udf(batch):
+                pass
+
+    @pytest.mark.asyncio
+    async def test_error_rows_bypass_udf(self):
+        """Error rows with __inference_error__ bypass the UDF entirely."""
+
+        class FailOnMissingValueUDF(StatefulStageUDF):
+            async def udf(
+                self, rows: list[Dict[str, Any]]
+            ) -> AsyncIterator[Dict[str, Any]]:
+                for row in rows:
+                    # Would crash on error rows missing 'value' field
+                    yield {
+                        self.IDX_IN_BATCH_COLUMN: row[self.IDX_IN_BATCH_COLUMN],
+                        "processed": row["value"] * 2,
+                    }
+
+        udf = FailOnMissingValueUDF(
+            data_column="__data", expected_input_keys=None  # Skip validation
+        )
+
+        batch = {
+            "__data": [
+                {"value": 1},  # Normal row
+                {"__inference_error__": "ValueError: prompt too long"},  # Error row
+                {"value": 3},  # Normal row
+            ]
+        }
+
+        results = []
+        async for result in udf(batch):
+            results.extend(result["__data"])
+
+        assert len(results) == 3
+
+        # Normal rows are processed
+        assert results[0]["processed"] == 2
+        assert results[2]["processed"] == 6
+
+        # Error row passes through unchanged
+        assert results[1]["__inference_error__"] == "ValueError: prompt too long"
+        assert "processed" not in results[1]
+
+    @pytest.mark.asyncio
+    async def test_all_error_rows_in_batch(self):
+        """Batch with all error rows should pass through without calling UDF."""
+
+        class FailIfCalledUDF(StatefulStageUDF):
+            async def udf(
+                self, rows: list[Dict[str, Any]]
+            ) -> AsyncIterator[Dict[str, Any]]:
+                raise AssertionError("UDF should not be called for all-error batch")
+                yield  # Make this a generator
+
+        udf = FailIfCalledUDF(data_column="__data", expected_input_keys=None)
+
+        batch = {
+            "__data": [
+                {"__inference_error__": "error 1"},
+                {"__inference_error__": "error 2"},
+            ]
+        }
+
+        results = []
+        async for result in udf(batch):
+            results.extend(result["__data"])
+
+        assert len(results) == 2
+        assert results[0]["__inference_error__"] == "error 1"
+        assert results[1]["__inference_error__"] == "error 2"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/cpu/stages/test_tokenize_stage.py

@@ -0,0 +1,101 @@
+import shutil
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray.llm._internal.batch.stages.tokenize_stage import DetokenizeUDF, TokenizeUDF
+from ray.llm._internal.common.utils.download_utils import get_model_entrypoint
+
+
+@pytest.fixture
+def mock_tokenizer_setup():
+    with patch(
+        "ray.llm._internal.batch.stages.tokenize_stage.get_cached_tokenizer"
+    ) as mock_get_tokenizer, patch("transformers.AutoTokenizer") as mock_auto_tokenizer:
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.side_effect = lambda texts: {
+            "input_ids": [[1, 2, 3] for _ in texts]
+        }
+        mock_get_tokenizer.return_value = mock_tokenizer
+        mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
+        yield mock_tokenizer
+
+
+@pytest.mark.asyncio
+async def test_tokenize_udf_basic(mock_tokenizer_setup):
+    mock_tokenizer = mock_tokenizer_setup
+    mock_tokenizer.return_value = [
+        {"input_ids": [1, 2, 3]},
+        {"input_ids": [4, 5, 6]},
+    ]
+
+    udf = TokenizeUDF(
+        data_column="__data", model="test-model", expected_input_keys=["prompt"]
+    )
+    batch = {"__data": [{"prompt": "Hello"}, {"prompt": "World"}]}
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 2
+    assert all(result["tokenized_prompt"] == [1, 2, 3] for result in results)
+    assert all(
+        original["prompt"] == result["prompt"]
+        for original, result in zip(batch, results)
+    )
+
+
+@pytest.mark.asyncio
+async def test_detokenize_udf_basic(mock_tokenizer_setup):
+    mock_tokenizer = mock_tokenizer_setup
+    mock_tokenizer.batch_decode.return_value = ["Hello", "World"]
+
+    udf = DetokenizeUDF(
+        data_column="__data",
+        model="test-model",
+        expected_input_keys=["generated_tokens"],
+    )
+    batch = {
+        "__data": [
+            {"generated_tokens": [1, 2, 3]},
+            {"generated_tokens": [4, 5, 6]},
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 2
+    assert results[0]["generated_text"] == "Hello"
+    assert results[1]["generated_text"] == "World"
+    mock_tokenizer.batch_decode.assert_called_once_with(
+        [[1, 2, 3], [4, 5, 6]], skip_special_tokens=True
+    )
+
+
+@pytest.mark.parametrize(
+    "udf_cls, expected_input_keys",
+    [
+        (TokenizeUDF, ["prompt"]),
+        (DetokenizeUDF, ["generated_tokens"]),
+    ],
+)
+def test_trust_remote_code(model_internlm2_1_8b, udf_cls, expected_input_keys):
+    model_entry = get_model_entrypoint(model_internlm2_1_8b)
+    if model_entry != model_internlm2_1_8b:
+        shutil.rmtree(model_entry, ignore_errors=True)
+
+    udf = udf_cls(
+        data_column="__data",
+        model=model_internlm2_1_8b,
+        expected_input_keys=expected_input_keys,
+        trust_remote_code=True,
+    )
+    assert udf.tokenizer is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/gpu/processor/test_serve_deployment_proc.py

@@ -0,0 +1,300 @@
+import sys
+from typing import Any, Dict
+
+import pytest
+
+import ray
+from ray import serve
+from ray.data import ActorPoolStrategy
+from ray.data.llm import ServeDeploymentProcessorConfig, build_processor
+from ray.llm._internal.batch.processor import ProcessorBuilder
+from ray.serve.llm.openai_api_models import ChatCompletionRequest, CompletionRequest
+
+
+@pytest.mark.parametrize(
+    "dtype_mapping", [None, {"CompletionRequest": CompletionRequest}]
+)
+def test_serve_deployment_processor(dtype_mapping):
+    app_name = "test_serve_deployment_processor_app"
+    deployment_name = "test_serve_deployment_name"
+
+    config_kwargs = dict(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        batch_size=16,
+        concurrency=1,
+    )
+    if dtype_mapping is not None:
+        config_kwargs["dtype_mapping"] = dtype_mapping
+    config = ServeDeploymentProcessorConfig(**config_kwargs)
+
+    processor = ProcessorBuilder.build(config)
+    assert processor.list_stage_names() == [
+        "ServeDeploymentStage",
+    ]
+
+    stage = processor.get_stage_by_name("ServeDeploymentStage")
+    assert stage.fn_constructor_kwargs == {
+        "deployment_name": deployment_name,
+        "app_name": app_name,
+        "dtype_mapping": dtype_mapping,
+        "should_continue_on_error": False,
+        "request_timeout_s": None,
+    }
+
+    assert "compute" in stage.map_batches_kwargs
+    assert isinstance(stage.map_batches_kwargs["compute"], ActorPoolStrategy)
+    assert stage.map_batches_kwargs["compute"].min_size == 1
+    assert stage.map_batches_kwargs["compute"].max_size == 1
+
+
+def test_simple_serve_deployment(serve_cleanup):
+    @serve.deployment
+    class SimpleServeDeployment:
+        # ServeDeploymentStageUDF expects an async generator.
+        async def add(self, request: Dict[str, Any]):
+            yield {"result": request["x"] + 1}
+
+    app_name = "simple_serve_deployment_app"
+    deployment_name = "SimpleServeDeployment"
+
+    serve.run(SimpleServeDeployment.bind(), name=app_name)
+
+    config = ServeDeploymentProcessorConfig(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        batch_size=16,
+        concurrency=1,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            method="add",
+            dtype=None,  # Empty dtype since output is already dict format
+            request_kwargs=dict(x=row["id"]),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["result"],
+            id=row["id"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"]})
+    ds = processor(ds)
+
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+    assert all(out["resp"] == out["id"] + 1 for out in outs)
+
+
+def test_serve_deployment_continue_on_error(serve_cleanup):
+    @serve.deployment
+    class FailingServeDeployment:
+        async def process(self, request: Dict[str, Any]):
+            x = request["x"]
+            if x % 10 == 0:  # Fail every 10th row
+                raise ValueError(f"Intentional failure for x={x}")
+            yield {"result": x * 2}
+
+    app_name = "failing_serve_deployment_app"
+    deployment_name = "FailingServeDeployment"
+
+    serve.run(FailingServeDeployment.bind(), name=app_name)
+
+    config = ServeDeploymentProcessorConfig(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        batch_size=16,
+        concurrency=1,
+        should_continue_on_error=True,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            method="process",
+            dtype=None,
+            request_kwargs=dict(x=row["id"]),
+        ),
+        # Error rows will bypass this postprocess and return raw data with
+        # __inference_error__ set. Only success rows get resp/id keys.
+        postprocess=lambda row: dict(
+            resp=row.get("result"),
+            id=row.get("id"),
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"]})
+    ds = processor(ds)
+
+    outs = ds.take_all()
+    assert len(outs) == 60
+
+    # Check __inference_error__ directly
+    errors = [o for o in outs if o.get("__inference_error__", "")]
+    successes = [o for o in outs if not o.get("__inference_error__", "")]
+
+    assert len(errors) == 6, f"Expected 6 errors, got {len(errors)}: {errors[:3]}..."
+    assert len(successes) == 54
+
+    for e in errors:
+        error_msg = e["__inference_error__"]
+        assert "ValueError" in error_msg, f"Expected ValueError in: {error_msg}"
+        assert (
+            "Intentional failure" in error_msg
+        ), f"Expected 'Intentional failure' in: {error_msg}"
+
+    for s in successes:
+        assert s.get("resp") is not None, f"Missing resp in success row: {s}"
+
+
+def test_completion_model(model_opt_125m, create_model_opt_125m_deployment):
+    deployment_name, app_name = create_model_opt_125m_deployment
+    config = ServeDeploymentProcessorConfig(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        dtype_mapping={
+            "CompletionRequest": CompletionRequest,
+        },
+        batch_size=16,
+        concurrency=1,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            method="completions",
+            dtype="CompletionRequest",
+            request_kwargs=dict(
+                model=model_opt_125m,
+                prompt=row["prompt"],
+                stream=False,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["choices"][0]["text"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"prompt": f"Hello {x['id']}"})
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+def test_multi_turn_completion_model(model_opt_125m, create_model_opt_125m_deployment):
+    deployment_name, app_name = create_model_opt_125m_deployment
+
+    config1 = ServeDeploymentProcessorConfig(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        dtype_mapping={
+            "CompletionRequest": CompletionRequest,
+        },
+        # Use lower batch size to reduce resource usage as there are multiple processors
+        batch_size=4,
+        concurrency=1,
+    )
+
+    processor1 = build_processor(
+        config1,
+        preprocess=lambda row: dict(
+            dtype="CompletionRequest",
+            method="completions",
+            request_kwargs=dict(
+                model=model_opt_125m,
+                prompt=row["prompt"],
+                stream=False,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            prompt=row["choices"][0]["text"],
+        ),
+    )
+
+    config2 = ServeDeploymentProcessorConfig(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        dtype_mapping={
+            "CompletionRequest": CompletionRequest,
+        },
+        batch_size=4,
+        concurrency=1,
+    )
+
+    processor2 = build_processor(
+        config2,
+        preprocess=lambda row: dict(
+            dtype="CompletionRequest",
+            method="completions",
+            request_kwargs=dict(
+                model=model_opt_125m,
+                prompt=row["prompt"],
+                stream=False,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["choices"][0]["text"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"prompt": f"Hello {x['id']}"})
+    ds = processor1(ds)
+    ds = processor2(ds)
+
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+def test_chat_model(model_opt_125m, create_model_opt_125m_deployment):
+    deployment_name, app_name = create_model_opt_125m_deployment
+    config = ServeDeploymentProcessorConfig(
+        deployment_name=deployment_name,
+        app_name=app_name,
+        dtype_mapping={
+            "ChatCompletionRequest": ChatCompletionRequest,
+        },
+        batch_size=16,
+        concurrency=1,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            dtype="ChatCompletionRequest",
+            method="chat",
+            request_kwargs=dict(
+                model=model_opt_125m,
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant"},
+                    {"role": "user", "content": f"Hello {row['id']}"},
+                ],
+                stream=False,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["choices"][0]["message"]["content"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"]})
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/gpu/processor/test_sglang_engine_proc.py

@@ -0,0 +1,126 @@
+"""This test suite does not need sglang to be installed."""
+
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray.data.llm import SGLangEngineProcessorConfig
+from ray.llm._internal.batch.constants import SGLangTaskType
+from ray.llm._internal.batch.processor import ProcessorBuilder
+from ray.llm._internal.batch.processor.sglang_engine_proc import (
+    build_sglang_engine_processor,
+)
+
+
+def test_sglang_engine_processor(gpu_type, model_llama_3_2_216M):
+    config = SGLangEngineProcessorConfig(
+        model_source=model_llama_3_2_216M,
+        engine_kwargs=dict(
+            context_length=8192,
+            tp_size=2,
+            dp_size=2,
+            disable_cuda_graph=True,
+            dtype="half",  # Older GPUs (e.g. T4) don't support bfloat16
+        ),
+        runtime_env=dict(
+            env_vars=dict(
+                RANDOM_ENV_VAR="12345",
+            ),
+        ),
+        accelerator_type=gpu_type,
+        concurrency=4,
+        batch_size=64,
+        max_concurrent_batches=4,
+        max_pending_requests=111,
+        chat_template_stage=True,
+        tokenize_stage=True,
+        detokenize_stage=True,
+    )
+    processor = ProcessorBuilder.build(config)
+    assert processor.list_stage_names() == [
+        "ChatTemplateStage",
+        "TokenizeStage",
+        "SGLangEngineStage",
+        "DetokenizeStage",
+    ]
+
+    stage = processor.get_stage_by_name("SGLangEngineStage")
+    assert stage.fn_constructor_kwargs == {
+        "model": model_llama_3_2_216M,
+        "engine_kwargs": {
+            "context_length": 8192,
+            "tp_size": 2,
+            "dp_size": 2,
+            "disable_cuda_graph": True,
+            "dtype": "half",
+            "task": SGLangTaskType.GENERATE,
+        },
+        "task_type": SGLangTaskType.GENERATE,
+        "max_pending_requests": 111,
+    }
+
+    runtime_env = stage.map_batches_kwargs.pop("runtime_env")
+    assert "env_vars" in runtime_env
+    assert runtime_env["env_vars"]["RANDOM_ENV_VAR"] == "12345"
+    compute = stage.map_batches_kwargs.pop("compute")
+    assert isinstance(compute, ray.data._internal.compute.ActorPoolStrategy)
+    assert stage.map_batches_kwargs == {
+        "zero_copy_batch": True,
+        "max_concurrency": 4,
+        "accelerator_type": gpu_type,
+        "num_gpus": 4,  # Based on tp_size=2, dp_size=2 in engine_kwargs
+    }
+
+
+class TestSGLangEngineProcessorConfig:
+    def test_build_processor_autoconfig_failure_with_trust_remote_code(self):
+        config = SGLangEngineProcessorConfig(
+            model_source="nonexistent-org/nonexistent-model",
+            engine_kwargs={"trust_remote_code": True},
+        )
+
+        processor = build_sglang_engine_processor(config)
+        assert processor is not None
+
+    def test_build_processor_import_error_with_trust_remote_code(self):
+        config = SGLangEngineProcessorConfig(
+            model_source="org/model-with-custom-code",
+            engine_kwargs={"trust_remote_code": True},
+        )
+
+        with (
+            patch(
+                "ray.llm._internal.batch.processor.sglang_engine_proc."
+                "download_model_files",
+                return_value="/tmp/fake_model_dir",
+            ),
+            patch(
+                "ray.llm._internal.batch.processor.sglang_engine_proc."
+                "transformers.AutoConfig.from_pretrained",
+                side_effect=ModuleNotFoundError("custom modeling module missing"),
+            ),
+        ):
+            processor = build_sglang_engine_processor(config)
+
+        assert processor is not None
+
+    def test_build_processor_download_error_with_trust_remote_code(self):
+        config = SGLangEngineProcessorConfig(
+            model_source="org/model-with-custom-code",
+            engine_kwargs={"trust_remote_code": True},
+        )
+
+        with patch(
+            "ray.llm._internal.batch.processor.sglang_engine_proc."
+            "download_model_files",
+            side_effect=RuntimeError("download failed"),
+        ):
+            processor = build_sglang_engine_processor(config)
+
+        assert processor is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/gpu/processor/test_vllm_engine_proc.py

@@ -0,0 +1,621 @@
+import sys
+
+import pydantic
+import pytest
+from transformers import AutoTokenizer
+
+import ray
+from ray.data.llm import build_processor, vLLMEngineProcessorConfig
+from ray.llm._internal.batch.constants import vLLMTaskType
+from ray.llm._internal.batch.processor import ProcessorBuilder
+from ray.llm._internal.batch.stages.configs import (
+    ChatTemplateStageConfig,
+    DetokenizeStageConfig,
+    PrepareMultimodalStageConfig,
+    TokenizerStageConfig,
+)
+
+
+@pytest.mark.parametrize(
+    "tensor_parallel_size, expected_distributed_executor_backend",
+    [(1, "uni"), (2, "ray")],
+)
+def test_vllm_engine_processor(
+    gpu_type,
+    model_opt_125m,
+    tensor_parallel_size,
+    expected_distributed_executor_backend,
+):
+    config = vLLMEngineProcessorConfig(
+        model_source=model_opt_125m,
+        engine_kwargs=dict(
+            max_model_len=8192,
+            tensor_parallel_size=tensor_parallel_size,
+        ),
+        runtime_env=dict(
+            env_vars=dict(
+                RANDOM_ENV_VAR="12345",
+            ),
+        ),
+        accelerator_type=gpu_type,
+        concurrency=4,
+        batch_size=64,
+        max_pending_requests=111,
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        detokenize_stage=DetokenizeStageConfig(enabled=True),
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=True),
+    )
+    processor = ProcessorBuilder.build(config)
+    assert processor.list_stage_names() == [
+        "PrepareMultimodalStage",
+        "ChatTemplateStage",
+        "TokenizeStage",
+        "vLLMEngineStage",
+        "DetokenizeStage",
+    ]
+
+    stage = processor.get_stage_by_name("vLLMEngineStage")
+    assert stage.fn_constructor_kwargs == {
+        "model": model_opt_125m,
+        "engine_kwargs": {
+            "max_model_len": 8192,
+            "distributed_executor_backend": expected_distributed_executor_backend,
+            "tensor_parallel_size": tensor_parallel_size,
+            "task_type": vLLMTaskType.GENERATE,
+        },
+        "task_type": vLLMTaskType.GENERATE,
+        "max_pending_requests": 111,
+        "dynamic_lora_loading_path": None,
+        "max_concurrent_batches": 8,
+        "batch_size": 64,
+        "should_continue_on_error": False,
+        "log_engine_metrics": True,
+    }
+
+    runtime_env = stage.map_batches_kwargs.pop("runtime_env")
+    assert "env_vars" in runtime_env
+    assert runtime_env["env_vars"]["RANDOM_ENV_VAR"] == "12345"
+    compute = stage.map_batches_kwargs.pop("compute")
+    assert isinstance(compute, ray.data._internal.compute.ActorPoolStrategy)
+
+    if expected_distributed_executor_backend == "ray":
+        ray_remote_args_fn = stage.map_batches_kwargs.pop("ray_remote_args_fn")
+        assert ray_remote_args_fn is not None
+        assert stage.map_batches_kwargs == {
+            "zero_copy_batch": True,
+            "max_concurrency": 8,
+            "accelerator_type": gpu_type,
+            "num_gpus": 0,
+        }
+    else:
+        assert "ray_remote_args_fn" not in stage.map_batches_kwargs
+        assert stage.map_batches_kwargs == {
+            "zero_copy_batch": True,
+            "max_concurrency": 8,
+            "accelerator_type": gpu_type,
+            "num_gpus": 1,
+        }
+
+
+def test_vllm_engine_processor_task_override(model_opt_125m):
+    config = vLLMEngineProcessorConfig(
+        model_source=model_opt_125m,
+        engine_kwargs=dict(
+            task_type=vLLMTaskType.EMBED,
+        ),
+        task_type=vLLMTaskType.GENERATE,
+        concurrency=4,
+        batch_size=64,
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        detokenize_stage=DetokenizeStageConfig(enabled=True),
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=True),
+    )
+    processor = ProcessorBuilder.build(config)
+    stage = processor.get_stage_by_name("vLLMEngineStage")
+
+    assert stage.fn_constructor_kwargs["task_type"] == vLLMTaskType.GENERATE
+    assert (
+        stage.fn_constructor_kwargs["engine_kwargs"]["task_type"]
+        == vLLMTaskType.GENERATE
+    )
+
+
+def test_vllm_engine_processor_invalid_task(model_opt_125m):
+    with pytest.raises(
+        pydantic.ValidationError, match="Invalid task type: invalid_task"
+    ):
+        vLLMEngineProcessorConfig(
+            model_source=model_opt_125m,
+            engine_kwargs=dict(
+                task_type=vLLMTaskType.EMBED,
+            ),
+            task_type="invalid_task",
+            concurrency=4,
+            batch_size=64,
+            chat_template_stage=ChatTemplateStageConfig(enabled=True),
+            tokenize_stage=TokenizerStageConfig(enabled=True),
+            detokenize_stage=DetokenizeStageConfig(enabled=True),
+            prepare_multimodal_stage=PrepareMultimodalStageConfig(enabled=True),
+        )
+
+
+def test_vllm_engine_processor_placement_group(gpu_type, model_opt_125m):
+    config = vLLMEngineProcessorConfig(
+        model_source=model_opt_125m,
+        engine_kwargs=dict(
+            max_model_len=8192,
+        ),
+        accelerator_type=gpu_type,
+        concurrency=4,
+        batch_size=64,
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        placement_group_config=dict(bundles=[{"CPU": 1, "GPU": 1}]),
+    )
+    processor = ProcessorBuilder.build(config)
+    stage = processor.get_stage_by_name("vLLMEngineStage")
+
+    stage.map_batches_kwargs.pop("runtime_env")
+    stage.map_batches_kwargs.pop("compute")
+
+    assert stage.map_batches_kwargs == {
+        "zero_copy_batch": True,
+        "max_concurrency": 8,
+        "accelerator_type": gpu_type,
+        "num_cpus": 1,
+        "num_gpus": 1,
+    }
+
+
+@pytest.mark.parametrize(
+    "engine_kwargs_extra,expected_num_bundles",
+    [
+        ({"tensor_parallel_size": 2}, 2),
+        (
+            {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
+            4,
+        ),
+        ({}, 1),  # Default case: tp=1, pp=1 → executor_backend="uni"
+    ],
+)
+def test_vllm_engine_processor_bundle_per_worker(
+    gpu_type, model_opt_125m, engine_kwargs_extra, expected_num_bundles
+):
+    """Test bundle_per_worker auto-expands based on tp*pp."""
+    engine_kwargs = dict(max_model_len=8192)
+    engine_kwargs.update(engine_kwargs_extra)
+    config = vLLMEngineProcessorConfig(
+        model_source=model_opt_125m,
+        engine_kwargs=engine_kwargs,
+        accelerator_type=gpu_type,
+        concurrency=4,
+        batch_size=64,
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        placement_group_config={"bundle_per_worker": {"CPU": 1, "GPU": 1}},
+    )
+    processor = ProcessorBuilder.build(config)
+    stage = processor.get_stage_by_name("vLLMEngineStage")
+
+    stage.map_batches_kwargs.pop("runtime_env")
+    stage.map_batches_kwargs.pop("compute")
+
+    expected_kwargs = {
+        "zero_copy_batch": True,
+        "max_concurrency": 8,
+        "accelerator_type": gpu_type,
+    }
+
+    if expected_num_bundles > 1:
+        # TP/PP > 1 -> executor_backend="ray"
+        ray_remote_args_fn = stage.map_batches_kwargs.pop("ray_remote_args_fn")
+        assert ray_remote_args_fn.args[0] == expected_num_bundles
+        assert ray_remote_args_fn.args[1] == gpu_type
+        expected_bundles = [{"CPU": 1, "GPU": 1}] * expected_num_bundles
+        assert ray_remote_args_fn.args[2]["bundles"] == expected_bundles
+        assert ray_remote_args_fn.args[2]["strategy"] == "PACK"
+        expected_kwargs["num_gpus"] = 0
+    else:
+        # TP=1, PP=1 -> executor_backend="uni"
+        expected_kwargs["num_cpus"] = 1
+        expected_kwargs["num_gpus"] = 1
+
+    assert stage.map_batches_kwargs == expected_kwargs
+
+
+def test_vllm_engine_processor_bundle_per_worker_conflict(gpu_type, model_opt_125m):
+    """Test that specifying both bundle_per_worker and bundles raises error."""
+    with pytest.raises(ValueError, match="Cannot specify both"):
+        vLLMEngineProcessorConfig(
+            model_source=model_opt_125m,
+            engine_kwargs=dict(max_model_len=8192),
+            accelerator_type=gpu_type,
+            placement_group_config={
+                "bundle_per_worker": {"CPU": 1, "GPU": 1},
+                "bundles": [{"CPU": 1, "GPU": 1}],
+            },
+        )
+
+
+def test_prepare_multimodal_stage_vllm_engine_processor(gpu_type, model_smolvlm_256m):
+    config = vLLMEngineProcessorConfig(
+        model_source=model_smolvlm_256m,
+        engine_kwargs=dict(
+            max_model_len=8192,
+        ),
+        accelerator_type=gpu_type,
+        concurrency=1,
+        batch_size=16,
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(
+            enabled=True,
+            model_config_kwargs=dict(
+                allowed_local_media_path="/tmp",
+            ),
+        ),
+    )
+    processor = ProcessorBuilder.build(config)
+
+    assert "PrepareMultimodalStage" in processor.list_stage_names()
+    stage = processor.get_stage_by_name("PrepareMultimodalStage")
+    fn_kwargs = stage.fn_constructor_kwargs
+
+    assert "model_config_kwargs" in fn_kwargs
+    model_config_kwargs = fn_kwargs["model_config_kwargs"]
+    assert model_config_kwargs["allowed_local_media_path"] == "/tmp"
+    assert model_config_kwargs["model"] == model_smolvlm_256m
+
+
+@pytest.mark.parametrize("backend", ["uni", "mp", "ray"])
+def test_generation_model(gpu_type, model_opt_125m, backend):
+    # OPT models don't have chat template, so we use ChatML template
+    # here to demonstrate the usage of custom chat template.
+    chat_template = """
+{% if messages[0]['role'] == 'system' %}
+    {% set offset = 1 %}
+{% else %}
+    {% set offset = 0 %}
+{% endif %}
+
+{{ bos_token }}
+{% for message in messages %}
+    {% if (message['role'] == 'user') != (loop.index0 % 2 == offset) %}
+        {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}
+    {% endif %}
+
+    {{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}
+{% endfor %}
+
+{% if add_generation_prompt %}
+    {{ '<|im_start|>assistant\n' }}
+{% endif %}
+    """
+
+    processor_config = vLLMEngineProcessorConfig(
+        model_source=model_opt_125m,
+        engine_kwargs=dict(
+            enable_prefix_caching=False,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=2048,
+            max_model_len=2048,
+            # Skip CUDA graph capturing to reduce startup time.
+            enforce_eager=True,
+            distributed_executor_backend=backend,
+        ),
+        batch_size=16,
+        accelerator_type=gpu_type,
+        concurrency=1,
+        chat_template_stage=ChatTemplateStageConfig(
+            enabled=True, chat_template=chat_template
+        ),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        detokenize_stage=DetokenizeStageConfig(enabled=True),
+    )
+
+    processor = build_processor(
+        processor_config,
+        preprocess=lambda row: dict(
+            messages=[
+                {"role": "system", "content": "You are a calculator"},
+                {"role": "user", "content": f"{row['id']} ** 3 = ?"},
+            ],
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=50,
+                detokenize=False,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+def test_generation_model_tokenized_prompt(gpu_type, model_opt_125m):
+    tokenizer = AutoTokenizer.from_pretrained(model_opt_125m, trust_remote_code=True)
+
+    processor_config = vLLMEngineProcessorConfig(
+        model_source=model_opt_125m,
+        engine_kwargs=dict(
+            enable_prefix_caching=False,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=2048,
+            max_model_len=2048,
+            enforce_eager=True,
+        ),
+        batch_size=16,
+        accelerator_type=gpu_type,
+        concurrency=1,
+        chat_template_stage=ChatTemplateStageConfig(enabled=False),
+        tokenize_stage=TokenizerStageConfig(enabled=False),
+        detokenize_stage=DetokenizeStageConfig(enabled=False),
+    )
+
+    def preprocess(row):
+        prompt_text = f"Calculate {row['id']} ** 3"
+
+        return dict(
+            tokenized_prompt=tokenizer(prompt_text)["input_ids"],
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=50,
+            ),
+        )
+
+    processor = build_processor(
+        processor_config,
+        preprocess=preprocess,
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+def test_embedding_model(gpu_type, model_smolvlm_256m):
+    processor_config = vLLMEngineProcessorConfig(
+        model_source=model_smolvlm_256m,
+        task_type="embed",
+        engine_kwargs=dict(
+            enable_prefix_caching=False,
+            enable_chunked_prefill=False,
+            max_model_len=2048,
+            # Skip CUDA graph capturing to reduce startup time.
+            enforce_eager=True,
+        ),
+        batch_size=16,
+        accelerator_type=gpu_type,
+        concurrency=1,
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        detokenize_stage=DetokenizeStageConfig(enabled=False),
+    )
+
+    processor = build_processor(
+        processor_config,
+        preprocess=lambda row: dict(
+            messages=[
+                {"role": "system", "content": "You are a calculator"},
+                {"role": "user", "content": f"{row['id']} ** 3 = ?"},
+            ],
+        ),
+        postprocess=lambda row: {
+            "resp": row["embeddings"],
+            "prompt": row["prompt"],
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+    assert all("prompt" in out for out in outs)
+
+
+def test_classification_model(gpu_type):
+    processor_config = vLLMEngineProcessorConfig(
+        model_source="HuggingFaceTB/fineweb-edu-classifier",
+        task_type="classify",
+        engine_kwargs=dict(
+            max_model_len=512,  # Model only supports up to 512 tokens
+        ),
+        batch_size=16,
+        accelerator_type=gpu_type,
+        concurrency=1,
+        chat_template_stage=ChatTemplateStageConfig(enabled=False),
+        tokenize_stage=TokenizerStageConfig(enabled=True),
+        detokenize_stage=DetokenizeStageConfig(enabled=False),
+    )
+
+    processor = build_processor(
+        processor_config,
+        preprocess=lambda row: dict(
+            prompt="This is a great educational content.",
+            tokenization_kwargs={"truncation": True, "max_length": 512},
+        ),
+        postprocess=lambda row: {
+            "probs": float(row["embeddings"][0])
+            if row.get("embeddings") is not None and len(row["embeddings"]) > 0
+            else None,
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("probs" in out for out in outs)
+
+
+@pytest.mark.parametrize("input_raw_image_data", [True, False])
+@pytest.mark.parametrize("decouple_tokenizer", [True, False])
+def test_vision_model(
+    gpu_type, model_smolvlm_256m, image_asset, input_raw_image_data, decouple_tokenizer
+):
+    image_url, image_pil = image_asset
+    llm_processor_config = vLLMEngineProcessorConfig(
+        model_source=model_smolvlm_256m,
+        task_type="generate",
+        engine_kwargs=dict(
+            # Skip CUDA graph capturing to reduce startup time.
+            enforce_eager=True,
+            # CI uses T4 GPU which does not support bfloat16.
+            dtype="half",
+            limit_mm_per_prompt={"image": 1},
+        ),
+        batch_size=16,
+        accelerator_type=gpu_type,
+        concurrency=1,
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(
+            enabled=True,
+            chat_template_content_format="openai",
+        ),
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=decouple_tokenizer),
+        detokenize_stage=DetokenizeStageConfig(enabled=decouple_tokenizer),
+    )
+    llm_processor = build_processor(
+        llm_processor_config,
+        preprocess=lambda row: dict(
+            messages=[
+                {"role": "system", "content": "You are an assistant"},
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"Say {row['val']} words about this image.",
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": image_url},
+                            "uuid": "image-1-id",  # UUID will not be included in the output as it's only used for internal caching
+                        }
+                        if input_raw_image_data
+                        else {
+                            "type": "image_pil",
+                            "image_pil": image_pil,
+                            "uuid": "image-2-id",
+                        },
+                    ],
+                },
+            ],
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=50,
+                detokenize=not decouple_tokenizer,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = llm_processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+@pytest.mark.parametrize("input_raw_audio_data", [True, False])
+def test_audio_model(
+    gpu_type, model_qwen_2_5_omni_3b, audio_asset, input_raw_audio_data
+):
+    audio_url, audio_data = audio_asset
+    llm_processor_config = vLLMEngineProcessorConfig(
+        model_source=model_qwen_2_5_omni_3b,
+        task_type="generate",
+        engine_kwargs=dict(
+            enforce_eager=True,
+            limit_mm_per_prompt={"audio": 1},
+        ),
+        batch_size=16,
+        accelerator_type=gpu_type,
+        concurrency=1,
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(
+            enabled=True,
+            chat_template_content_format="openai",
+        ),
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=False),
+        detokenize_stage=DetokenizeStageConfig(enabled=False),
+    )
+
+    llm_processor = build_processor(
+        llm_processor_config,
+        preprocess=lambda row: dict(
+            messages=[
+                {"role": "system", "content": "You are an assistant"},
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"Describe this audio in {row['val']} words.",
+                        },
+                        {
+                            "type": "input_audio",
+                            "input_audio": {"data": audio_data, "format": "wav"},
+                        }
+                        if input_raw_audio_data
+                        else {
+                            "type": "audio_url",
+                            "audio_url": {"url": audio_url},
+                        },
+                    ],
+                },
+            ],
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=50,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = llm_processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)
+
+
+class TestVLLMEngineProcessorConfig:
+    def test_build_processor_autoconfig_failure(self):
+        config = vLLMEngineProcessorConfig(
+            model_source="nonexistent-org/nonexistent-model",
+        )
+
+        processor = build_processor(config)
+        assert processor is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/gpu/stages/test_serve_deployment_stage.py

@@ -0,0 +1,542 @@
+import asyncio
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray.exceptions import RayActorError
+from ray.llm._internal.batch.stages.serve_deployment_stage import (
+    ServeDeploymentStageUDF,
+)
+from ray.serve._private.common import DeploymentID
+from ray.serve.exceptions import BackPressureError, DeploymentUnavailableError
+from ray.serve.llm.openai_api_models import ChatCompletionRequest, CompletionRequest
+
+
+@pytest.fixture
+def mock_serve_deployment_handle():
+    """Mock the serve deployment handle and its methods."""
+    with patch("ray.serve.get_deployment_handle") as mock_get_handle:
+        mock_handle = MagicMock()
+        mock_handle.options.return_value = mock_handle
+
+        # Mock the chat and completions methods
+        mock_handle.chat = MagicMock()
+        mock_handle.completions = MagicMock()
+
+        mock_get_handle.return_value = mock_handle
+        yield mock_handle
+
+
+@pytest.mark.parametrize(
+    "method,test_data",
+    [
+        (
+            "completions",
+            [
+                {
+                    "method": "completions",
+                    "dtype": "CompletionRequest",
+                    "request_kwargs": {"prompt": "Hello", "temperature": 0.7},
+                },
+            ],
+        ),
+        (
+            "chat",
+            [
+                {
+                    "method": "chat",
+                    "dtype": "ChatCompletionRequest",
+                    "request_kwargs": {
+                        "messages": [
+                            {
+                                "role": "system",
+                                "content": "You are a helpful assistant",
+                            },
+                            {"role": "user", "content": "Hello 1"},
+                        ]
+                    },
+                },
+            ],
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_serve_deployment_udf_methods(
+    mock_serve_deployment_handle, method, test_data
+):
+    """Test both completions and chat methods."""
+    # Create a mock response that will be returned directly
+    mock_response = {"test": "response"}
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            yield mock_response
+
+        return mock_async_iterator()
+
+    getattr(mock_serve_deployment_handle, method).remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={
+            "CompletionRequest": CompletionRequest,
+            "ChatCompletionRequest": ChatCompletionRequest,
+        },
+    )
+
+    batch = {"__data": test_data}
+
+    responses = []
+    async for response in udf(batch):
+        responses.append(response)
+
+    assert len(responses) == 1
+    assert "__data" in responses[0]
+    assert len(responses[0]["__data"]) == len(test_data)
+
+    for i, item in enumerate(responses[0]["__data"]):
+        assert "batch_uuid" in item
+        assert "time_taken" in item
+        assert item["request_id"] == str(i)
+        assert "test" in item  # From the mock response
+
+    assert getattr(mock_serve_deployment_handle, method).remote.call_count == len(
+        test_data
+    )
+
+
+@pytest.mark.asyncio
+async def test_serve_deployment_invalid_method(mock_serve_deployment_handle):
+    """Test that invalid method raises error at runtime."""
+    # Set up the mock to simulate a method that doesn't exist
+    mock_serve_deployment_handle.invalid_method = None
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={
+            "CompletionRequest": CompletionRequest,
+        },
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "invalid_method",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "Hello", "temperature": 0.7},
+            }
+        ]
+    }
+
+    with pytest.raises(
+        ValueError, match="Method invalid_method not found in the serve deployment."
+    ):
+        async for _ in udf(batch):
+            pass
+
+
+@pytest.mark.parametrize(
+    "dtype_mapping", [None, {"ChatCompletionRequest": ChatCompletionRequest}]
+)
+@pytest.mark.asyncio
+async def test_serve_deployment_missing_dtype(
+    mock_serve_deployment_handle, dtype_mapping
+):
+    """Test that missing dtype raises error at runtime."""
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping=dtype_mapping,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "Hello", "temperature": 0.7},
+            }
+        ]
+    }
+
+    with pytest.raises(
+        ValueError,
+        match="CompletionRequest must be provided in ServeDeploymentProcessorConfig's dtype_mapping.",
+    ):
+        async for _ in udf(batch):
+            pass
+
+
+# ============================================================================
+# Error handling tests for should_continue_on_error
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_default_raises_on_error(mock_serve_deployment_handle):
+    """Default behavior (should_continue_on_error=False) raises on inference error."""
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            raise ValueError("prompt too long")
+            yield  # Make it a generator
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=False,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test", "temperature": 0.7},
+            }
+        ]
+    }
+
+    with pytest.raises(ValueError, match="prompt too long"):
+        async for _ in udf(batch):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_continue_on_error_yields_error_row(
+    mock_serve_deployment_handle,
+):
+    """With should_continue_on_error=True, errors yield rows with __inference_error__."""
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            raise ValueError("prompt too long")
+            yield  # Make it a generator
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=True,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test prompt", "temperature": 0.7},
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 1
+    assert "__inference_error__" in results[0]
+    assert "ValueError" in results[0]["__inference_error__"]
+    assert "prompt too long" in results[0]["__inference_error__"]
+    # Error rows include request_kwargs snippet for debuggability
+    assert "request_kwargs" in results[0]
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_mixed_success_and_error(mock_serve_deployment_handle):
+    """Mixed batch: some rows succeed, some fail."""
+    call_count = 0
+
+    def mock_remote_call(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        current_call = call_count
+
+        async def mock_async_iterator():
+            # Second call fails
+            if current_call == 2:
+                raise ValueError("prompt too long")
+            yield {"generated_text": f"Response {current_call}"}
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=True,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "first", "temperature": 0.7},
+            },
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "second", "temperature": 0.7},
+            },
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "third", "temperature": 0.7},
+            },
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 3
+
+    errors = [r for r in results if r.get("__inference_error__", "") != ""]
+    successes = [r for r in results if r.get("__inference_error__", "") == ""]
+
+    assert len(errors) == 1
+    assert len(successes) == 2
+    assert "ValueError" in errors[0]["__inference_error__"]
+
+
+@pytest.mark.parametrize(
+    "fatal_error",
+    [
+        RayActorError(error_msg="Actor died"),
+        BackPressureError(num_queued_requests=100, max_queued_requests=50),
+        DeploymentUnavailableError(
+            deployment_id=DeploymentID(name="test", app_name="test_app")
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_serve_udf_fatal_errors_always_propagate(
+    mock_serve_deployment_handle, fatal_error
+):
+    """Fatal errors (RayActorError, BackPressureError, etc.) always propagate."""
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            raise fatal_error
+            yield  # Make it a generator
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=True,  # Even with this True, fatal errors propagate
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test", "temperature": 0.7},
+            }
+        ]
+    }
+
+    with pytest.raises(type(fatal_error)):
+        async for _ in udf(batch):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_unknown_errors_propagate(mock_serve_deployment_handle):
+    """Unknown errors propagate even with should_continue_on_error=True."""
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            raise RuntimeError("unexpected system error")
+            yield
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=True,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test", "temperature": 0.7},
+            }
+        ]
+    }
+
+    with pytest.raises(RuntimeError, match="unexpected system error"):
+        async for _ in udf(batch):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_success_with_continue_on_error_includes_none_error(
+    mock_serve_deployment_handle,
+):
+    """Successful rows with should_continue_on_error=True have __inference_error__=None."""
+    mock_response = {"generated_text": "Hello!"}
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            yield mock_response
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=True,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test", "temperature": 0.7},
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 1
+    assert results[0]["__inference_error__"] == ""
+    assert results[0]["generated_text"] == "Hello!"
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_timeout_raises_by_default(mock_serve_deployment_handle):
+    """With a request_timeout_s and default error handling, a slow request raises."""
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            await asyncio.sleep(10)
+            yield {"generated_text": "too late"}
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        request_timeout_s=0.05,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test", "temperature": 0.7},
+            }
+        ]
+    }
+
+    with pytest.raises(TimeoutError, match="timed out after 0.05s"):
+        async for _ in udf(batch):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_serve_udf_timeout_recovers_with_continue_on_error(
+    mock_serve_deployment_handle,
+):
+    """A timed-out request becomes an error row when should_continue_on_error=True."""
+
+    def mock_remote_call(*args, **kwargs):
+        async def mock_async_iterator():
+            await asyncio.sleep(10)
+            yield {"generated_text": "too late"}
+
+        return mock_async_iterator()
+
+    mock_serve_deployment_handle.completions.remote.side_effect = mock_remote_call
+
+    udf = ServeDeploymentStageUDF(
+        data_column="__data",
+        expected_input_keys=["method", "request_kwargs"],
+        deployment_name="test_deployment",
+        app_name="test_app",
+        dtype_mapping={"CompletionRequest": CompletionRequest},
+        should_continue_on_error=True,
+        request_timeout_s=0.05,
+    )
+
+    batch = {
+        "__data": [
+            {
+                "method": "completions",
+                "dtype": "CompletionRequest",
+                "request_kwargs": {"prompt": "test prompt", "temperature": 0.7},
+            }
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 1
+    assert "TimeoutError" in results[0]["__inference_error__"]
+    assert "request_kwargs" in results[0]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/gpu/stages/test_sglang_engine_stage.py

@@ -0,0 +1,280 @@
+"""This test suite does not need sglang to be installed."""
+
+import asyncio
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from ray.data import ActorPoolStrategy
+from ray.llm._internal.batch.stages.sglang_engine_stage import (
+    SGLangEngineStage,
+    SGLangEngineStageUDF,
+    SGLangEngineWrapper,
+    SGLangTaskType,
+)
+
+
+@pytest.fixture
+def mock_sglang_wrapper():
+    with patch(
+        "ray.llm._internal.batch.stages.sglang_engine_stage.SGLangEngineWrapper"
+    ) as mock_wrapper:
+        # Create a mock instance that will be returned by the wrapper class
+        mock_instance = MagicMock()
+        mock_instance.generate_async = AsyncMock()
+        mock_instance.shutdown = MagicMock()
+
+        # Configure the mock instance's behavior
+        async def mock_generate(row):
+            return (
+                MagicMock(
+                    request_id=0,
+                    prompt=row["prompt"],
+                    prompt_token_ids=None,
+                    params=row["sampling_params"],
+                    idx_in_batch=row["__idx_in_batch"],
+                ),
+                {
+                    "prompt": row["prompt"],
+                    "prompt_token_ids": None,
+                    "num_input_tokens": 3,
+                    "generated_tokens": None,
+                    "generated_text": f"Response to: {row['prompt']}",
+                    "num_generated_tokens": 3,
+                },
+                0.1,  # time_taken_llm
+            )
+
+        mock_instance.generate_async.side_effect = mock_generate
+
+        # Make the wrapper class return our mock instance
+        mock_wrapper.return_value = mock_instance
+        yield mock_wrapper
+
+
+@pytest.fixture
+def mock_sgl_engine():
+    """Mock the SGLang engine and its _generate_async method."""
+    with (
+        patch(
+            "ray.llm._internal.batch.stages.sglang_engine_stage.SGLangEngineWrapper._generate_async"
+        ) as mock_generate_async,
+    ):
+        try:
+            import sglang  # noqa: F401
+        except ImportError:
+            # Mock sglang module if it's not installed in test env.
+            mock_sgl = MagicMock()
+            mock_sgl.Engine = AsyncMock()
+            sys.modules["sglang"] = mock_sgl
+        num_running_requests = 0
+        request_lock = asyncio.Lock()
+
+        # Configure mock engine's generate behavior to simulate delay
+        async def mock_generate(request):
+            nonlocal num_running_requests
+            async with request_lock:
+                num_running_requests += 1
+
+            # This will be checked in tests that use max_pending_requests
+            max_pending_requests = getattr(mock_generate, "max_pending_requests", -1)
+            if max_pending_requests > 0:
+                assert num_running_requests <= max_pending_requests
+
+            await asyncio.sleep(0.1)  # Reduced sleep time for faster tests
+
+            async with request_lock:
+                num_running_requests -= 1
+
+            # Create a mock SGLang output
+            return {
+                "prompt": request.prompt,
+                "prompt_token_ids": None,
+                "text": f"Response to: {request.prompt}",
+                "meta_info": {
+                    "prompt_tokens": 3,
+                    "completion_tokens": request.params.get("max_new_tokens", 3),
+                    "finish_reason": "stop",
+                },
+                "output_ids": [4, 5, 6],
+            }
+
+        mock_generate_async.side_effect = mock_generate
+        yield mock_generate_async
+
+
+def test_sglang_engine_stage_post_init(gpu_type, model_llama_3_2_216M):
+    stage = SGLangEngineStage(
+        fn_constructor_kwargs=dict(
+            model=model_llama_3_2_216M,
+            engine_kwargs=dict(
+                tp_size=2,
+                dp_size=2,
+            ),
+            task_type=SGLangTaskType.GENERATE,
+            max_pending_requests=10,
+        ),
+        map_batches_kwargs=dict(
+            zero_copy_batch=True,
+            compute=ActorPoolStrategy(size=1),
+            max_concurrency=4,
+            accelerator_type=gpu_type,
+        ),
+    )
+
+    assert stage.fn_constructor_kwargs == {
+        "model": model_llama_3_2_216M,
+        "task_type": SGLangTaskType.GENERATE,
+        "max_pending_requests": 10,
+        "engine_kwargs": {
+            "tp_size": 2,
+            "dp_size": 2,
+        },
+    }
+
+    compute = stage.map_batches_kwargs.pop("compute")
+    assert isinstance(compute, ActorPoolStrategy)
+    assert compute.min_size == 1
+    assert compute.max_size == 1
+
+    assert stage.map_batches_kwargs == {
+        "zero_copy_batch": True,
+        "max_concurrency": 4,
+        "accelerator_type": gpu_type,
+        "num_gpus": 4,
+    }
+
+
+@pytest.mark.asyncio
+async def test_sglang_engine_udf_basic(mock_sglang_wrapper, model_llama_3_2_216M):
+    # Create UDF instance - it will use the mocked wrapper
+    udf = SGLangEngineStageUDF(
+        data_column="__data",
+        expected_input_keys=["prompt", "sampling_params"],
+        model=model_llama_3_2_216M,
+        task_type=SGLangTaskType.GENERATE,
+        engine_kwargs={
+            # Test that this should be overridden by the stage.
+            "model": "random-model",
+            # When reaching SGLangEngineStageUDF, this kwargs has been inserted by the processor
+            # even though it is unset in the processor config.
+            "task": SGLangTaskType.GENERATE,
+        },
+    )
+
+    assert udf.model is not None
+    assert udf.task_type == SGLangTaskType.GENERATE
+    assert udf.engine_kwargs["task"] == SGLangTaskType.GENERATE
+    assert udf.max_pending_requests == -1  # Default value for SGLang
+
+    # Test batch processing
+    batch = {
+        "__data": [
+            {"prompt": "Hello", "sampling_params": {"temperature": 0.7}},
+            {"prompt": "World", "sampling_params": {"temperature": 0.7}},
+        ]
+    }
+
+    responses = []
+    async for response in udf(batch):
+        responses.extend(response["__data"])
+
+    assert len(responses) == 2
+    assert all("batch_uuid" in r for r in responses)
+    assert all("time_taken_llm" in r for r in responses)
+    # The output order is not guaranteed.
+    assert responses[0]["prompt"] in ["Hello", "World"]
+    assert responses[1]["prompt"] in ["Hello", "World"]
+    assert responses[0]["prompt"] != responses[1]["prompt"]
+
+    # Verify the wrapper was constructed with correct arguments
+    mock_sglang_wrapper.assert_called_once_with(
+        model=udf.model,
+        idx_in_batch_column="__idx_in_batch",
+        max_pending_requests=-1,
+        task=SGLangTaskType.GENERATE,
+    )
+
+
+@pytest.mark.parametrize("max_pending_requests,batch_size", [(2, 10), (-1, 5)])
+@pytest.mark.asyncio
+async def test_sglang_wrapper(
+    mock_sgl_engine, model_llama_3_2_216M, max_pending_requests, batch_size
+):
+    """Test the SGLang wrapper with different configurations."""
+    mock_generate_async = mock_sgl_engine
+
+    # Set the max_pending_requests for assertion in the mock
+    mock_generate_async.side_effect.max_pending_requests = max_pending_requests
+
+    # Create wrapper with configured max_pending_requests
+    wrapper = SGLangEngineWrapper(
+        model=model_llama_3_2_216M,
+        idx_in_batch_column="__idx_in_batch",
+        max_pending_requests=max_pending_requests,
+        skip_tokenizer_init=False,
+    )
+
+    # Create batch requests with different sampling parameters
+    batch = [
+        {
+            "__idx_in_batch": i,
+            "prompt": f"Test {i}",
+            "sampling_params": {
+                "max_new_tokens": i + 5,
+                "temperature": 0.7,
+            },
+        }
+        for i in range(batch_size)
+    ]
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+    results = await asyncio.gather(*tasks)
+
+    # Verify all requests were processed
+    assert mock_generate_async.call_count == batch_size
+
+    # Verify the outputs match expected values
+    for i, (request, output, time_taken_llm) in enumerate(results):
+        assert output["prompt"] == f"Test {i}"
+        assert output["num_generated_tokens"] == i + 5  # max_new_tokens we set
+        assert time_taken_llm > 0
+
+
+@pytest.mark.asyncio
+async def test_sglang_error_handling(model_llama_3_2_216M):
+    """Test error handling when SGLang is not available."""
+    with patch.dict(sys.modules, {"sglang": None}):
+        with pytest.raises(ImportError, match="SGLang is not installed"):
+            SGLangEngineWrapper(
+                model=model_llama_3_2_216M,
+                idx_in_batch_column="__idx_in_batch",
+            )
+
+
+@pytest.mark.asyncio
+async def test_sglang_invalid_task_type(model_llama_3_2_216M, mock_sgl_engine):
+    """Test handling of invalid task types."""
+    wrapper = SGLangEngineWrapper(
+        model=model_llama_3_2_216M,
+        idx_in_batch_column="__idx_in_batch",
+        task=SGLangTaskType.GENERATE,
+    )
+
+    # Create a task type that doesn't exist in the prepare_llm_request method
+    invalid_task_type = "invalid_task"
+    wrapper.task_type = invalid_task_type
+
+    with pytest.raises(ValueError, match=f"Unsupported task type: {invalid_task_type}"):
+        await wrapper._prepare_llm_request(
+            {
+                "prompt": "Hello",
+                "sampling_params": {"temperature": 0.7},
+                "__idx_in_batch": 0,
+            }
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/gpu/stages/test_vllm_engine_stage.py

@@ -0,0 +1,992 @@
+import asyncio
+import json
+import math
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from pydantic import BaseModel
+
+from ray.data import ActorPoolStrategy
+from ray.llm._internal.batch.constants import vLLMTaskType
+from ray.llm._internal.batch.stages.vllm_engine_stage import (
+    vLLMEngineRequest,
+    vLLMEngineStage,
+    vLLMEngineStageUDF,
+    vLLMEngineWrapper,
+    vLLMOutputData,
+)
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+
+@pytest.fixture
+def mock_vllm_wrapper():
+    with patch(
+        "ray.llm._internal.batch.stages.vllm_engine_stage.vLLMEngineWrapper"
+    ) as mock_wrapper:
+        # Create a mock instance that will be returned by the wrapper class
+        mock_instance = MagicMock()
+        mock_instance.generate_async = AsyncMock()
+        mock_instance.shutdown = MagicMock()
+
+        # Configure the mock instance's behavior
+        async def mock_generate(row):
+            return (
+                MagicMock(
+                    request_id=0,
+                    prompt=row["prompt"],
+                    prompt_token_ids=None,
+                    multimodal_data=None,
+                    params=row["sampling_params"],
+                    idx_in_batch=row["__idx_in_batch"],
+                ),
+                {
+                    "prompt": row["prompt"],
+                    "prompt_token_ids": [1, 2, 3],
+                    "num_input_tokens": 3,
+                    "generated_tokens": [4, 5, 6],
+                    "generated_text": f"Response to: {row['prompt']}",
+                    "num_generated_tokens": 3,
+                    "time_per_token": 0.1,
+                },
+                0.1,  # time_taken_llm
+            )
+
+        mock_instance.generate_async.side_effect = mock_generate
+
+        # Configure the scheduler config mock
+        mock_scheduler_config = MagicMock()
+        mock_scheduler_config.max_num_seqs = 128  # Current vLLM default
+        mock_instance.get_scheduler_config.return_value = mock_scheduler_config
+
+        # Configure the engine mock
+        mock_engine = MagicMock()
+        mock_engine.do_log_stats = AsyncMock()
+        mock_instance.engine = mock_engine
+
+        # Make the wrapper class return our mock instance
+        mock_wrapper.return_value = mock_instance
+        yield mock_wrapper
+
+
+def test_vllm_engine_stage_post_init(gpu_type, model_llama_3_2_216M):
+    stage = vLLMEngineStage(
+        fn_constructor_kwargs=dict(
+            model=model_llama_3_2_216M,
+            engine_kwargs=dict(
+                tensor_parallel_size=2,
+                pipeline_parallel_size=2,
+                distributed_executor_backend="ray",
+            ),
+            task_type=vLLMTaskType.GENERATE,
+            max_pending_requests=10,
+        ),
+        map_batches_kwargs=dict(
+            zero_copy_batch=True,
+            compute=ActorPoolStrategy(size=1),
+            max_concurrency=4,
+            accelerator_type=gpu_type,
+        ),
+    )
+
+    assert stage.fn_constructor_kwargs == {
+        "model": model_llama_3_2_216M,
+        "task_type": vLLMTaskType.GENERATE,
+        "max_pending_requests": 10,
+        "engine_kwargs": {
+            "tensor_parallel_size": 2,
+            "pipeline_parallel_size": 2,
+            "distributed_executor_backend": "ray",
+        },
+    }
+    ray_remote_args_fn = stage.map_batches_kwargs.pop("ray_remote_args_fn")
+    compute = stage.map_batches_kwargs.pop("compute")
+    assert isinstance(compute, ActorPoolStrategy)
+    assert compute.min_size == 1
+    assert compute.max_size == 1
+
+    assert stage.map_batches_kwargs == {
+        "zero_copy_batch": True,
+        "max_concurrency": 4,
+        "accelerator_type": gpu_type,
+        "num_gpus": 0,
+    }
+    scheduling_strategy = ray_remote_args_fn()["scheduling_strategy"]
+    assert isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy)
+
+    bundle_specs = scheduling_strategy.placement_group.bundle_specs
+    assert len(bundle_specs) == 4
+    for bundle_spec in bundle_specs:
+        assert bundle_spec[f"accelerator_type:{gpu_type}"] == 0.001
+        assert bundle_spec["CPU"] == 1.0
+        assert bundle_spec["GPU"] == 1.0
+
+
+@pytest.mark.asyncio
+async def test_vllm_engine_udf_basic(mock_vllm_wrapper, model_llama_3_2_216M):
+    # Simulate vLLM's resolved state when the user sets `max_num_seqs=100`:
+    # the wrapper owns the resolution of `max_pending_requests`, and the UDF
+    # reads the resolved value back.
+    expected_max_pending_requests = math.ceil(100 * 1.1)
+    mock_vllm_wrapper.return_value.max_pending_requests = expected_max_pending_requests
+
+    # Create UDF instance - it will use the mocked wrapper
+    udf = vLLMEngineStageUDF(
+        data_column="__data",
+        expected_input_keys=["prompt", "sampling_params"],
+        model=model_llama_3_2_216M,
+        task_type=vLLMTaskType.GENERATE,
+        batch_size=32,
+        max_concurrent_batches=4,
+        engine_kwargs={
+            # Test that this should be overridden by the stage.
+            "model": "random-model",
+            # This is overridden in the processor, so it remains unchanged when we bypass
+            # the processor and pass it directly to the stage via vLLMEngineStageUDF.
+            "task_type": vLLMTaskType.EMBED,
+            "max_num_seqs": 100,
+            "disable_log_stats": False,
+        },
+    )
+
+    assert udf.model == model_llama_3_2_216M
+    assert udf.task_type == vLLMTaskType.GENERATE
+    assert udf.engine_kwargs["task_type"] == vLLMTaskType.EMBED
+    assert udf.engine_kwargs["max_num_seqs"] == 100
+    assert udf.max_pending_requests == expected_max_pending_requests
+
+    # Test batch processing
+    batch = {
+        "__data": [
+            {"prompt": "Hello", "sampling_params": {"temperature": 0.7}},
+            {"prompt": "World", "sampling_params": {"temperature": 0.7}},
+        ]
+    }
+
+    responses = []
+    async for response in udf(batch):
+        responses.extend(response["__data"])
+
+    assert len(responses) == 2
+    assert all("batch_uuid" in r for r in responses)
+    assert all("time_taken_llm" in r for r in responses)
+    # The output order is not guaranteed.
+    assert responses[0]["prompt"] in ["Hello", "World"]
+    assert responses[1]["prompt"] in ["Hello", "World"]
+    assert responses[0]["prompt"] != responses[1]["prompt"]
+
+    # Verify the wrapper was constructed with correct arguments. The UDF
+    # passes `max_pending_requests=None` straight through when the caller
+    # doesn't supply it; the wrapper resolves the default from vLLM's
+    # resolved engine config.
+    mock_vllm_wrapper.assert_called_once_with(
+        model=model_llama_3_2_216M,
+        model_source=model_llama_3_2_216M,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=False,
+        max_pending_requests=None,
+        task_type=vLLMTaskType.EMBED,
+        max_num_seqs=100,
+        dynamic_lora_loading_path=None,
+        enable_log_requests=False,
+        log_engine_metrics=True,
+    )
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_semaphore(model_llama_3_2_216M):
+    from vllm.outputs import CompletionOutput, RequestOutput
+
+    max_pending_requests = 2
+
+    with (
+        patch("vllm.AsyncLLMEngine") as mock_engine,
+        patch(
+            "ray.llm._internal.batch.stages.vllm_engine_stage.vLLMEngineWrapper._generate_async"
+        ) as mock_generate_async,
+    ):
+        mock_engine.from_engine_args.return_value = AsyncMock()
+        num_running_requests = 0
+        request_lock = asyncio.Lock()
+
+        # Configure mock engine's generate behavior to simulate delay
+        async def mock_generate(request):
+            nonlocal num_running_requests
+            async with request_lock:
+                num_running_requests += 1
+
+            assert num_running_requests <= max_pending_requests
+            await asyncio.sleep(0.3)
+
+            async with request_lock:
+                num_running_requests -= 1
+
+            return RequestOutput(
+                request_id="test",
+                prompt="test",
+                prompt_token_ids=[1, 2, 3],
+                prompt_logprobs=None,
+                metrics=None,
+                outputs=[
+                    CompletionOutput(
+                        index=0,
+                        text="test response",
+                        token_ids=[4, 5, 6],
+                        cumulative_logprob=None,
+                        logprobs=None,
+                    )
+                ],
+                finished=True,
+            )
+
+        mock_generate_async.side_effect = mock_generate
+
+        # Create wrapper with max 2 pending requests
+        wrapper = vLLMEngineWrapper(
+            model=model_llama_3_2_216M,
+            model_source=model_llama_3_2_216M,
+            idx_in_batch_column="__idx_in_batch",
+            disable_log_stats=True,
+            max_pending_requests=max_pending_requests,
+        )
+
+        # Create 10 requests
+        batch = [
+            {"__idx_in_batch": i, "prompt": f"Test {i}", "sampling_params": {}}
+            for i in range(10)
+        ]
+
+        tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+        await asyncio.gather(*tasks)
+
+        # Verify all requests were processed
+        assert mock_generate_async.call_count == 10
+
+        # Clean up GPU memory
+        wrapper.shutdown()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "task_type",
+    [
+        vLLMTaskType.GENERATE,
+        vLLMTaskType.EMBED,
+        vLLMTaskType.CLASSIFY,
+        vLLMTaskType.SCORE,
+    ],
+)
+async def test_vllm_wrapper_forwards_lora_request(task_type):
+    """Regression test: lora_request must be forwarded to the vLLM engine.
+
+    A prior bug populated vLLMEngineRequest.lora_request correctly but never
+    passed it to the vLLM engine, causing per-row LoRA adapters to be
+    silently dropped. GENERATE dispatches to engine.generate(); pooling task
+    types (EMBED / CLASSIFY / SCORE) dispatch to engine.encode().
+    """
+
+    async def finished_stream():
+        output = MagicMock()
+        output.finished = True
+        yield output
+
+    wrapper = vLLMEngineWrapper.__new__(vLLMEngineWrapper)
+    wrapper.engine = MagicMock()
+    wrapper.engine.generate = MagicMock(return_value=finished_stream())
+    wrapper.engine.encode = MagicMock(return_value=finished_stream())
+    wrapper.task_type = task_type
+
+    sentinel_lora = object()
+    request = vLLMEngineRequest(
+        request_id=0,
+        idx_in_batch=0,
+        prompt="hello",
+        prompt_token_ids=None,
+        multimodal_data=None,
+        mm_processor_kwargs=None,
+        multimodal_uuids=None,
+        params=MagicMock(),
+        tokenization_kwargs=None,
+        lora_request=sentinel_lora,
+    )
+
+    await wrapper._generate_async(request)
+
+    expected = (
+        wrapper.engine.generate
+        if task_type == vLLMTaskType.GENERATE
+        else wrapper.engine.encode
+    )
+    assert expected.call_args.kwargs.get("lora_request") is sentinel_lora
+
+
+def _make_bare_wrapper():
+    """Build a vLLMEngineWrapper without invoking __init__ (which boots vLLM)."""
+    wrapper = vLLMEngineWrapper.__new__(vLLMEngineWrapper)
+    wrapper.request_id = 0
+    wrapper.idx_in_batch_column = "__idx_in_batch"
+    wrapper.task_type = vLLMTaskType.GENERATE
+    wrapper._image_row_column_warning_logged = False
+    wrapper.model = "test-model"
+    wrapper.lora_lock = asyncio.Lock()
+    wrapper.lora_name_to_request = {}
+    return wrapper
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_legacy_image_warns_and_routes():
+    """Legacy `image` row column should warn once and be routed to multimodal_data."""
+    wrapper = _make_bare_wrapper()
+
+    sentinel_image = object()
+    row = {
+        "__idx_in_batch": 0,
+        "prompt": "hi",
+        "image": [sentinel_image],
+        "sampling_params": {"max_tokens": 1, "temperature": 0.0},
+    }
+
+    with patch(
+        "ray.llm._internal.batch.stages.vllm_engine_stage.logger.warning"
+    ) as mock_warning:
+        request = await wrapper._prepare_llm_request(row)
+
+    assert request.multimodal_data == {"image": [sentinel_image]}
+    assert any(
+        "image" in str(call.args[0]) and "deprecated" in str(call.args[0]).lower()
+        for call in mock_warning.call_args_list
+    )
+    assert wrapper._image_row_column_warning_logged is True
+
+    # Second call must not log the deprecation warning again.
+    with patch(
+        "ray.llm._internal.batch.stages.vllm_engine_stage.logger.warning"
+    ) as mock_warning2:
+        await wrapper._prepare_llm_request(
+            {
+                "__idx_in_batch": 1,
+                "prompt": "hi again",
+                "image": [sentinel_image],
+                "sampling_params": {"max_tokens": 1, "temperature": 0.0},
+            }
+        )
+    assert not any(
+        "deprecated" in str(call.args[0]).lower()
+        for call in mock_warning2.call_args_list
+    )
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_legacy_image_merges_into_existing_multimodal_data():
+    """Legacy `image` should merge into an explicit multimodal_data dict."""
+    wrapper = _make_bare_wrapper()
+
+    img = object()
+    audio = object()
+    row = {
+        "__idx_in_batch": 0,
+        "prompt": "hi",
+        "image": [img],
+        "multimodal_data": {"audio": [audio]},
+        "sampling_params": {"max_tokens": 1, "temperature": 0.0},
+    }
+
+    request = await wrapper._prepare_llm_request(row)
+    assert request.multimodal_data == {"audio": [audio], "image": [img]}
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_legacy_image_conflict_with_multimodal_data_raises():
+    """Setting both legacy `image` and multimodal_data['image'] must raise."""
+    wrapper = _make_bare_wrapper()
+
+    legacy_img = object()
+    modern_img = object()
+    row = {
+        "__idx_in_batch": 0,
+        "prompt": "hi",
+        "image": [legacy_img],
+        "multimodal_data": {"image": [modern_img]},
+        "sampling_params": {"max_tokens": 1, "temperature": 0.0},
+    }
+
+    with pytest.raises(ValueError, match="multimodal_data"):
+        await wrapper._prepare_llm_request(row)
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_legacy_image_empty_list_is_noop():
+    """Empty legacy `image=[]` should be skipped, not merged or conflict."""
+    wrapper = _make_bare_wrapper()
+
+    modern_img = object()
+    row = {
+        "__idx_in_batch": 0,
+        "prompt": "hi",
+        "image": [],
+        "multimodal_data": {"image": [modern_img]},
+        "sampling_params": {"max_tokens": 1, "temperature": 0.0},
+    }
+
+    request = await wrapper._prepare_llm_request(row)
+    assert request.multimodal_data == {"image": [modern_img]}
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_generate(model_llama_3_2_216M):
+    # TODO: Test v1 engine. The issue is once vLLM is imported with v0,
+    # we cannot configure it to use v1, so we need a separate test for v1.
+
+    wrapper = vLLMEngineWrapper(
+        model=model_llama_3_2_216M,
+        model_source=model_llama_3_2_216M,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=True,
+        max_pending_requests=10,
+        # Skip CUDA graph capturing to reduce the start time.
+        enforce_eager=True,
+        gpu_memory_utilization=0.8,
+        max_model_len=2048,
+        task_type=vLLMTaskType.GENERATE,
+        # Older GPUs (e.g. T4) don't support bfloat16.
+        dtype="half",
+    )
+
+    batch = [
+        {
+            "__idx_in_batch": 0,
+            "prompt": "Hello",
+            "sampling_params": {
+                "max_tokens": 10,
+                "temperature": 0.7,
+                "ignore_eos": True,
+            },
+        },
+        {
+            "__idx_in_batch": 1,
+            "prompt": "World",
+            "sampling_params": {
+                "max_tokens": 5,
+                "temperature": 0.7,
+                "ignore_eos": True,
+            },
+        },
+    ]
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+
+    for resp in asyncio.as_completed(tasks):
+        request, output, time_taken_llm = await resp
+        params = request.params
+        max_tokens = params.max_tokens
+        assert max_tokens == output["num_generated_tokens"]
+        assert time_taken_llm > 0
+
+    # Clean up GPU memory
+    wrapper.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_embed(model_opt_125m):
+    wrapper = vLLMEngineWrapper(
+        model=model_opt_125m,
+        model_source=model_opt_125m,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=True,
+        max_pending_requests=10,
+        # Skip CUDA graph capturing to reduce the start time.
+        enforce_eager=True,
+        gpu_memory_utilization=0.8,
+        max_model_len=2048,
+        task_type=vLLMTaskType.EMBED,
+        # Older GPUs (e.g. T4) don't support bfloat16.
+        dtype="half",
+    )
+
+    batch = [
+        {"__idx_in_batch": 0, "prompt": "Hello World"},
+        {"__idx_in_batch": 1, "prompt": "How are you?"},
+    ]
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+
+    for resp in asyncio.as_completed(tasks):
+        _, output, time_taken_llm = await resp
+        assert output["embeddings"].shape == (768,)
+        assert time_taken_llm > 0
+
+    # Clean up GPU memory
+    wrapper.shutdown()
+
+
+@pytest.mark.parametrize(
+    "pooling_params,tokenization_kwargs,expect_same_output",
+    [
+        ({}, None, True),
+        # Truncation via tokenization_kwargs.
+        (None, {"truncation": True, "max_length": 3}, False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_vllm_wrapper_embed_pooling_params(
+    model_opt_125m, pooling_params, tokenization_kwargs, expect_same_output
+):
+    prompt = "Hello! How's the weather?"
+    wrapper = vLLMEngineWrapper(
+        model=model_opt_125m,
+        model_source=model_opt_125m,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=True,
+        max_pending_requests=10,
+        # Skip CUDA graph capturing to reduce the start time.
+        enforce_eager=True,
+        gpu_memory_utilization=0.8,
+        max_model_len=2048,
+        task_type=vLLMTaskType.EMBED,
+    )
+
+    row_with_params = {
+        "__idx_in_batch": 0,
+        "prompt": prompt,
+    }
+    if pooling_params is not None:
+        row_with_params["pooling_params"] = pooling_params
+    if tokenization_kwargs is not None:
+        row_with_params["tokenization_kwargs"] = tokenization_kwargs
+
+    batch = [
+        row_with_params,
+        {
+            "__idx_in_batch": 1,
+            "prompt": prompt,
+            # By default, no pooling params are applied.
+        },
+    ]
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+
+    outputs = {}
+    for resp in asyncio.as_completed(tasks):
+        request, output, time_taken_llm = await resp
+        idx = request.idx_in_batch
+        outputs[idx] = output
+
+        assert output["embeddings"].shape == (768,)
+        assert time_taken_llm > 0
+
+    assert (
+        outputs[0]["embeddings"] == outputs[1]["embeddings"]
+    ).all() == expect_same_output
+
+    # Clean up GPU memory
+    wrapper.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_embed_long_prompt(model_opt_125m):
+    # Preferred path: tokenization_kwargs truncation.
+    pooling_params = None
+    tokenization_kwargs = {"truncation": True, "max_length": 2048}
+    # Sufficiently long prompt to trigger truncation to max_model_len
+    max_model_len = 2048
+    prompt = "Hello! How's the weather?" * 10_000
+    wrapper = vLLMEngineWrapper(
+        model=model_opt_125m,
+        model_source=model_opt_125m,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=True,
+        max_pending_requests=10,
+        # Skip CUDA graph capturing to reduce the start time.
+        enforce_eager=True,
+        gpu_memory_utilization=0.8,
+        max_model_len=max_model_len,
+        task_type=vLLMTaskType.EMBED,
+    )
+
+    row = {"__idx_in_batch": 0, "prompt": prompt}
+    if pooling_params is not None:
+        row["pooling_params"] = pooling_params
+    if tokenization_kwargs is not None:
+        row["tokenization_kwargs"] = tokenization_kwargs
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row))]
+
+    for resp in asyncio.as_completed(tasks):
+        _, output, time_taken_llm = await resp
+        assert output["embeddings"].shape == (768,)
+        assert output["num_input_tokens"] == max_model_len
+        assert time_taken_llm > 0
+
+    # Clean up GPU memory
+    wrapper.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_lora(model_llama_3_2_216M, model_llama_3_2_216M_lora):
+    wrapper = vLLMEngineWrapper(
+        model=model_llama_3_2_216M,
+        model_source=model_llama_3_2_216M,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=True,
+        max_pending_requests=10,
+        # Skip CUDA graph capturing to reduce the start time.
+        enforce_eager=True,
+        task_type=vLLMTaskType.GENERATE,
+        max_model_len=2048,
+        enable_lora=True,
+        max_lora_rank=16,
+    )
+
+    batch = [
+        {
+            "__idx_in_batch": 0,
+            "prompt": "Hello",
+            "sampling_params": {
+                "max_tokens": 10,
+                "temperature": 0.7,
+                "ignore_eos": True,
+            },
+            "model": model_llama_3_2_216M_lora,
+        },
+        {
+            "__idx_in_batch": 1,
+            "prompt": "World",
+            "sampling_params": {
+                "max_tokens": 5,
+                "temperature": 0.7,
+                "ignore_eos": True,
+            },
+        },
+    ]
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+
+    for resp in asyncio.as_completed(tasks):
+        request, output, time_taken_llm = await resp
+        params = request.params
+        max_tokens = params.max_tokens
+        assert max_tokens == output["num_generated_tokens"]
+        assert time_taken_llm > 0
+
+    # Clean up GPU memory
+    wrapper.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_vllm_wrapper_json(model_llama_3_2_1B_instruct):
+    """Test JSON output with the structured_outputs sampling param."""
+
+    class AnswerModel(BaseModel):
+        answer: int
+        explain: str
+
+    json_schema = AnswerModel.model_json_schema()
+
+    wrapper = vLLMEngineWrapper(
+        model=model_llama_3_2_1B_instruct,
+        model_source=model_llama_3_2_1B_instruct,
+        idx_in_batch_column="__idx_in_batch",
+        disable_log_stats=True,
+        max_pending_requests=10,
+        # Skip CUDA graph capturing to reduce the start time.
+        enforce_eager=True,
+        task_type=vLLMTaskType.GENERATE,
+        max_model_len=2048,
+        structured_outputs_config={"backend": "xgrammar"},
+        seed=42,
+    )
+
+    batch = [
+        {
+            "__idx_in_batch": 0,
+            "prompt": "Answer 2 ** 3 + 5. Return the answer in JSON. Expected fields: 'answer', 'explain'.",
+            "sampling_params": {
+                "max_tokens": 100,
+                "temperature": 0.7,
+                "structured_outputs": {"json": json_schema},
+            },
+        },
+    ]
+
+    tasks = [asyncio.create_task(wrapper.generate_async(row)) for row in batch]
+
+    for resp in asyncio.as_completed(tasks):
+        _, output, time_taken_llm = await resp
+        json_obj = json.loads(output["generated_text"])
+        assert "answer" in json_obj
+        assert isinstance(json_obj["answer"], int)
+        assert "explain" in json_obj
+        assert isinstance(json_obj["explain"], str)
+        assert time_taken_llm > 0
+
+    # Clean up GPU memory
+    wrapper.shutdown()
+
+
+def test_vllm_output_data_logprobs():
+    """Test that logprobs and prompt_logprobs are correctly extracted."""
+    from vllm.logprobs import Logprob
+    from vllm.outputs import CompletionOutput, RequestOutput
+
+    logprobs = [
+        {
+            123: Logprob(logprob=-0.5, rank=1, decoded_token="hello"),
+            456: Logprob(logprob=-1.2, rank=2, decoded_token="hi"),
+        },
+        {
+            789: Logprob(logprob=-0.3, rank=1, decoded_token="world"),
+            999: Logprob(logprob=-1.5, rank=2, decoded_token="earth"),
+        },
+    ]
+
+    prompt_logprobs = [
+        None,
+        {
+            111: Logprob(logprob=-0.1, rank=1, decoded_token="test"),
+            222: Logprob(logprob=-0.8, rank=2, decoded_token="demo"),
+        },
+    ]
+
+    request_output = RequestOutput(
+        request_id="test",
+        prompt="test prompt",
+        prompt_token_ids=[1, 2],
+        prompt_logprobs=prompt_logprobs,
+        outputs=[
+            CompletionOutput(
+                index=0,
+                text="hello world",
+                token_ids=[123, 789],
+                cumulative_logprob=-0.8,
+                logprobs=logprobs,
+            )
+        ],
+        finished=True,
+    )
+
+    output_data = vLLMOutputData.from_vllm_engine_output(request_output)
+
+    expected_logprobs = [
+        {
+            123: {"logprob": -0.5, "rank": 1, "decoded_token": "hello"},
+            456: {"logprob": -1.2, "rank": 2, "decoded_token": "hi"},
+        },
+        {
+            789: {"logprob": -0.3, "rank": 1, "decoded_token": "world"},
+            999: {"logprob": -1.5, "rank": 2, "decoded_token": "earth"},
+        },
+    ]
+    assert output_data.logprobs == expected_logprobs
+
+    expected_prompt_logprobs = [
+        None,
+        {
+            111: {"logprob": -0.1, "rank": 1, "decoded_token": "test"},
+            222: {"logprob": -0.8, "rank": 2, "decoded_token": "demo"},
+        },
+    ]
+    assert output_data.prompt_logprobs == expected_prompt_logprobs
+
+    dumped = output_data.model_dump()
+    assert dumped["logprobs"] == expected_logprobs
+    assert dumped["prompt_logprobs"] == expected_prompt_logprobs
+
+
+def test_vllm_output_data_no_logprobs():
+    """Test that None logprobs are handled correctly when not requested."""
+    from vllm.outputs import CompletionOutput, RequestOutput
+
+    request_output = RequestOutput(
+        request_id="test",
+        prompt="test prompt",
+        prompt_token_ids=[1, 2],
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=0,
+                text="test response",
+                token_ids=[4, 5, 6],
+                cumulative_logprob=None,
+                logprobs=None,
+            )
+        ],
+        finished=True,
+    )
+
+    output_data = vLLMOutputData.from_vllm_engine_output(request_output)
+
+    assert output_data.logprobs is None
+    assert output_data.prompt_logprobs is None
+
+    dumped = output_data.model_dump()
+    assert dumped["logprobs"] is None
+    assert dumped["prompt_logprobs"] is None
+
+
+@pytest.mark.asyncio
+async def test_vllm_udf_default_raises_on_error(mock_vllm_wrapper):
+    """Default behavior (should_continue_on_error=False) raises on inference error."""
+    mock_vllm_wrapper.return_value.generate_async.side_effect = ValueError(
+        "prompt too long"
+    )
+
+    udf = vLLMEngineStageUDF(
+        data_column="__data",
+        expected_input_keys=["prompt", "sampling_params"],
+        model="/tmp/fake-model",
+        task_type=vLLMTaskType.GENERATE,
+        batch_size=32,
+        max_concurrent_batches=4,
+        engine_kwargs={},
+        should_continue_on_error=False,
+    )
+
+    batch = {"__data": [{"prompt": "test", "sampling_params": {"temperature": 0.7}}]}
+
+    with pytest.raises(ValueError, match="prompt too long"):
+        async for _ in udf(batch):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_vllm_udf_should_continue_on_error_yields_error_row(mock_vllm_wrapper):
+    """With should_continue_on_error=True, errors yield rows with __inference_error__."""
+    mock_vllm_wrapper.return_value.generate_async.side_effect = ValueError(
+        "prompt too long"
+    )
+
+    udf = vLLMEngineStageUDF(
+        data_column="__data",
+        expected_input_keys=["prompt", "sampling_params"],
+        model="/tmp/fake-model",
+        task_type=vLLMTaskType.GENERATE,
+        batch_size=32,
+        max_concurrent_batches=4,
+        engine_kwargs={},
+        should_continue_on_error=True,
+    )
+
+    batch = {
+        "__data": [{"prompt": "test prompt", "sampling_params": {"temperature": 0.7}}]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 1
+    assert "__inference_error__" in results[0]
+    assert "ValueError" in results[0]["__inference_error__"]
+    assert "prompt too long" in results[0]["__inference_error__"]
+    # Error rows include the original prompt for debuggability
+    assert results[0]["prompt"] == "test prompt"
+
+
+@pytest.mark.asyncio
+async def test_vllm_udf_mixed_success_and_error(mock_vllm_wrapper):
+    """Mixed batch: some rows succeed, some fail."""
+    call_count = 0
+
+    async def mock_generate(row):
+        nonlocal call_count
+        call_count += 1
+        idx = row["__idx_in_batch"]
+        if idx == 1:
+            raise ValueError("prompt too long")
+        output_data = vLLMOutputData(
+            prompt=row["prompt"],
+            prompt_token_ids=None,
+            num_input_tokens=0,
+        )
+        return (
+            MagicMock(
+                request_id=idx,
+                prompt=row["prompt"],
+                params=row["sampling_params"],
+                idx_in_batch=idx,
+            ),
+            output_data.model_dump(),
+            0.1,
+        )
+
+    mock_vllm_wrapper.return_value.generate_async.side_effect = mock_generate
+
+    udf = vLLMEngineStageUDF(
+        data_column="__data",
+        expected_input_keys=["prompt", "sampling_params"],
+        model="/tmp/fake-model",
+        task_type=vLLMTaskType.GENERATE,
+        batch_size=32,
+        max_concurrent_batches=4,
+        engine_kwargs={},
+        should_continue_on_error=True,
+    )
+
+    batch = {
+        "__data": [
+            {"prompt": "first", "sampling_params": {"temperature": 0.7}},
+            {"prompt": "second", "sampling_params": {"temperature": 0.7}},
+            {"prompt": "third", "sampling_params": {"temperature": 0.7}},
+        ]
+    }
+
+    results = []
+    async for result in udf(batch):
+        results.extend(result["__data"])
+
+    assert len(results) == 3
+
+    errors = [r for r in results if r.get("__inference_error__", "") != ""]
+    successes = [r for r in results if r.get("__inference_error__", "") == ""]
+
+    assert len(errors) == 1
+    assert len(successes) == 2
+    assert "ValueError" in errors[0]["__inference_error__"]
+
+    # Verify schema consistency
+    error_keys = set(errors[0].keys())
+    success_keys = set(successes[0].keys())
+    assert error_keys == success_keys
+
+
+@pytest.mark.asyncio
+async def test_vllm_udf_fatal_error_exits_actor(mock_vllm_wrapper):
+    """Fatal errors (EngineDeadError) trigger actor exit for recovery, not error rows."""
+    from vllm.v1.engine.exceptions import EngineDeadError
+
+    mock_vllm_wrapper.return_value.generate_async.side_effect = EngineDeadError()
+
+    udf = vLLMEngineStageUDF(
+        data_column="__data",
+        expected_input_keys=["prompt", "sampling_params"],
+        model="/tmp/fake-model",
+        task_type=vLLMTaskType.GENERATE,
+        batch_size=32,
+        max_concurrent_batches=4,
+        engine_kwargs={},
+        should_continue_on_error=True,  # Even with this True, fatal errors should not yield error rows
+    )
+
+    batch = {"__data": [{"prompt": "test", "sampling_params": {"temperature": 0.7}}]}
+
+    # Fatal errors trigger actor exit for recovery (not error rows, not simple re-raise).
+    # We use os._exit(1) instead of ray.actor.exit_actor() because:
+    # - os._exit(1) -> SYSTEM_ERROR -> RaySystemError -> task IS retried
+    # - ray.actor.exit_actor() -> INTENDED_USER_EXIT -> ActorDiedError -> NOT retried
+    # We mock os._exit to verify it was called with exit code 1.
+    with patch(
+        "ray.llm._internal.batch.stages.vllm_engine_stage.os._exit"
+    ) as mock_os_exit:
+        # Don't actually exit - let code continue and fail naturally
+        # The important thing is verifying os._exit was called
+        try:
+            async for _ in udf(batch):
+                pass
+        except Exception:
+            pass  # Code may fail after mock returns None - that's OK for this test
+
+        mock_os_exit.assert_called_once_with(1)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/llm/tests/batch/observability/usage_telemetry/test_usage.py

@@ -0,0 +1,125 @@
+import sys
+
+import pytest
+
+import ray
+from ray._common.usage.usage_lib import TagKey
+from ray.llm._internal.batch.observability.usage_telemetry.usage import (
+    BatchModelTelemetry,
+    get_or_create_telemetry_agent,
+)
+from ray.llm._internal.batch.processor import ProcessorBuilder
+from ray.llm._internal.batch.processor.http_request_proc import (
+    HttpRequestProcessorConfig,
+)
+from ray.llm._internal.batch.processor.vllm_engine_proc import (
+    vLLMEngineProcessorConfig,
+)
+
+
+@ray.remote(num_cpus=0)
+class FakeTelemetryRecorder:
+    def __init__(self):
+        self._telemetry = {}
+
+    def record(self, key, value):
+        self._telemetry[key] = value
+
+    def telemetry(self):
+        return self._telemetry
+
+
+def test_push_telemetry_report():
+    recorder = FakeTelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        recorder.record.remote(key, value)
+
+    telemetry_agent = get_or_create_telemetry_agent()
+    ray.get(telemetry_agent.remote_telemetry_agent._reset.remote())
+    telemetry_agent._update_record_tag_func(record_tag_func)
+
+    config = vLLMEngineProcessorConfig(
+        model_source="facebook/opt-125m",
+        engine_kwargs=dict(
+            max_model_len=8192,
+        ),
+        runtime_env=dict(
+            env_vars=dict(
+                RANDOM_ENV_VAR="12345",
+            ),
+        ),
+        accelerator_type="A10G",
+        concurrency=4,
+        batch_size=64,
+        max_pending_requests=111,
+        chat_template_stage=True,
+        tokenize_stage=True,
+        detokenize_stage=True,
+        prepare_multimodal_stage=True,
+    )
+    _ = ProcessorBuilder.build(config)
+
+    _ = ProcessorBuilder.build(
+        HttpRequestProcessorConfig(
+            url="http://localhost:8000",
+            headers={"Authorization": "Bearer 1234567890"},
+            qps=2,
+            concurrency=4,
+            batch_size=64,
+        )
+    )
+
+    # Ensure that the telemetry is correct after pushing the reports.
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert telemetry == {
+        TagKey.LLM_BATCH_PROCESSOR_CONFIG_NAME: "vLLMEngineProcessorConfig,HttpRequestProcessorConfig",
+        TagKey.LLM_BATCH_MODEL_ARCHITECTURE: "OPTForCausalLM,",
+        TagKey.LLM_BATCH_SIZE: "64,64",
+        TagKey.LLM_BATCH_ACCELERATOR_TYPE: "A10G,",
+        TagKey.LLM_BATCH_CONCURRENCY: "4,4",
+        TagKey.LLM_BATCH_TASK_TYPE: "generate,",
+        TagKey.LLM_BATCH_PIPELINE_PARALLEL_SIZE: "1,0",
+        TagKey.LLM_BATCH_TENSOR_PARALLEL_SIZE: "1,0",
+        TagKey.LLM_BATCH_DATA_PARALLEL_SIZE: "1,0",
+    }, f"actual telemetry: {telemetry}"
+
+
+def test_telemetry_dedups_by_model_identity():
+    """Distinct models sharing reported fields stay separate; identical builds merge."""
+    recorder = FakeTelemetryRecorder.remote()
+
+    def record_tag_func(key, value):
+        ray.get(recorder.record.remote(key, value))
+
+    telemetry_agent = get_or_create_telemetry_agent()
+    ray.get(telemetry_agent.remote_telemetry_agent._reset.remote())
+    telemetry_agent._update_record_tag_func(record_tag_func)
+
+    # Two distinct models with identical reported fields (same architecture/config).
+    common = dict(
+        model_architecture="LlamaForCausalLM",
+        batch_size=64,
+        concurrency=4,
+        task_type="generate",
+    )
+    telemetry_agent.push_telemetry_report(
+        BatchModelTelemetry(model_id_hash="hash_a", **common)
+    )
+    telemetry_agent.push_telemetry_report(
+        BatchModelTelemetry(model_id_hash="hash_b", **common)
+    )
+    # A repeated identical build of model A must not add a third entry.
+    telemetry_agent.push_telemetry_report(
+        BatchModelTelemetry(model_id_hash="hash_a", **common)
+    )
+
+    telemetry = ray.get(recorder.telemetry.remote())
+    assert (
+        telemetry[TagKey.LLM_BATCH_MODEL_ARCHITECTURE]
+        == "LlamaForCausalLM,LlamaForCausalLM"
+    ), f"actual telemetry: {telemetry}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))