chore: import upstream snapshot with attribution

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