chore: import upstream snapshot with attribution

doc/source/data/doc_code/custom_datasource_example.py

@@ -0,0 +1,79 @@
+# flake8: noqa
+# fmt: off
+
+from typing import Iterator, Union, List
+
+import pyarrow
+
+from ray.data.block import Block
+
+# __datasource_constructor_start__
+from ray.data.datasource import FileBasedDatasource
+
+class ImageDatasource(FileBasedDatasource):
+    def __init__(self, paths: Union[str, List[str]], *, mode: str):
+        super().__init__(
+            paths,
+            file_extensions=["png", "jpg", "jpeg", "bmp", "gif", "tiff"],
+        )
+
+        self.mode = mode  # Specify read options in the constructor
+# __datasource_constructor_end__
+
+# __read_stream_start__
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import io
+        import numpy as np
+        from PIL import Image
+        from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+
+        data = f.readall()
+        image = Image.open(io.BytesIO(data))
+        image = image.convert(self.mode)
+
+        # Each block contains one row
+        builder = DelegatingBlockBuilder()
+        array = np.asarray(image)
+        item = {"image": array}
+        builder.add(item)
+        yield builder.build()
+# __read_stream_end__
+
+# __read_datasource_start__
+import ray
+
+ds = ray.data.read_datasource(
+    ImageDatasource("s3://anonymous@ray-example-data/batoidea", mode="RGB")
+)
+# __read_datasource_end__
+
+
+from typing import Any, Dict
+import pyarrow
+
+# __datasink_constructor_start__
+from ray.data.datasource import RowBasedFileDatasink
+
+class ImageDatasink(RowBasedFileDatasink):
+    def __init__(self, path: str, column: str, file_format: str):
+        super().__init__(path, file_format=file_format)
+
+        self.column = column
+        self.file_format = file_format  # Specify write options in the constructor
+# __datasink_constructor_end__
+
+# __write_row_to_file_start__
+    def write_row_to_file(self, row: Dict[str, Any], file: pyarrow.NativeFile):
+        import io
+        from PIL import Image
+
+        # PIL can't write to a NativeFile, so we have to write to a buffer first.
+        image = Image.fromarray(row[self.column])
+        buffer = io.BytesIO()
+        image.save(buffer, format=self.file_format)
+        file.write(buffer.getvalue())
+# __write_row_to_file_end__
+
+# __write_datasink_start__
+ds.write_datasink(ImageDatasink("/tmp/results", column="image", file_format="png"))
+# __write_datasink_end__

doc/source/data/doc_code/key_concepts.py

@@ -0,0 +1,27 @@
+# flake8: noqa
+
+# fmt: off
+# __resource_allocation_1_begin__
+import ray
+from ray import tune
+
+# This workload will use spare cluster resources for execution.
+def objective(*args):
+    ray.data.range(10).show()
+
+# Create a cluster with 4 CPU slots available.
+ray.init(num_cpus=4)
+
+# By setting `max_concurrent_trials=3`, this ensures the cluster will always
+# have a sparse CPU for Dataset. Try setting `max_concurrent_trials=4` here,
+# and notice that the experiment will appear to hang.
+tuner = tune.Tuner(
+    tune.with_resources(objective, {"cpu": 1}),
+    tune_config=tune.TuneConfig(
+        num_samples=1,
+        max_concurrent_trials=3
+    )
+)
+tuner.fit()
+# __resource_allocation_1_end__
+# fmt: on

doc/source/data/doc_code/working-with-llms/basic_llm_example.py

