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__]))