@@ -0,0 +1,462 @@
+"""
+This file serves as a documentation example and CI test for basic LLM batch inference.
+
+"""
+
+# __basic_llm_example_start__
+import os
+import shutil
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# __basic_config_example_start__
+# Basic vLLM configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,  # Reduce if CUDA OOM occurs
+        "max_model_len": 4096,  # Constrain to fit test GPU memory
+    },
+    concurrency=1,
+    batch_size=64,
+)
+# __basic_config_example_end__
+
+processor = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        messages=[
+            {"role": "system", "content": "You are a bot that responds with haikus."},
+            {"role": "user", "content": row["item"]},
+        ],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=250,
+        ),
+    ),
+    postprocess=lambda row: dict(
+        answer=row["generated_text"],
+        **row,  # This will return all the original columns in the dataset.
+    ),
+)
+
+ds = ray.data.from_items(["Start of the haiku is: Complete this for me..."])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            ds = processor(ds)
+            ds.show(limit=1)
+        else:
+            print("Skipping basic LLM run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping basic LLM run due to environment error: {e}")
+
+# __hf_token_config_example_start__
+# Configuration with Hugging Face token
+config_with_token = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    runtime_env={"env_vars": {"HF_TOKEN": "your_huggingface_token"}},
+    concurrency=1,
+    batch_size=64,
+)
+# __hf_token_config_example_end__
+
+# __parallel_config_example_start__
+# Model parallelism configuration for larger models
+# tensor_parallel_size=2: Split model across 2 GPUs for tensor parallelism
+# pipeline_parallel_size=2: Use 2 pipeline stages (total 4 GPUs needed)
+# Total GPUs required = tensor_parallel_size * pipeline_parallel_size = 4
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_model_len": 16384,
+        "tensor_parallel_size": 2,
+        "pipeline_parallel_size": 2,
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 2048,
+    },
+    concurrency=1,
+    batch_size=32,
+    accelerator_type="L4",
+)
+# __parallel_config_example_end__
+
+# __runai_config_example_start__
+# RunAI streamer configuration for optimized model loading
+# Note: Install vLLM with runai dependencies: pip install -U "vllm[runai]>=0.10.1"
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "load_format": "runai_streamer",
+        "max_model_len": 16384,
+    },
+    concurrency=1,
+    batch_size=64,
+)
+# __runai_config_example_end__
+
+# __lora_config_example_start__
+# Multi-LoRA configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_lora": True,
+        "max_lora_rank": 32,
+        "max_loras": 1,
+        "max_model_len": 16384,
+    },
+    concurrency=1,
+    batch_size=32,
+)
+# __lora_config_example_end__
+
+# __s3_config_example_start__
+# S3 hosted model configuration
+s3_config = vLLMEngineProcessorConfig(
+    model_source="s3://your-bucket/your-model-path/",
+    engine_kwargs={
+        "load_format": "runai_streamer",
+        "max_model_len": 16384,
+    },
+    concurrency=1,
+    batch_size=64,
+)
+# __s3_config_example_end__
+
+base_dir = "/tmp/llm_checkpoint_demo"
+input_path = os.path.join(base_dir, "input")
+output_path = os.path.join(base_dir, "output")
+checkpoint_path = os.path.join(base_dir, "checkpoint")
+
+# Reset directories
+for path in (input_path, output_path, checkpoint_path):
+    shutil.rmtree(path, ignore_errors=True)
+    os.makedirs(path)
+
+# __row_level_fault_tolerance_config_example_start__
+# Row-level fault tolerance configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    concurrency=1,
+    batch_size=64,
+    should_continue_on_error=True,
+)
+# __row_level_fault_tolerance_config_example_end__
+
+# Seed the input directory because Ray Data V2's `read_parquet` errors on empty dirs
+ray.data.from_items(
+    [{"id": i, "message": f"Question {i}: What is 2 + 2?"} for i in range(4)]
+).write_parquet(input_path)
+
+# __checkpoint_config_setup_example_start__
+from ray.data.checkpoint import CheckpointConfig
+
+ctx = ray.data.DataContext.get_current()
+ctx.checkpoint_config = CheckpointConfig(
+    id_column="id",
+    checkpoint_path=checkpoint_path,
+    delete_checkpoint_on_success=False,
+)
+# __checkpoint_config_setup_example_end__
+
+# __checkpoint_usage_example_start__
+processor_config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 4096,
+    },
+    concurrency=1,
+    batch_size=16,
+)
+
+processor = build_processor(
+    processor_config,
+    preprocess=lambda row: dict(
+        id=row["id"], # Preserve the ID column for checkpointing
+        messages=[{"role": "user", "content": row["message"]}],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=10,
+        ),
+    ),
+    postprocess=lambda row: {
+        "id": row["id"], # Preserve the ID column for checkpointing
+        "answer": row.get("generated_text"),
+    },
+)
+
+ds = ray.data.read_parquet(input_path)
+ds = processor(ds)
+ds.write_parquet(output_path)
+# __checkpoint_usage_example_end__
+
+
+# __gpu_memory_config_example_start__
+# GPU memory management configuration
+# If you encounter CUDA out of memory errors, try these optimizations:
+config_memory_optimized = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_model_len": 8192,
+        "max_num_batched_tokens": 2048,
+        "enable_chunked_prefill": True,
+        "gpu_memory_utilization": 0.85,
+        "block_size": 16,
+    },
+    concurrency=1,
+    batch_size=16,
+)
+
+# For very large models or limited GPU memory:
+config_minimal_memory = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_model_len": 4096,
+        "max_num_batched_tokens": 1024,
+        "enable_chunked_prefill": True,
+        "gpu_memory_utilization": 0.75,
+    },
+    concurrency=1,
+    batch_size=8,
+)
+# __gpu_memory_config_example_end__
+
+# __embedding_config_example_start__
+# Embedding model configuration
+embedding_config = vLLMEngineProcessorConfig(
+    model_source="sentence-transformers/all-MiniLM-L6-v2",
+    task_type="embed",
+    engine_kwargs=dict(
+        enable_prefix_caching=False,
+        enable_chunked_prefill=False,
+        max_model_len=256,
+        enforce_eager=True,
+    ),
+    batch_size=32,
+    concurrency=1,
+    chat_template_stage=False,  # Skip chat templating for embeddings
+    detokenize_stage=False,     # Skip detokenization for embeddings
+)
+
+# Example usage for embeddings
+def create_embedding_processor():
+    return build_processor(
+        embedding_config,
+        preprocess=lambda row: dict(prompt=row["text"]),
+        postprocess=lambda row: {
+            "text": row["prompt"],
+            "embedding": row["embeddings"],
+        },
+    )
+
+
+# __embedding_config_example_end__
+
+# __classification_config_example_start__
+# Sequence classification model configuration
+# Use task_type="classify" for classification models (e.g., sentiment, quality scoring)
+# Use task_type="score" for cross-encoder scoring models
+classification_config = vLLMEngineProcessorConfig(
+    model_source="nvidia/nemocurator-fineweb-nemotron-4-edu-classifier",
+    task_type="classify",
+    engine_kwargs=dict(
+        max_model_len=512,
+        enforce_eager=True,
+    ),
+    batch_size=8,
+    concurrency=1,
+    chat_template_stage=False,
+    detokenize_stage=False,
+)
+
+
+# Example usage for classification
+def create_classification_processor():
+    return build_processor(
+        classification_config,
+        preprocess=lambda row: dict(prompt=row["text"]),
+        postprocess=lambda row: {
+            "text": row["prompt"],
+            # Classification models return logits in the 'embeddings' field
+            "score": float(row["embeddings"][0])
+            if row.get("embeddings") is not None and len(row["embeddings"]) > 0
+            else None,
+        },
+    )
+
+
+# __classification_config_example_end__
+
+# __shared_vllm_engine_config_example_start__
+import ray
+from ray import serve
+from ray.data.llm import ServeDeploymentProcessorConfig, build_processor
+from ray.serve.llm import (
+    LLMConfig,
+    ModelLoadingConfig,
+    build_llm_deployment,
+)
+from ray.serve.llm.openai_api_models import CompletionRequest
+
+llm_config = LLMConfig(
+    model_loading_config=ModelLoadingConfig(
+        model_id="facebook/opt-1.3b",
+        model_source="facebook/opt-1.3b",
+    ),
+    deployment_config=dict(
+        name="demo_deployment_config",
+        autoscaling_config=dict(
+            min_replicas=1,
+            max_replicas=1,
+        ),
+    ),
+    engine_kwargs=dict(
+        enable_prefix_caching=True,
+        enable_chunked_prefill=True,
+        max_num_batched_tokens=4096,
+    ),
+)
+
+APP_NAME = "demo_app"
+DEPLOYMENT_NAME = "demo_deployment"
+override_serve_options = dict(name=DEPLOYMENT_NAME)
+
+llm_app = build_llm_deployment(
+    llm_config, override_serve_options=override_serve_options
+)
+app = serve.run(llm_app, name=APP_NAME)
+config = ServeDeploymentProcessorConfig(
+    deployment_name=DEPLOYMENT_NAME,
+    app_name=APP_NAME,
+    dtype_mapping={
+        "CompletionRequest": CompletionRequest,
+    },
+    concurrency=1,
+    batch_size=64,
+)
+
+processor1 = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        method="completions",
+        dtype="CompletionRequest",
+        request_kwargs=dict(
+            model="facebook/opt-1.3b",
+            prompt=f"This is a prompt for {row['id']}",
+            stream=False,
+        ),
+    ),
+    postprocess=lambda row: dict(
+        prompt=row["choices"][0]["text"],
+    ),
+)
+
+processor2 = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        method="completions",
+        dtype="CompletionRequest",
+        request_kwargs=dict(
+            model="facebook/opt-1.3b",
+            prompt=row["prompt"],
+            stream=False,
+        ),
+    ),
+    postprocess=lambda row: row,
+)
+
+ds = ray.data.range(10)
+ds = processor2(processor1(ds))
+print(ds.take_all())
+# __shared_vllm_engine_config_example_end__
+
+# __cross_node_parallelism_config_example_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+        "pipeline_parallel_size": 4,
+        "tensor_parallel_size": 4,
+        "distributed_executor_backend": "ray",
+    },
+    batch_size=32,
+    concurrency=1,
+)
+# __cross_node_parallelism_config_example_end__
+
+# __custom_placement_group_strategy_config_example_start__
+# Simple: specify resources per worker, auto-replicated by TP*PP (4 workers here)
+# Alternative: use "bundles": [{"GPU": 1}] * 4 for explicit bundle control
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+        "pipeline_parallel_size": 2,
+        "tensor_parallel_size": 2,
+        "distributed_executor_backend": "ray",
+    },
+    batch_size=32,
+    concurrency=1,
+    placement_group_config={
+        "bundle_per_worker": {"GPU": 1},
+        "strategy": "STRICT_PACK",
+    },
+)
+# __custom_placement_group_strategy_config_example_end__
+
+# __concurrent_config_example_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+    },
+    concurrency=10,
+    batch_size=64,
+)
+# __concurrent_config_example_end__
+
+
+# __concurrent_config_fixed_pool_example_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+    },
+    concurrency=(10, 10),
+    batch_size=64,
+)
+# __concurrent_config_fixed_pool_example_end__
+
+# __concurrent_batches_tuning_example_start__
+# Tuning concurrent batch processing
+# Configure both parameters together for optimal throughput
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+    },
+    batch_size=64,
+    # Dataset-level concurrency (number of actor replicas)
+    concurrency=1,
+    # Number of batches that can run concurrently per actor (default: 8)
+    max_concurrent_batches=8,
+    # Number of tasks Ray Data queues per actor (default: 16)
+    # Increase to keep actor task queue saturated
+    experimental={"max_tasks_in_flight_per_actor": 16},
+)
+# __concurrent_batches_tuning_example_end__
+# __basic_llm_example_end__

doc/source/data/doc_code/working-with-llms/classification_example.py

@@ -0,0 +1,69 @@
+"""
+Classification batch inference with Ray Data LLM.
+
+Uses sequence classification models for content classifiers and sentiment analyzers.
+"""
+
+# Dependency setup
+import subprocess
+import sys
+
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "ray[llm]"])
+subprocess.check_call(
+    [sys.executable, "-m", "pip", "install", "--upgrade", "transformers"]
+)
+subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy==1.26.4"])
+
+
+# __classification_example_start__
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# Configure vLLM for a sequence classification model
+classification_config = vLLMEngineProcessorConfig(
+    model_source="nvidia/nemocurator-fineweb-nemotron-4-edu-classifier",
+    task_type="classify",  # Use 'classify' for sequence classification models
+    engine_kwargs=dict(
+        max_model_len=512,
+        enforce_eager=True,
+    ),
+    batch_size=8,
+    concurrency=1,
+    chat_template_stage=False,
+    detokenize_stage=False,
+)
+
+classification_processor = build_processor(
+    classification_config,
+    preprocess=lambda row: dict(prompt=row["text"]),
+    postprocess=lambda row: {
+        "text": row["prompt"],
+        # Classification models return logits in the 'embeddings' field
+        "edu_score": float(row["embeddings"][0])
+        if row.get("embeddings") is not None and len(row["embeddings"]) > 0
+        else None,
+    },
+)
+
+# Sample texts with varying educational quality
+texts = [
+    "lol that was so funny haha",
+    "Photosynthesis converts light energy into chemical energy.",
+    "Newton's laws describe the relationship between forces and motion.",
+]
+ds = ray.data.from_items([{"text": text} for text in texts])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            classified_ds = classification_processor(ds)
+            classified_ds.show(limit=3)
+        else:
+            print("Skipping classification run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping classification run due to environment error: {e}")
+# __classification_example_end__
+
+

doc/source/data/doc_code/working-with-llms/custom_tokenizer_example.py

@@ -0,0 +1,195 @@
+"""
+Documentation example and test for custom tokenizer batch inference.
+
+Demonstrates how to use vLLM's tokenizer infrastructure for models whose
+tokenizers are not natively supported by HuggingFace (e.g. Mistral Tekken,
+DeepSeek-V3.2, Grok-2 tiktoken).
+
+This example uses a standard model to demonstrate the pattern. For models
+that truly require vLLM's custom tokenizer (e.g. deepseek-ai/DeepSeek-V3-0324),
+replace the model ID and adjust tokenizer_mode accordingly.
+"""
+
+# __custom_chat_template_start__
+from typing import Any, Dict, List
+from vllm.tokenizers import get_tokenizer
+
+
+class VLLMChatTemplate:
+    """Apply a chat template using vLLM's tokenizer."""
+
+    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
+        self.tokenizer = get_tokenizer(
+            model_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=True,
+        )
+
+    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        prompts: List[str] = []
+        all_messages: List[List[Dict[str, Any]]] = []
+
+        for messages in batch["messages"]:
+            if hasattr(messages, "tolist"):
+                messages = messages.tolist()
+            all_messages.append(messages)
+
+            add_generation_prompt = messages[-1]["role"] == "user"
+            prompt = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=add_generation_prompt,
+                continue_final_message=not add_generation_prompt,
+            )
+            prompts.append(prompt)
+
+        return {
+            "prompt": prompts,
+            "messages": all_messages,
+            "sampling_params": batch["sampling_params"],
+        }
+
+
+# __custom_chat_template_end__
+
+
+# __custom_tokenize_start__
+class VLLMTokenize:
+    """Tokenize text prompts using vLLM's tokenizer."""
+
+    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
+        self.tokenizer = get_tokenizer(
+            model_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=True,
+        )
+
+    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        all_tokenized: List[List[int]] = [
+            self.tokenizer.encode(prompt) for prompt in batch["prompt"]
+        ]
+
+        return {
+            "tokenized_prompt": all_tokenized,
+            "messages": batch["messages"],
+            "sampling_params": batch["sampling_params"],
+        }
+
+
+# __custom_tokenize_end__
+
+
+# __custom_detokenize_start__
+class VLLMDetokenize:
+    """Detokenize generated token IDs using vLLM's tokenizer."""
+
+    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
+        self.tokenizer = get_tokenizer(
+            model_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=True,
+        )
+
+    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        decoded: List[str] = []
+        for tokens in batch["generated_tokens"]:
+            if hasattr(tokens, "tolist"):
+                tokens = tokens.tolist()
+            decoded.append(self.tokenizer.decode(tokens, skip_special_tokens=True))
+
+        return {
+            **batch,
+            "generated_text_custom": decoded,
+        }
+
+
+# __custom_detokenize_end__
+
+
+def run_custom_tokenizer_example():
+    import ray
+    from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+    # Input dataset with sampling_params per row.
+    ds = ray.data.from_items(
+        [
+            {
+                "messages": [
+                    {"role": "user", "content": "What is the capital of France?"}
+                ],
+                "sampling_params": {"max_tokens": 256, "temperature": 0.7},
+            },
+            {
+                "messages": [
+                    {"role": "user", "content": "Write a haiku about computing."}
+                ],
+                "sampling_params": {"max_tokens": 256, "temperature": 0.7},
+            },
+        ]
+    )
+
+    # __custom_tokenizer_pipeline_start__
+    MODEL_ID = "unsloth/Llama-3.1-8B-Instruct"
+
+    config = vLLMEngineProcessorConfig(
+        model_source=MODEL_ID,
+        engine_kwargs=dict(
+            max_model_len=4096,
+            trust_remote_code=True,
+            tokenizer_mode="auto",
+        ),
+        batch_size=4,
+        concurrency=1,
+        # Disable built-in stages -- we handle them via map_batches.
+        chat_template_stage=False,
+        tokenize_stage=False,
+        detokenize_stage=False,
+    )
+
+    processor = build_processor(
+        config,
+        postprocess=lambda row: {
+            "generated_text": row.get("generated_text", ""),
+            "generated_tokens": row.get("generated_tokens", []),
+            "num_input_tokens": row.get("num_input_tokens", 0),
+            "num_generated_tokens": row.get("num_generated_tokens", 0),
+        },
+    )
+
+    ds = ds.map_batches(
+        VLLMChatTemplate,
+        fn_constructor_kwargs={"model_id": MODEL_ID},
+        concurrency=1,
+        batch_size=4,
+    )
+
+    ds = ds.map_batches(
+        VLLMTokenize,
+        fn_constructor_kwargs={"model_id": MODEL_ID},
+        concurrency=1,
+        batch_size=4,
+    )
+
+    ds = processor(ds)
+
+    ds = ds.map_batches(
+        VLLMDetokenize,
+        fn_constructor_kwargs={"model_id": MODEL_ID},
+        concurrency=1,
+        batch_size=4,
+    )
+
+    # __custom_tokenizer_pipeline_end__
+    ds.show(limit=2)
+
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_custom_tokenizer_example()
+        else:
+            print("Skipping custom tokenizer example (no GPU available)")
+    except Exception as e:
+        print(f"Skipping custom tokenizer example: {e}")

doc/source/data/doc_code/working-with-llms/embedding_example.py

@@ -0,0 +1,63 @@
+"""
+Documentation example and test for embedding model batch inference.
+
+"""
+
+import subprocess
+import sys
+
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "ray[llm]"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy==1.26.4"])
+
+
+def run_embedding_example():
+    # __embedding_example_start__
+    import ray
+    from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+    embedding_config = vLLMEngineProcessorConfig(
+        model_source="sentence-transformers/all-MiniLM-L6-v2",
+        task_type="embed",
+        engine_kwargs=dict(
+            enable_prefix_caching=False,
+            enable_chunked_prefill=False,
+            max_model_len=256,
+            enforce_eager=True,
+        ),
+        batch_size=32,
+        concurrency=1,
+        chat_template_stage=False,  # Skip chat templating for embeddings
+        detokenize_stage=False,     # Skip detokenization for embeddings
+    )
+
+    embedding_processor = build_processor(
+        embedding_config,
+        preprocess=lambda row: dict(prompt=row["text"]),
+        postprocess=lambda row: {
+            "text": row["prompt"],
+            "embedding": row["embeddings"],
+        },
+    )
+
+    texts = [
+        "Hello world",
+        "This is a test sentence",
+        "Embedding models convert text to vectors",
+    ]
+    ds = ray.data.from_items([{"text": text} for text in texts])
+
+    embedded_ds = embedding_processor(ds)
+    embedded_ds.show(limit=1)
+    # __embedding_example_end__
+
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_embedding_example()
+        else:
+            print("Skipping embedding example (no GPU available)")
+    except Exception as e:
+        print(f"Skipping embedding example: {e}")

doc/source/data/doc_code/working-with-llms/minimal_quickstart.py

@@ -0,0 +1,62 @@
+"""
+Quickstart: vLLM + Ray Data batch inference.
+
+1. Installation
+2. Dataset creation
+3. Processor configuration
+4. Running inference
+5. Getting results
+"""
+
+# __minimal_vllm_quickstart_start__
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# Initialize Ray
+ray.init()
+
+# simple dataset
+ds = ray.data.from_items([
+    {"prompt": "What is machine learning?"},
+    {"prompt": "Explain neural networks in one sentence."},
+])
+
+# Minimal vLLM configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    concurrency=1,  # 1 vLLM engine replica
+    batch_size=32,  # 32 samples per batch
+    engine_kwargs={
+        "max_model_len": 4096, # Fit into test GPU memory
+    }
+)
+
+# Build processor
+# preprocess: converts input row to format expected by vLLM (OpenAI chat format)
+# postprocess: extracts generated text from vLLM output
+processor = build_processor(
+    config,
+    preprocess=lambda row: {
+        "messages": [{"role": "user", "content": row["prompt"]}],
+        "sampling_params": {"temperature": 0.7, "max_tokens": 100},
+    },
+    postprocess=lambda row: {
+        "prompt": row["prompt"],
+        "response": row["generated_text"],
+    },
+)
+
+# inference
+ds = processor(ds)
+
+# iterate through the results
+for result in ds.iter_rows():
+    print(f"Q: {result['prompt']}")
+    print(f"A: {result['response']}\n")
+
+# Alternative ways to get results:
+# results = ds.take(10)  # Get first 10 results
+# ds.show(limit=5)       # Print first 5 results
+# ds.write_parquet("output.parquet")  # Save to file
+# __minimal_vllm_quickstart_end__
+

doc/source/data/doc_code/working-with-llms/omni_audio_example.py

@@ -0,0 +1,212 @@
+"""
+This file serves as a documentation example and CI test for VLM batch inference with audio.
+
+Structure:
+1. Infrastructure setup: Dataset compatibility patches, dependency handling
+2. Docs example (between __vlm_audio_example_start/end__): Embedded in Sphinx docs via literalinclude
+3. Test validation and cleanup
+"""
+
+
+'''
+# __audio_message_format_example_start__
+"""Supported audio input formats: audio URL, audio binary data"""
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "Provide a detailed description of the audio."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe what happens in this audio."},
+                # Option 1: Provide audio URL
+                {"type": "audio_url", "audio_url": {"url": "https://example.com/audio.wav"}},
+                # Option 2: Provide audio binary data
+                {"type": "input_audio", "input_audio": {"data": audio_base64, "format": "wav"}},
+            ]
+        },
+    ]
+}
+# __audio_message_format_example_end__
+'''
+
+# __omni_audio_example_start__
+import ray
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+)
+
+# __omni_audio_config_example_start__
+audio_processor_config = vLLMEngineProcessorConfig(
+    model_source="Qwen/Qwen2.5-Omni-3B",
+    task_type="generate",
+    engine_kwargs=dict(
+        limit_mm_per_prompt={"audio": 1},
+    ),
+    batch_size=16,
+    accelerator_type="L4",
+    concurrency=1,
+    prepare_multimodal_stage={
+        "enabled": True,
+        "chat_template_content_format": "openai",
+    },
+    chat_template_stage=True,
+    tokenize_stage=True,
+    detokenize_stage=True,
+)
+# __omni_audio_config_example_end__
+
+
+# __omni_audio_preprocess_example_start__
+def audio_preprocess(row: dict) -> dict:
+    """
+    Preprocessing function for audio-language model inputs.
+
+    Converts dataset rows into the format expected by the Omni model:
+    - System prompt for analysis instructions
+    - User message with text and audio content
+    - Sampling parameters
+    """
+    return {
+        "messages": [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that analyzes audio. "
+                "Listen to the audio carefully and provide detailed descriptions.",
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": row["text"],
+                    },
+                    {
+                        "type": "input_audio",
+                        "input_audio": {
+                            "data": row["audio_data"],
+                            "format": "wav",
+                        },
+                    },
+                ],
+            },
+        ],
+        "sampling_params": {
+            "temperature": 0.3,
+            "max_tokens": 150,
+            "detokenize": False,
+        },
+    }
+
+
+def audio_postprocess(row: dict) -> dict:
+    return {
+        "resp": row["generated_text"],
+    }
+
+
+# __omni_audio_preprocess_example_end__
+
+
+def load_audio_dataset():
+# __omni_audio_load_dataset_example_start__
+    """
+    Load audio dataset from MRSAudio Hugging Face dataset.
+    """
+    try:
+        from datasets import load_dataset
+        from huggingface_hub import hf_hub_download
+        import base64
+
+        dataset_name = "MRSAudio/MRSAudio"
+
+        dataset = load_dataset(dataset_name, split="train")
+
+        audio_items = []
+
+        # Limit to first 10 samples for the example
+        num_samples = min(10, len(dataset))
+        for i in range(num_samples):
+            item = dataset[i]
+
+            audio_path = hf_hub_download(
+                repo_id=dataset_name, filename=item["path"], repo_type="dataset"
+            )
+
+            with open(audio_path, "rb") as f:
+                audio_bytes = f.read()
+
+            audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
+            audio_items.append(
+                {
+                    "audio_data": audio_base64,
+                    "text": item.get("text", "Describe this audio."),
+                }
+            )
+
+        audio_dataset = ray.data.from_items(audio_items)
+        return audio_dataset
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None
+# __omni_audio_load_dataset_example_end__
+
+
+def create_omni_audio_config():
+    """Create Omni audio configuration."""
+    return vLLMEngineProcessorConfig(
+        model_source="Qwen/Qwen2.5-Omni-3B",
+        task_type="generate",
+        engine_kwargs=dict(
+            enforce_eager=True,
+            limit_mm_per_prompt={"audio": 1},
+        ),
+        batch_size=16,
+        accelerator_type="L4",
+        concurrency=1,
+        prepare_multimodal_stage={
+            "enabled": True,
+            "chat_template_content_format": "openai",
+        },
+        chat_template_stage=True,
+        tokenize_stage=True,
+        detokenize_stage=True,
+    )
+
+def run_omni_audio_example():
+# __omni_audio_run_example_start__
+    """Run the complete Omni audio example workflow."""
+    config = create_omni_audio_config()
+    audio_dataset = load_audio_dataset()
+
+    if audio_dataset:
+        # Build processor with preprocessing and postprocessing
+        processor = build_processor(
+            config, preprocess=audio_preprocess, postprocess=audio_postprocess
+        )
+
+        print("Omni audio processor configured successfully")
+        print(f"Model: {config.model_source}")
+        print(f"Has multimodal support: {config.prepare_multimodal_stage.get('enabled', False)}")
+        result = processor(audio_dataset).take_all()
+        return config, processor, result
+# __omni_audio_run_example_end__
+    return None, None, None
+
+
+# __omni_audio_example_end__
+
+if __name__ == "__main__":
+    # Run the example Omni audio workflow only if GPU is available
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_omni_audio_example()
+        else:
+            print("Skipping Omni audio example run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping Omni audio example run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/openai_api_example.py

@@ -0,0 +1,99 @@
+"""
+This file serves as a documentation example and CI test for OpenAI API batch inference.
+
+"""
+
+import os
+from ray.data.llm import HttpRequestProcessorConfig, build_processor
+
+
+def run_openai_example():
+    # __openai_example_start__
+    import ray
+
+    OPENAI_KEY = os.environ["OPENAI_API_KEY"]
+    ds = ray.data.from_items(["Hand me a haiku."])
+
+    config = HttpRequestProcessorConfig(
+        url="https://api.openai.com/v1/chat/completions",
+        headers={"Authorization": f"Bearer {OPENAI_KEY}"},
+        qps=1,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            payload=dict(
+                model="gpt-4o-mini",
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a bot that responds with haikus.",
+                    },
+                    {"role": "user", "content": row["item"]},
+                ],
+                temperature=0.0,
+                max_tokens=150,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            response=row["http_response"]["choices"][0]["message"]["content"]
+        ),
+    )
+
+    ds = processor(ds)
+    print(ds.take_all())
+    # __openai_example_end__
+
+
+def run_openai_demo():
+    """Run the OpenAI API configuration demo."""
+    print("OpenAI API Configuration Demo")
+    print("=" * 30)
+    print("\nExample configuration:")
+    print("config = HttpRequestProcessorConfig(")
+    print("    url='https://api.openai.com/v1/chat/completions',")
+    print("    headers={'Authorization': f'Bearer {OPENAI_KEY}'},")
+    print("    qps=1,")
+    print(")")
+    print("\nThe processor handles:")
+    print("- Preprocessing: Convert text to OpenAI API format")
+    print("- HTTP requests: Send batched requests to OpenAI")
+    print("- Postprocessing: Extract response content")
+
+
+def preprocess_for_openai(row):
+    """Preprocess function for OpenAI API requests."""
+    return dict(
+        payload=dict(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": row["item"]},
+            ],
+            temperature=0.0,
+            max_tokens=150,
+        )
+    )
+
+
+def postprocess_openai_response(row):
+    """Postprocess function for OpenAI API responses."""
+    return dict(response=row["http_response"]["choices"][0]["message"]["content"])
+
+
+if __name__ == "__main__":
+    # Run live call if API key is set; otherwise show demo with mock output
+    if "OPENAI_API_KEY" in os.environ:
+        run_openai_example()
+    else:
+        # Mock response without API key
+        print(
+            [
+                {
+                    "response": (
+                        "Autumn leaves whisper\nSoft code flows in quiet lines\nBugs fall one by one"
+                    )
+                }
+            ]
+        )

doc/source/data/doc_code/working-with-llms/tokenization_aggregation_example.py

@@ -0,0 +1,53 @@
+"""
+This file serves as a documentation example for aggregated tokenization.
+
+"""
+
+import ray
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+    TokenizerStageConfig,
+    DetokenizeStageConfig,
+)
+
+# __aggregated_tokenization_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={"max_model_len": 4096},
+    concurrency=1,
+    batch_size=64,
+    tokenize_stage=TokenizerStageConfig(enabled=False),
+    detokenize_stage=DetokenizeStageConfig(enabled=False),
+)
+
+processor = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        messages=[
+            {"role": "user", "content": row["item"]},
+        ],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=250,
+            # Let the vLLM engine handle detokenization
+            detokenize=True,
+        ),
+    ),
+    postprocess=lambda row: dict(resp=row["generated_text"]),
+)
+# __aggregated_tokenization_end__
+
+ds = ray.data.from_items(["Hello world!"])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            ds = processor(ds)
+            ds.show(limit=1)
+        else:
+            print("Skipping aggregated run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping aggregated run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/tokenization_disaggregation_example.py

@@ -0,0 +1,48 @@
+"""
+This file serves as a documentation example for disaggregated tokenization.
+
+"""
+
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# __disaggregated_tokenization_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={"max_model_len": 4096},
+    concurrency=1,
+    batch_size=64,
+    tokenize_stage=True,
+    detokenize_stage=True,
+)
+
+processor = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        messages=[
+            {"role": "user", "content": row["item"]},
+        ],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=250,
+            # Let the vLLMEngineProcessor's CPU detokenize stage handle detokenization
+            detokenize=False,
+        ),
+    ),
+    postprocess=lambda row: dict(resp=row["generated_text"]),
+)
+# __disaggregated_tokenization_end__
+
+ds = ray.data.from_items(["Hello world!"])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            ds = processor(ds)
+            ds.show(limit=1)
+        else:
+            print("Skipping disaggregated run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping disaggregated run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/vlm_image_example.py

@@ -0,0 +1,217 @@
+"""
+This file serves as a documentation example and CI test for VLM batch inference with images.
+
+Structure:
+1. Infrastructure setup: Dataset compatibility patches, dependency handling
+2. Docs example (between __vlm_image_example_start/end__): Embedded in Sphinx docs via literalinclude
+3. Test validation and cleanup
+"""
+
+
+'''
+# __image_message_format_example_start__
+"""Supported image input formats: image URL, PIL Image object"""
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "Provide a detailed description of the image."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe what happens in this image."},
+                # Option 1: Provide image URL
+                {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}},
+                # Option 2: Provide PIL Image object
+                {"type": "image_pil", "image_pil": PIL.Image.open("path/to/image.jpg")}
+            ]
+        },
+    ]
+}
+# __image_message_format_example_end__
+'''
+
+
+# __vlm_image_example_start__
+import ray
+from PIL import Image
+from io import BytesIO
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+)
+from huggingface_hub import HfFileSystem
+
+# Load "LMMs-Eval-Lite" dataset from Hugging Face using HfFileSystem
+path = "hf://datasets/lmms-lab/LMMs-Eval-Lite/coco2017_cap_val/"
+fs = HfFileSystem()
+vision_dataset = ray.data.read_parquet(path, filesystem=fs)
+
+# __vlm_config_example_start__
+vision_processor_config = vLLMEngineProcessorConfig(
+    model_source="Qwen/Qwen2.5-VL-3B-Instruct",
+    engine_kwargs=dict(
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+        max_model_len=4096,
+        trust_remote_code=True,
+        limit_mm_per_prompt={"image": 1},
+    ),
+    batch_size=16,
+    concurrency=1,
+    prepare_multimodal_stage=True,
+)
+# __vlm_config_example_end__
+
+
+# __vlm_preprocess_example_start__
+def vision_preprocess(row: dict) -> dict:
+    """
+    Preprocessing function for vision-language model inputs.
+
+    Converts dataset rows into the format expected by the VLM:
+    - System prompt for analysis instructions
+    - User message with text and image content
+    - Multiple choice formatting
+    - Sampling parameters
+    """
+    choice_indices = ["A", "B", "C", "D", "E", "F", "G", "H"]
+
+    return {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "Analyze the image and question carefully, using step-by-step reasoning. "
+                    "First, describe any image provided in detail. Then, present your reasoning. "
+                    "And finally your final answer in this format: Final Answer: <answer> "
+                    "where <answer> is: The single correct letter choice A, B, C, D, E, F, etc. when options are provided. "
+                    "Only include the letter. Your direct answer if no options are given, as a single phrase or number. "
+                    "IMPORTANT: Remember, to end your answer with Final Answer: <answer>."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": row["question"] + "\n\n"},
+                    {
+                        "type": "image_pil",
+                        "image_pil": Image.open(BytesIO(row["image"]["bytes"])),
+                    },
+                    {
+                        "type": "text",
+                        "text": "\n\nChoices:\n"
+                        + "\n".join(
+                            [
+                                f"{choice_indices[i]}. {choice}"
+                                for i, choice in enumerate(row["answer"])
+                            ]
+                        ),
+                    },
+                ],
+            },
+        ],
+        "sampling_params": {
+            "temperature": 0.3,
+            "max_tokens": 150,
+            "detokenize": False,
+        },
+        # Include original data for reference
+        "original_data": {
+            "question": row["question"],
+            "answer_choices": row["answer"],
+            "image_size": row["image"].get("width", 0) if row["image"] else 0,
+        },
+    }
+
+
+def vision_postprocess(row: dict) -> dict:
+    return {
+        "resp": row["generated_text"],
+    }
+
+
+# __vlm_preprocess_example_end__
+
+
+def load_vision_dataset():
+# __vlm_image_load_dataset_example_start__
+    """
+    Load vision dataset from Hugging Face.
+
+    This function loads the LMMs-Eval-Lite dataset which contains:
+    - Images with associated questions
+    - Multiple choice answers
+    - Various visual reasoning tasks
+    """
+    try:
+        from huggingface_hub import HfFileSystem
+
+        # Load "LMMs-Eval-Lite" dataset from Hugging Face using HfFileSystem
+        path = "hf://datasets/lmms-lab/LMMs-Eval-Lite/coco2017_cap_val/"
+        fs = HfFileSystem()
+        vision_dataset = ray.data.read_parquet(path, filesystem=fs)
+
+        return vision_dataset
+    except ImportError:
+        print(
+            "huggingface_hub package not available. Install with: pip install huggingface_hub"
+        )
+        return None
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None
+# __vlm_image_load_dataset_example_end__
+
+
+def create_vlm_config():
+    """Create VLM configuration."""
+    return vLLMEngineProcessorConfig(
+        model_source="Qwen/Qwen2.5-VL-3B-Instruct",
+        engine_kwargs=dict(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            max_model_len=4096,
+            trust_remote_code=True,
+            limit_mm_per_prompt={"image": 1},
+        ),
+        batch_size=1,
+        concurrency=1,
+        prepare_multimodal_stage=True,
+    )
+
+
+def run_vlm_example():
+# __vlm_run_example_start__
+    """Run the complete VLM example workflow."""
+    config = create_vlm_config()
+    vision_dataset = load_vision_dataset()
+
+    if vision_dataset:
+        # Build processor with preprocessing and postprocessing
+        processor = build_processor(
+            config, preprocess=vision_preprocess, postprocess=vision_postprocess
+        )
+
+        print("VLM processor configured successfully")
+        print(f"Model: {config.model_source}")
+        result = processor(vision_dataset).take_all()
+        return config, processor, result
+# __vlm_run_example_end__
+    return None, None, None
+
+
+# __vlm_image_example_end__
+
+if __name__ == "__main__":
+    # Run the example VLM workflow only if GPU is available
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_vlm_example()
+        else:
+            print("Skipping VLM example run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping VLM example run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/vlm_video_example.py

@@ -0,0 +1,241 @@
+"""
+This file serves as a documentation example and CI test for VLM batch inference with videos.
+
+Structure:
+1. Infrastructure setup: Dataset compatibility patches, dependency handling
+2. Docs example (between __vlm_video_example_start/end__): Embedded in Sphinx docs via literalinclude
+3. Test validation and cleanup
+"""
+
+import os
+
+'''
+# __video_message_format_example_start__
+"""Supported video input formats: video URL"""
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "Provide a detailed description of the video."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe what happens in this video."},
+                # Provide video URL
+                {"type": "video_url", "video_url": {"url": "https://example.com/video.mp4"}},
+            ]
+        },
+    ]
+}
+# __video_message_format_example_end__
+'''
+
+# __vlm_video_example_start__
+import ray
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+)
+
+
+# __vlm_video_config_example_start__
+video_processor_config = vLLMEngineProcessorConfig(
+    model_source="Qwen/Qwen3-VL-4B-Instruct",
+    engine_kwargs=dict(
+        tensor_parallel_size=4,
+        pipeline_parallel_size=1,
+        trust_remote_code=True,
+        limit_mm_per_prompt={"video": 1},
+        mm_processor_kwargs={
+            "size": {
+                "shortest_edge": 65536,
+                "longest_edge": 20 * 1088 * 1920,
+            },
+            "do_sample_frames": False,
+        },
+    ),
+    batch_size=1,
+    accelerator_type="L4",
+    concurrency=1,
+    prepare_multimodal_stage={
+        "enabled": True,
+        "model_config_kwargs": dict(
+            # See available model config kwargs at https://docs.vllm.ai/en/latest/api/vllm/config/#vllm.config.ModelConfig
+            allowed_local_media_path="/tmp",
+            media_io_kwargs={"video": {"num_frames": 20, "fps": 2}},
+        ),
+    },
+    chat_template_stage=True,
+    tokenize_stage=True,
+    detokenize_stage=True,
+)
+# __vlm_video_config_example_end__
+
+
+# __vlm_video_preprocess_example_start__
+def video_preprocess(row: dict) -> dict:
+    """
+    Preprocessing function for video-language model inputs.
+
+    Converts dataset rows into the format expected by the VLM:
+    - System prompt for analysis instructions
+    - User message with text and video content
+    - Sampling parameters
+    - Multimodal processor kwargs for video processing
+    """
+    return {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "You are a helpful assistant that analyzes videos. "
+                    "Watch the video carefully and provide detailed descriptions."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": row["text"],
+                    },
+                    {
+                        "type": "video_url",
+                        "video_url": {"url": row["video_url"]},
+                    },
+                ],
+            },
+        ],
+        "sampling_params": {
+            "temperature": 0.3,
+            "max_tokens": 150,
+            "detokenize": False,
+        },
+    }
+
+
+def video_postprocess(row: dict) -> dict:
+    return {
+        "resp": row["generated_text"],
+    }
+
+
+# __vlm_video_preprocess_example_end__
+
+
+def load_video_dataset():
+# __vlm_video_load_dataset_example_start__
+    """
+    Load video dataset from ShareGPTVideo Hugging Face dataset.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+        import tarfile
+        from pathlib import Path
+
+        dataset_name = "ShareGPTVideo/train_raw_video"
+
+        tar_path = hf_hub_download(
+            repo_id=dataset_name,
+            filename="activitynet/chunk_0.tar.gz",
+            repo_type="dataset",
+        )
+
+        extract_dir = "/tmp/sharegpt_videos"
+        os.makedirs(extract_dir, exist_ok=True)
+
+        if not any(Path(extract_dir).glob("*.mp4")):
+            with tarfile.open(tar_path, "r:gz") as tar:
+                tar.extractall(extract_dir)
+
+        video_files = list(Path(extract_dir).rglob("*.mp4"))
+
+        # Limit to first 10 videos for the example
+        video_files = video_files[:10]
+
+        video_dataset = ray.data.from_items(
+            [
+                {
+                    "video_path": str(video_file),
+                    "video_url": f"file://{video_file}",
+                    "text": "Describe what happens in this video.",
+                }
+                for video_file in video_files
+            ]
+        )
+
+        return video_dataset
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None
+# __vlm_video_load_dataset_example_end__
+
+def create_vlm_video_config():
+    """Create VLM video configuration."""
+    return vLLMEngineProcessorConfig(
+        model_source="Qwen/Qwen3-VL-4B-Instruct",
+        engine_kwargs=dict(
+            tensor_parallel_size=4,
+            pipeline_parallel_size=1,
+            trust_remote_code=True,
+            limit_mm_per_prompt={"video": 1},
+            mm_processor_kwargs={
+                "size": {
+                    "shortest_edge": 65536,
+                    "longest_edge": 20 * 1088 * 1920,
+                },
+                "do_sample_frames": False,
+            },
+        ),
+        batch_size=1,
+        accelerator_type="L4",
+        concurrency=1,
+        prepare_multimodal_stage={
+            "enabled": True,
+            "model_config_kwargs": dict(
+                # See available model config kwargs at https://docs.vllm.ai/en/latest/api/vllm/config/#vllm.config.ModelConfig
+                allowed_local_media_path="/tmp",
+                media_io_kwargs={"video": {"num_frames": 20, "fps": 2}},
+            ),
+        },
+        chat_template_stage=True,
+        tokenize_stage=True,
+        detokenize_stage=True,
+    )
+
+
+def run_vlm_video_example():
+# __vlm_video_run_example_start__
+    """Run the complete VLM video example workflow."""
+    config = create_vlm_video_config()
+    video_dataset = load_video_dataset()
+
+    if video_dataset:
+        # Build processor with preprocessing and postprocessing
+        processor = build_processor(
+            config, preprocess=video_preprocess, postprocess=video_postprocess
+        )
+
+        print("VLM video processor configured successfully")
+        print(f"Model: {config.model_source}")
+        print(f"Has multimodal support: {config.prepare_multimodal_stage.get('enabled', False)}")
+        result = processor(video_dataset).take_all()
+        return config, processor, result
+# __vlm_video_run_example_end__
+    return None, None, None
+
+
+# __vlm_video_example_end__
+
+if __name__ == "__main__":
+    # Run the example VLM video workflow only if GPU is available
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_vlm_video_example()
+        else:
+            print("Skipping VLM video example run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping VLM video example run due to environment error: {e}")