chore: import upstream snapshot with attribution

release/llm_tests/batch/llm_2x_4xl4.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g6.12xlarge
+      min_nodes: 2
+      max_nodes: 2
+      market_type: ON_DEMAND

release/llm_tests/batch/llm_4xl4.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g6.12xlarge
+      min_nodes: 0
+      max_nodes: 1
+      market_type: ON_DEMAND

release/llm_tests/batch/llm_single_node_benchmark_l4.yaml

@@ -0,0 +1,13 @@
+# Single-node compute config for Ray Data LLM baseline benchmark
+# Instance: g6.xlarge (1x NVIDIA L4 GPU, 24GB VRAM)
+
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.large
+
+worker_nodes:
+    - instance_type: g6.xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/llm_tests/batch/requirements.in

@@ -0,0 +1,5 @@
+datasets==4.4.1
+# Keep in sync with huggingface-hub in the llm-cu130 image (transformers 5.x
+# requires huggingface-hub>=1.5.0); the byod layer installs this with --no-deps
+# on top of the image, so a mismatch would downgrade the image's copy.
+huggingface-hub==1.13.0; python_version == "3.12"

release/llm_tests/batch/test_batch_multi_node_vllm.py

@@ -0,0 +1,60 @@
+import pytest
+
+import ray
+from ray.data.llm import build_processor, vLLMEngineProcessorConfig
+
+
+@pytest.fixture(autouse=True)
+def cleanup_ray_resources():
+    """Automatically cleanup Ray resources between tests to prevent conflicts."""
+    yield
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size",
+    [
+        (2, 4),
+        (4, 2),
+    ],
+)
+def test_vllm_multi_node(tp_size, pp_size):
+    config = vLLMEngineProcessorConfig(
+        model_source="facebook/opt-1.3b",
+        engine_kwargs=dict(
+            enable_prefix_caching=True,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=4096,
+            pipeline_parallel_size=pp_size,
+            tensor_parallel_size=tp_size,
+            distributed_executor_backend="ray",
+        ),
+        tokenize_stage=False,
+        detokenize_stage=False,
+        concurrency=1,
+        batch_size=64,
+        chat_template_stage=False,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            prompt=f"You are a calculator. {row['id']} ** 3 = ?",
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=20,
+                detokenize=True,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["generated_text"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = processor(ds)
+    ds = ds.materialize()
+
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)

release/llm_tests/batch/test_batch_sglang.py

@@ -0,0 +1,126 @@
+import sys
+
+import pytest
+
+import ray
+from ray.data.llm import SGLangEngineProcessorConfig, build_processor
+
+
+def test_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 = SGLangEngineProcessorConfig(
+        model_source="unsloth/Llama-3.2-1B-Instruct",
+        engine_kwargs=dict(
+            context_length=2048,
+            disable_cuda_graph=True,
+            dtype="half",
+        ),
+        batch_size=16,
+        concurrency=1,
+        chat_template_stage={"enabled": True, "chat_template": chat_template},
+        tokenize_stage=True,
+        detokenize_stage=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_new_tokens=50,  # SGLang uses max_new_tokens instead of max_tokens
+            ),
+        ),
+        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)
+
+
+@pytest.mark.parametrize(
+    "tp_size,dp_size,concurrency",
+    [
+        (2, 1, 2),
+        (2, 2, 1),
+    ],
+)
+def test_sglang_llama_parallel(tp_size, dp_size, concurrency):
+    """Test SGLang with Llama model using different parallelism configurations."""
+    runtime_env = {}
+
+    processor_config = SGLangEngineProcessorConfig(
+        model_source="unsloth/Llama-3.2-1B-Instruct",
+        engine_kwargs=dict(
+            context_length=2048,
+            tp_size=tp_size,
+            dp_size=dp_size,
+            dtype="half",
+        ),
+        runtime_env=runtime_env,
+        tokenize_stage=True,
+        detokenize_stage=True,
+        batch_size=16,
+        concurrency=concurrency,
+    )
+
+    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_new_tokens=50,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+        },
+    )
+
+    ds = ray.data.range(120)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = processor(ds)
+    ds = ds.materialize()
+
+    # Verify results
+    outs = ds.take_all()
+    assert len(outs) == 120
+    assert all("resp" in out for out in outs)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

release/llm_tests/batch/test_batch_single_node_vllm.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python
+"""
+Single-node vLLM baseline benchmark for Ray Data LLM batch inference.
+
+Measures throughput and supports env-driven thresholds and
+JSON artifact output.
+"""
+import json
+import os
+import sys
+
+import pytest
+
+import ray
+from ray.llm._internal.batch.benchmark.dataset import ShareGPTDataset
+from ray.llm._internal.batch.benchmark.benchmark_processor import (
+    Mode,
+    VLLM_SAMPLING_PARAMS,
+    benchmark,
+)
+
+
+# Benchmark constants
+NUM_REQUESTS = 1000
+MODEL_ID = "facebook/opt-1.3b"
+BATCH_SIZE = 64
+CONCURRENCY = 1
+
+
+@pytest.fixture(autouse=True)
+def disable_vllm_compile_cache(monkeypatch):
+    """Disable vLLM compile cache to avoid cache corruption."""
+    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
+
+
+@pytest.fixture(autouse=True)
+def cleanup_ray_resources():
+    """Cleanup Ray resources between tests."""
+    yield
+    ray.shutdown()
+
+
+def _get_float_env(name: str, default: float | None = None) -> float | None:
+    value = os.getenv(name)
+    if value is None or value == "":
+        return default
+    try:
+        return float(value)
+    except ValueError:
+        raise AssertionError(f"Invalid float for {name}: {value}")
+
+
+def test_single_node_baseline_benchmark():
+    """
+    Single-node baseline benchmark: facebook/opt-1.3b, TP=1, PP=1, 1000 prompts.
+
+    Logs BENCHMARK_* metrics and optionally asserts perf thresholds from env:
+    - RAY_DATA_LLM_BENCHMARK_MIN_THROUGHPUT (req/s)
+    - RAY_DATA_LLM_BENCHMARK_MAX_LATENCY_S (seconds)
+    Writes JSON artifact to RAY_LLM_BENCHMARK_ARTIFACT_PATH if set.
+    """
+    # Dataset setup
+    dataset_path = os.getenv(
+        "RAY_LLM_BENCHMARK_DATASET_PATH", "/tmp/ray_llm_benchmark_dataset"
+    )
+
+    dataset = ShareGPTDataset(
+        dataset_path=dataset_path,
+        seed=0,
+        hf_dataset_id="Crystalcareai/Code-feedback-sharegpt-renamed",
+        hf_split="train",
+        truncate_prompt=2048,
+    )
+
+    print(f"Loading {NUM_REQUESTS} prompts from ShareGPT dataset...")
+    prompts = dataset.sample(num_requests=NUM_REQUESTS)
+    print(f"Loaded {len(prompts)} prompts")
+
+    ds = ray.data.from_items(prompts)
+
+    # Benchmark config (single node, TP=1, PP=1)
+    print(
+        f"\nBenchmark: {MODEL_ID}, batch={BATCH_SIZE}, concurrency={CONCURRENCY}, TP=1, PP=1"
+    )
+
+    # Use benchmark processor to run a single-node vLLM benchmark
+    result = benchmark(
+        Mode.VLLM_ENGINE,
+        ds,
+        batch_size=BATCH_SIZE,
+        concurrency=CONCURRENCY,
+        model=MODEL_ID,
+        sampling_params=VLLM_SAMPLING_PARAMS,
+        pipeline_parallel_size=1,
+        tensor_parallel_size=1,
+    )
+
+    result.show()
+
+    # Assertions and metrics
+    assert result.samples == len(prompts)
+    assert result.throughput > 0
+
+    print("\n" + "=" * 60)
+    print("BENCHMARK METRICS")
+    print("=" * 60)
+    print(f"BENCHMARK_THROUGHPUT: {result.throughput:.4f} req/s")
+    print(f"BENCHMARK_LATENCY: {result.elapsed_s:.4f} s")
+    print(f"BENCHMARK_SAMPLES: {result.samples}")
+    print("=" * 60)
+
+    # Optional thresholds to fail on regressions
+    min_throughput = _get_float_env("RAY_DATA_LLM_BENCHMARK_MIN_THROUGHPUT", 5)
+    max_latency_s = _get_float_env("RAY_DATA_LLM_BENCHMARK_MAX_LATENCY_S", 150)
+    if min_throughput is not None:
+        assert (
+            result.throughput >= min_throughput
+        ), f"Throughput regression: {result.throughput:.4f} < {min_throughput:.4f} req/s"
+    if max_latency_s is not None:
+        assert (
+            result.elapsed_s <= max_latency_s
+        ), f"Latency regression: {result.elapsed_s:.4f} > {max_latency_s:.4f} s"
+
+    # Optional JSON artifact emission for downstream ingestion
+    artifact_path = os.getenv("RAY_LLM_BENCHMARK_ARTIFACT_PATH")
+    if artifact_path:
+        metrics = {
+            "model": MODEL_ID,
+            "batch_size": BATCH_SIZE,
+            "concurrency": CONCURRENCY,
+            "samples": int(result.samples),
+            "throughput_req_per_s": float(result.throughput),
+            "elapsed_s": float(result.elapsed_s),
+        }
+        try:
+            os.makedirs(os.path.dirname(artifact_path), exist_ok=True)
+            with open(artifact_path, "w", encoding="utf-8") as f:
+                json.dump(metrics, f, indent=2, sort_keys=True)
+            print(f"Wrote benchmark artifact to: {artifact_path}")
+        except Exception as e:  # noqa: BLE001
+            print(
+                f"Warning: failed to write benchmark artifact to {artifact_path}: {e}"
+            )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

release/llm_tests/batch/test_batch_vllm.py

@@ -0,0 +1,779 @@
+import sys
+import logging
+import time
+
+import pytest
+
+import ray
+from ray.data.llm import (
+    build_processor,
+    vLLMEngineProcessorConfig,
+    ChatTemplateStageConfig,
+    DetokenizeStageConfig,
+    PrepareMultimodalStageConfig,
+    TokenizerStageConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+S3_ARTIFACT_ASSETS_URL = (
+    "https://air-example-data.s3.amazonaws.com/rayllm-ossci/assets/"
+)
+
+
+@pytest.fixture(autouse=True)
+def disable_vllm_compile_cache(monkeypatch):
+    """Automatically disable vLLM compile cache for all tests.
+
+    Avoids AssertionError due to torch compile cache corruption caused by
+    running multiple engines on the same node.
+    See: https://github.com/vllm-project/vllm/issues/18851, fix expected with
+    PyTorch 2.8.0
+    """
+    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
+
+
+@pytest.fixture(autouse=True)
+def add_buffer_time_between_tests():
+    """Add buffer time after each test to avoid resource conflicts, which cause
+    flakiness.
+    """
+    import gc
+
+    gc.collect()
+    time.sleep(15)
+
+
+@pytest.fixture(autouse=True)
+def cleanup_ray_resources():
+    """Automatically cleanup Ray resources between tests to prevent conflicts."""
+    yield
+    _cleanup_gpu_processes()
+    ray.shutdown()
+
+
+def _cleanup_gpu_processes():
+    """
+    Kill GPU processes on all nodes in the cluster. With Ray as the external orchestrator,
+    mp backend suffers from uncoordinated shutdown issues, leaving orphaned GPU processes.
+
+    TODO (jeffreywang): Remove this once https://github.com/vllm-project/vllm/pull/39846 lands.
+    """
+    if not ray.is_initialized():
+        return
+
+    @ray.remote(num_cpus=0)
+    def _remote_kill_gpu_processes():
+        import os
+        import signal
+
+        import pynvml
+
+        pids = set()
+        try:
+            pynvml.nvmlInit()
+            device_count = pynvml.nvmlDeviceGetCount()
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                for proc in pynvml.nvmlDeviceGetComputeRunningProcesses(handle):
+                    pids.add(proc.pid)
+            pynvml.nvmlShutdown()
+        except Exception:
+            pass
+
+        for pid in pids:
+            try:
+                os.kill(pid, signal.SIGKILL)
+            except (ProcessLookupError, ValueError):
+                pass
+
+    try:
+        nodes = ray.nodes()
+        refs = []
+        for node in nodes:
+            if not node.get("Alive", False):
+                continue
+            node_id = node["NodeID"]
+            refs.append(
+                _remote_kill_gpu_processes.options(
+                    scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+                        node_id=node_id, soft=False
+                    ),
+                ).remote()
+            )
+        if refs:
+            ray.get(refs, timeout=30)
+    except Exception as e:
+        logging.warning(f"Failed to kill GPU processes on remote nodes: {e}")
+
+
+@pytest.mark.asyncio
+async def test_vllm_multimodal_utils():
+    """Test vLLM's multimodal utilities.
+
+    This test is adapted from https://github.com/vllm-project/vllm/blob/main/tests/entrypoints/test_chat_utils.py.
+    `parse_chat_messages_async` is thoroughly tested in vLLM. This test serves as an
+    integration test to verify that the function isn't moved to an unexpected location and its signature isn't changed.
+    """
+    from vllm.config import ModelConfig
+    from vllm.entrypoints.chat_utils import parse_chat_messages_async
+
+    image_url = "https://air-example-data.s3.us-west-2.amazonaws.com/rayllm-ossci/assets/cherry_blossom.jpg"
+    image_uuid = str(hash(image_url))
+
+    conversation, mm_data, mm_uuids = await parse_chat_messages_async(
+        [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": image_url},
+                        "uuid": image_uuid,
+                    },
+                    {"type": "text", "text": "What's in the image?"},
+                ],
+            }
+        ],
+        ModelConfig(
+            "microsoft/Phi-3.5-vision-instruct",
+            runner="generate",
+            trust_remote_code=True,
+            limit_mm_per_prompt={"image": 2},
+        ),
+        content_format="string",
+    )
+
+    assert conversation == [
+        {"role": "user", "content": "<|image_1|>\nWhat's in the image?"}
+    ]
+
+    assert mm_data is not None
+    assert set(mm_data.keys()) == {"image"}
+
+    image_data = mm_data.get("image")
+    assert image_data is not None
+
+    assert isinstance(image_data, list) and len(image_data) == 1
+
+    assert mm_uuids is not None
+    assert "image" in mm_uuids
+
+    image_uuids = mm_uuids.get("image")
+    assert image_uuids is not None
+    assert isinstance(image_uuids, list) and len(image_uuids) == 1
+    assert image_uuids[0] == image_uuid
+
+
+def test_chat_template_with_vllm():
+    """Test vLLM with explicit chat template."""
+
+    processor_config = vLLMEngineProcessorConfig(
+        model_source="unsloth/Llama-3.2-1B-Instruct",
+        engine_kwargs=dict(
+            max_model_len=16384,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=2048,
+        ),
+        tokenize_stage=True,
+        detokenize_stage=True,
+        batch_size=16,
+        concurrency=1,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    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)
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size,concurrency",
+    [
+        (2, 1, 2),  # TP=2, concurrency=2
+        (1, 2, 2),  # PP=2, concurrency=2
+    ],
+)
+def test_vllm_llama_parallel(tp_size, pp_size, concurrency):
+    """Test vLLM with Llama model using different parallelism configurations."""
+
+    # vLLM v1 does not support decoupled tokenizer,
+    # but since the tokenizer is in a separate process,
+    # the overhead should be moderated.
+    tokenize = False
+    detokenize = False
+
+    processor_config = vLLMEngineProcessorConfig(
+        model_source="unsloth/Llama-3.2-1B-Instruct",
+        engine_kwargs=dict(
+            tensor_parallel_size=tp_size,
+            pipeline_parallel_size=pp_size,
+            max_model_len=16384,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=2048,
+        ),
+        tokenize_stage=tokenize,
+        detokenize_stage=detokenize,
+        batch_size=16,
+        accelerator_type=None,
+        concurrency=concurrency,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    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(120)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 5})
+    ds = processor(ds)
+    ds = ds.materialize()
+
+    # Verify results
+    outs = ds.take_all()
+    assert len(outs) == 120
+    assert all("resp" in out for out in outs)
+
+
+def test_vllm_llama_lora():
+    """Test vLLM with Llama model and LoRA adapter support.
+
+    Validates that the LoRA adapter is actually applied during generation by
+    using greedy sampling (temperature=0) and short, diverse, open-ended prompts,
+    then asserting that a significant number of base-model vs LoRA paired outputs differ.
+
+    If the LoRA adapter is not being applied, we expect very few pairs to differ
+    due to non-determinism from vLLM's dynamic batching and CUDA. If it is being applied,
+    we expect a significant fraction of pairs to differ due to the adapter's effect on the output.
+
+    The two regimes are separated by the `min_diff_fraction` threshold defined below.
+    """
+    # Minimum fraction of base/LoRA pairs that must differ to pass.
+    # Determined empirically.
+    min_diff_fraction = 0.5
+    model_source = "s3://air-example-data/llama-3.2-216M-dummy/"
+    lora_path = "s3://air-example-data/"
+    lora_name = "llama-3.2-216M-lora-dummy"
+    max_lora_rank = 32
+
+    # Short, diverse prompts — shorter prompts let the LoRA perturbation
+    # manifest earlier in generation before the base model's momentum
+    # dominates the output. Also prefer open-ended prompts to encourage LoRA-induced diversity.
+    prompts = [
+        "Hello world",
+        "The capital of France is",
+        "Once upon a time",
+        "1 + 1 =",
+        "def fibonacci(n):",
+        "The quick brown fox",
+        "In the beginning",
+        "To be or not to be",
+        "import numpy as np",
+        "The weather today is",
+        "My favorite color is",
+        "How to cook rice:",
+        "The meaning of life is",
+        "SELECT * FROM",
+        "Dear Sir or Madam,",
+        "Breaking news:",
+        "A long time ago in a galaxy",
+        "The first law of thermodynamics",
+        "class MyClass:",
+        "Roses are red,",
+        "According to recent studies,",
+        "Step 1: Preheat the oven",
+        "The president announced",
+        "In mathematics, a prime number",
+        "function hello() {",
+        "The cat sat on the",
+        "Water boils at",
+        "Happy birthday to",
+        "ERROR: NullPointerException",
+        "The mitochondria is the",
+    ]
+    num_pairs = len(prompts)
+
+    # The following controls are propagated to the vLLM worker to minimize non-determinism:
+    #   * engine_kwargs["seed"]: vLLM seeds its internal torch/np/random state.
+    #   * PYTHONHASHSEED: stabilizes dict/set iteration order in the worker.
+    #   * CUBLAS_WORKSPACE_CONFIG: forces deterministic cuBLAS workspace.
+    # Flash-attention kernels and the vLLM v1 async scheduler still introduce
+    # residual non-determinism we can't eliminate from the test, but in
+    # practice the regime separation (base/LoRA divergence signal vs noise)
+    # is very wide (~30/30 vs ~5/30), so half-of-num_pairs is a robust threshold.
+    # Note for future: VLLM_BATCH_INVARIANT=1 with attention_backend=FLASH_ATTN
+    # eliminates the remaining non-determinism entirely for perfect separation,
+    # but it requires a GPU compute capability (>=9.0) the CI runners don't meet, so
+    # we rely on the median threshold instead. Recommended if CI hardware is upgraded.
+    seed = 42
+
+    processor_config = vLLMEngineProcessorConfig(
+        model_source=model_source,
+        dynamic_lora_loading_path=lora_path,
+        engine_kwargs=dict(
+            max_model_len=4096,
+            enable_chunked_prefill=True,
+            enable_lora=True,
+            max_lora_rank=max_lora_rank,
+            seed=seed,
+        ),
+        tokenize_stage=True,
+        detokenize_stage=True,
+        # minimize non-determinism from vLLM's dynamic batching by sending (1 base + 1 LoRA) per batch
+        batch_size=2,
+        concurrency=1,
+        runtime_env={
+            "env_vars": {
+                "VLLM_DISABLE_COMPILE_CACHE": "1",
+                "PYTHONHASHSEED": "0",
+                "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
+            }
+        },
+    )
+
+    processor = build_processor(
+        processor_config,
+        preprocess=lambda row: dict(
+            model=model_source if row["id"] % 2 == 0 else lora_name,
+            messages=[{"role": "user", "content": row["prompt"]}],
+            sampling_params=dict(
+                temperature=0,
+                max_tokens=50,
+                detokenize=False,
+            ),
+        ),
+        postprocess=lambda row: {
+            "id": row["id"],
+            "resp": row["generated_text"],
+        },
+    )
+
+    # Build paired rows: id 2k = base, id 2k+1 = LoRA, same prompt.
+    rows = []
+    for i, prompt in enumerate(prompts):
+        rows.append({"id": 2 * i, "prompt": prompt})
+        rows.append({"id": 2 * i + 1, "prompt": prompt})
+
+    ds = ray.data.from_items(rows)
+    ds = processor(ds)
+    ds = ds.materialize()
+    outs = ds.take_all()
+
+    assert len(outs) == 2 * num_pairs
+    assert all("resp" in out for out in outs)
+
+    # LoRA exercise check: a significant fraction of pairs must differ.
+    by_id = {out["id"]: out["resp"] for out in outs}
+    diffs = sum(1 for k in range(num_pairs) if by_id[2 * k] != by_id[2 * k + 1])
+    min_diffs = int(num_pairs * min_diff_fraction)
+    assert diffs >= min_diffs, (
+        f"Only {diffs}/{num_pairs} base/LoRA pairs differ (need >= {min_diffs}) — "
+        "the LoRA adapter does not appear to be applied."
+    )
+
+
+@pytest.mark.parametrize(
+    "model_source,tp_size,pp_size,concurrency,sample_size,chat_template_content_format,apply_sys_msg_formatting",
+    [
+        # LLaVA model with TP=1, PP=1, concurrency=1
+        ("llava-hf/llava-1.5-7b-hf", 1, 1, 1, 60, "openai", False),
+        # Pixtral model with TP=2, PP=1, concurrency=2
+        ("mistral-community/pixtral-12b", 2, 1, 2, 60, "openai", True),
+    ],
+)
+def test_vllm_vision_language_models(
+    model_source,
+    tp_size,
+    pp_size,
+    concurrency,
+    sample_size,
+    chat_template_content_format,
+    apply_sys_msg_formatting,
+):
+    """Test vLLM with vision language models using different configurations."""
+
+    # vLLM v1 does not support decoupled tokenizer,
+    # but since the tokenizer is in a separate process,
+    # the overhead should be moderated.
+    tokenize = False
+    detokenize = False
+
+    llm_processor_config = vLLMEngineProcessorConfig(
+        model_source=model_source,
+        task_type="generate",
+        engine_kwargs=dict(
+            tensor_parallel_size=tp_size,
+            pipeline_parallel_size=pp_size,
+            max_model_len=4096,
+            enable_chunked_prefill=True,
+        ),
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(
+            enabled=True,
+            chat_template_content_format=chat_template_content_format,
+            apply_sys_msg_formatting=apply_sys_msg_formatting,
+        ),
+        chat_template_stage=True,
+        tokenize_stage=tokenize,
+        detokenize_stage=detokenize,
+        batch_size=16,
+        concurrency=concurrency,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+    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": S3_ARTIFACT_ASSETS_URL + "cherry_blossom.jpg"
+                            },
+                        },
+                    ],
+                },
+            ],
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=50,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+        },
+    )
+
+    ds = ray.data.range(sample_size)
+    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) == sample_size
+    assert all("resp" in out for out in outs)
+
+
+@pytest.mark.parametrize(
+    "multimodal_content",
+    [
+        {
+            "type": "image_url",
+            "image_url": {"url": S3_ARTIFACT_ASSETS_URL + "cherry_blossom.jpg"},
+        },
+        {
+            "type": "video_url",
+            "video_url": {"url": S3_ARTIFACT_ASSETS_URL + "free-videos.mp4"},
+        },
+    ],
+)
+def test_vllm_qwen_vl_multimodal(multimodal_content):
+    model_source = "Qwen/Qwen2.5-VL-3B-Instruct"
+
+    llm_processor_config = vLLMEngineProcessorConfig(
+        model_source=model_source,
+        task_type="generate",
+        engine_kwargs=dict(
+            enable_chunked_prefill=True,
+            distributed_executor_backend="ray",
+            # A single GPU won't be able to accomodate Qwen/Qwen2.5-VL-3B-Instruct's memory requirements
+            # due to vllm0.12.0 resource/profiling issues.
+            # Issue: https://github.com/vllm-project/vllm/issues/30521.
+            tensor_parallel_size=2,
+            pipeline_parallel_size=1,
+        ),
+        prepare_multimodal_stage=PrepareMultimodalStageConfig(
+            enabled=True,
+        ),
+        chat_template_stage=ChatTemplateStageConfig(enabled=True),
+        tokenize_stage=TokenizerStageConfig(enabled=False),
+        detokenize_stage=DetokenizeStageConfig(enabled=False),
+        batch_size=16,
+        concurrency=1,
+    )
+
+    llm_processor = build_processor(
+        llm_processor_config,
+        preprocess=lambda row: dict(
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=50,
+            ),
+            mm_processor_kwargs=dict(
+                min_pixels=28 * 28,
+                max_pixels=1280 * 28 * 28,
+                fps=1,
+            ),
+            messages=[
+                {"role": "system", "content": "You are an assistant"},
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"Describe this asset in {row['id']} sentences.",
+                        },
+                        multimodal_content,
+                    ],
+                },
+            ],
+        ),
+        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("concurrency", [1, 4])
+def test_async_udf_queue_capped(concurrency):
+    """
+    Test that the large object in input/output rows
+    are stored in object store and does not OOM.
+    """
+
+    processor_config = vLLMEngineProcessorConfig(
+        model_source="unsloth/Llama-3.2-1B-Instruct",
+        engine_kwargs=dict(
+            max_model_len=16384,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=2048,
+        ),
+        tokenize_stage=False,
+        detokenize_stage=False,
+        batch_size=4,
+        accelerator_type=None,
+        concurrency=concurrency,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    processor = build_processor(
+        processor_config,
+        preprocess=lambda row: dict(
+            # 1M emoji (4 bytes), should not leak to memory heap.
+            large_memory_to_carry_over="🤗" * 1_000_000,
+            messages=[
+                {"role": "system", "content": "You are a calculator"},
+                {"role": "user", "content": f"{row['id']} ** 3 = ?"},
+            ],
+            sampling_params=dict(
+                temperature=0.3,
+                # we don't care about the actual output
+                max_tokens=1,
+                detokenize=False,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp": row["generated_text"],
+            "large_memory_still_there": "large_memory_to_carry_over" in row,
+        },
+    )
+
+    ds = ray.data.range(12000)
+
+    def map_id_to_val_in_test_no_memory_leak(x):
+        return {"id": x["id"], "val": x["id"] + 5}
+
+    ds = ds.map(map_id_to_val_in_test_no_memory_leak)
+    ds = processor(ds)
+    ds = ds.materialize()
+
+    outs = ds.take_all()
+    assert all(out["large_memory_still_there"] for out in outs)
+
+
+@pytest.mark.parametrize(
+    "backend, placement_group_config",
+    [
+        # Custom placement group with STRICT_PACK strategy
+        (
+            "ray",
+            dict(bundles=[{"CPU": 1, "GPU": 1}] * 4, strategy="STRICT_PACK"),
+        ),
+        # Strategy omitted (PACK default). Omitted GPU now validates as 0.0; explicit GPU: 1 for this GPU job.
+        (
+            "ray",
+            dict(bundles=[{"CPU": 1, "GPU": 1}] * 4),
+        ),
+        # Empty placement group
+        (
+            "ray",
+            None,
+        ),
+        # Custom placement group with MP backend
+        (
+            "mp",
+            dict(bundles=[{"GPU": 1}] * 4),
+        ),
+        # Empty placement group with MP backend
+        (
+            "mp",
+            None,
+        ),
+    ],
+)
+def test_vllm_placement_group(backend, placement_group_config):
+    """Test vLLM with different placement group configurations."""
+
+    config = vLLMEngineProcessorConfig(
+        model_source="facebook/opt-1.3b",
+        engine_kwargs=dict(
+            enable_prefix_caching=True,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=4096,
+            pipeline_parallel_size=2,
+            tensor_parallel_size=2,
+            distributed_executor_backend=backend,
+        ),
+        tokenize_stage=False,
+        detokenize_stage=False,
+        concurrency=1,
+        batch_size=16,
+        chat_template_stage=False,
+        placement_group_config=placement_group_config,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            prompt=f"You are a calculator. {row['id']} ** 3 = ?",
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=20,
+                detokenize=True,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["generated_text"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    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_vllm_autoscaling_no_starvation():
+    """Test that chained vLLMEngineProcessor instances with autoscaling
+    concurrency can run without starving each other.
+    """
+    processor_config_1 = vLLMEngineProcessorConfig(
+        model_source="facebook/opt-1.3b",
+        chat_template_stage=False,
+        tokenize_stage=False,
+        detokenize_stage=False,
+        batch_size=16,
+        concurrency=(1, 4),
+    )
+
+    processor_config_2 = vLLMEngineProcessorConfig(
+        model_source="facebook/opt-1.3b",
+        chat_template_stage=False,
+        tokenize_stage=False,
+        detokenize_stage=False,
+        batch_size=16,
+        concurrency=(3, 4),
+    )
+
+    processor_1 = build_processor(
+        processor_config_1,
+        preprocess=lambda row: dict(
+            prompt=f"Calculate {row['id']} ** 2 = ",
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=30,
+                detokenize=True,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp_1": row["generated_text"],
+            "id": row.get("id", None),
+        },
+    )
+
+    processor_2 = build_processor(
+        processor_config_2,
+        preprocess=lambda row: dict(
+            prompt=f"Previous result: {row.get('resp_1', 'N/A')}. Now calculate its cube: ",
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=30,
+                detokenize=True,
+            ),
+        ),
+        postprocess=lambda row: {
+            "resp_2": row["generated_text"],
+            "resp_1": row.get("resp_1", None),
+            "id": row.get("id", None),
+        },
+    )
+
+    ds = ray.data.range(60)
+    ds = ds.map(lambda x: {"id": x["id"], "val": x["id"] + 1})
+
+    processed_ds = processor_2(processor_1(ds))
+    processed_ds = processed_ds.materialize()
+    results = processed_ds.take_all()
+
+    assert len(results) == 60
+    assert all("resp_1" in out for out in results)
+    assert all("resp_2" in out for out in results)
+    assert all("id" in out for out in results)
+    assert all(out.get("resp_1") for out in results)
+    assert all(out.get("resp_2") for out in results)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

release/llm_tests/benchmark/bm.py

@@ -0,0 +1,55 @@
+# There is a dead-lock issue that arises due to gevent's monkey-patching
+# https://github.com/ipython/ipython/issues/11730
+# Fix: We do this import first before anything else
+# ruff: noqa: E402
+import gevent.monkey
+
+gevent.monkey.patch_all()
+
+
+from .locust import LLMLoadTester
+from .configs import LoadTestConfig
+
+from typing import List, Optional
+import openai
+
+
+def run_bm(
+    api_url: str,
+    api_key: Optional[str] = None,
+    concurrency: Optional[List[int]] = None,
+    run_time: str = "1m",
+    prompt_tokens: int = 512,
+    max_tokens: int = 64,
+    stream: bool = False,
+    summary_file: str = "./results.csv",
+):
+    if api_key is None:
+        api_key = "NONE"
+
+    # Get model_id
+    client = openai.Client(base_url=f"{api_url}/v1", api_key=api_key)
+    models = client.models.list().model_dump()["data"]
+
+    if len(models) != 1:
+        raise ValueError("The service is expected to have only one model.")
+
+    model_id = models[0]["id"]
+    results = []
+    for n_users in concurrency:
+        config = LoadTestConfig(
+            host=api_url,
+            api_key=api_key,
+            provider="openai",
+            model=model_id,
+            stream=stream,
+            prompt_tokens=prompt_tokens,
+            max_tokens=max_tokens,
+            users=n_users,
+            run_time=run_time,
+            summary_file=summary_file,
+        )
+        tester = LLMLoadTester(config)
+        results.append(tester.run())
+
+    return results

release/llm_tests/benchmark/configs.py

@@ -0,0 +1,156 @@
+from typing import List, Optional
+from enum import Enum
+from pydantic import BaseModel, Field
+import argparse
+
+
+class DistributionType(str, Enum):
+    CONSTANT = "constant"
+    UNIFORM = "uniform"
+    EXPONENTIAL = "exponential"
+    NORMAL = "normal"
+
+
+class TokensDistributionType(str, Enum):
+    CONSTANT = "constant"
+    UNIFORM = "uniform"
+    EXPONENTIAL = "exponential"
+
+
+class LoadTestConfig(BaseModel):
+    provider: Optional[str] = Field(
+        None,
+        description="Which flavor of API to use. If not specified, we'll try to guess based on the URL and /v1/models output",
+    )
+
+    model: Optional[str] = Field(
+        None,
+        description="The model to use for generating text. If not specified we will pick the first model from the service as returned by /v1/models",
+    )
+
+    chat: bool = Field(True, description="Use /v1/chat/completions API")
+
+    prompt_tokens: int = Field(
+        512,
+        description="Length of the prompt in tokens",
+    )
+
+    prompt_chars: Optional[int] = Field(
+        None,
+        description="Length of the prompt in characters",
+    )
+
+    prompt_text: Optional[str] = Field(
+        None,
+        description="Prompt text to use instead of generating one. It can be a file reference starting with an ampersand, e.g. `@prompt.txt`",
+    )
+
+    prompt_randomize: bool = Field(
+        False,
+        description="Include a few random numbers in the generated prompt to avoid caching",
+    )
+
+    max_tokens: int = Field(
+        64,
+        description="Max number of tokens to generate. If max_tokens_distribution is non-constant this is going to be the mean",
+    )
+
+    max_tokens_cap: Optional[int] = Field(
+        None,
+        description="If max_tokens_distribution is non-constant, this truncates the distribition at the specified limit",
+    )
+
+    max_tokens_distribution: TokensDistributionType = Field(
+        TokensDistributionType.CONSTANT,
+        description="How to sample max_tokens on each request",
+    )
+
+    max_tokens_range: float = Field(
+        0.3,
+        description="Specifies the width of the distribution. Specified value `alpha` is relative to `max_tokens`",
+    )
+
+    stream: bool = Field(True, description="Use the streaming API")
+
+    api_key: Optional[str] = Field(
+        None,
+        description="Auth for the API",
+    )
+
+    temperature: float = Field(0.1, description="Temperature parameter for the API")
+
+    logprobs: Optional[int] = Field(
+        None,
+        description="Whether to ask for logprobs, it makes things slower for some providers but is necessary for token count in streaming",
+    )
+
+    summary_file: Optional[str] = Field(
+        None,
+        description="Append the line with the summary to the specified CSV file",
+    )
+
+    qps: Optional[float] = Field(
+        None,
+        description="Enabled 'fixed QPS' mode where requests are issues at the specified rate regardless of how long the processing takes",
+    )
+
+    qps_distribution: DistributionType = Field(
+        DistributionType.CONSTANT,
+        description="Must be used with qps. Specifies how to space out requests",
+    )
+
+    burst: Optional[float] = Field(
+        None,
+        description="Makes requests to arrive in bursts every specified number of seconds",
+    )
+
+    tokenizer: Optional[str] = Field(
+        None,
+        description="Specify HF tokenizer to use for validating the output of the model",
+    )
+
+    show_response: bool = Field(
+        False,
+        description="Print the result of each generation",
+    )
+
+    prompt_cache_max_len: int = Field(
+        0,
+        description="Maximum length of the prompt cache to use",
+    )
+
+    header: List[str] = Field(
+        default_factory=list,
+        description="Arbitrary headers to add to the inference request",
+    )
+
+    n: int = Field(
+        1,
+        description="How many sequences to generate (makes sense to use with non-zero temperature)",
+    )
+
+    host: Optional[str] = Field(
+        default=None,
+        description="Host to load test in the following format: http://10.21.32.33",
+    )
+
+    reset_stats: bool = Field(
+        default=True,
+        description="Determines if stats should be reset once hatching is complete",
+    )
+
+    users: int = Field(
+        default=None,
+        description="Number of concurrent users to spawn for benchmarking.",
+    )
+
+    run_time: str = Field(
+        default="30s",
+        description="The runtime it is in form of Ns, Nm, or Nh, for seconds, minutes, and hours.",
+    )
+
+    def to_namespace(self) -> argparse.Namespace:
+        """
+        Convert the model to an argparse.Namespace object
+        """
+        return argparse.Namespace(**self.dict())

release/llm_tests/benchmark/locust.py

@@ -0,0 +1,105 @@
+from locust.env import Environment
+from locust.stats import stats_printer, stats_history, print_stats
+import gevent
+from typing import Dict, Any
+import logging
+from locust.log import setup_logging
+
+from .load_test import LLMUser, events, collect_metrics
+from .configs import LoadTestConfig
+
+
+class LLMLoadTester:
+    """
+    Usage Example:
+
+    ```python
+        config = LoadTestConfig(
+            host="http://localhost:8000",
+            provider="vllm",
+            model="meta-llama/Meta-Llama-3.1-8B-Instruct",
+            api_key="NONE",
+            prompt_tokens=550,
+            max_tokens=150,
+            users=128,
+            run_time="1m",
+            summary_file="./vllm.csv"
+        )
+        tester = LLMLoadTester(config)
+        results = tester.run()
+    ```
+
+
+    """
+
+    def __init__(self, config: LoadTestConfig):
+        self.config = config
+
+    def _setup_environment(self) -> Environment:
+
+        setup_logging("INFO", None)
+
+        # Setup Environment and Runner
+        env = Environment(
+            user_classes=[LLMUser],
+            host=self.config.host,
+            reset_stats=self.config.reset_stats,
+            events=events,
+        )
+        env.parsed_options = self.config.to_namespace()
+
+        return env
+
+    def run(self) -> Dict[str, Any]:
+        try:
+            # Setup environment
+            env = self._setup_environment()
+            env.create_local_runner()
+
+            # Log test start
+            logging.info(f"Starting test with {self.config.users} users")
+
+            # Create greenlets for stats
+            stats_printer_greenlet = gevent.spawn(stats_printer(env.stats))
+            stats_history_greenlet = gevent.spawn(stats_history, env.runner)
+
+            # Start the test
+            env.runner.start(
+                user_count=self.config.users,
+                spawn_rate=self.config.users,
+            )
+
+            # Run for specified duration
+            gevent.sleep(self._parse_time(self.config.run_time))
+
+            # Stop the test
+            env.runner.quit()
+            entries = collect_metrics(env)
+
+            # Wait for greenlets
+            env.runner.greenlet.join()
+            stats_printer_greenlet.kill()
+            stats_history_greenlet.kill()
+
+            # Print final stats
+            print_stats(env.stats)
+
+            return entries
+
+        except Exception as e:
+            logging.error(f"Test failed: {str(e)}")
+            raise
+
+    @staticmethod
+    def _parse_time(time_str: str) -> int:
+        """Convert time string (e.g., '30s', '1m', '1h') to seconds"""
+        unit = time_str[-1]
+        value = int(time_str[:-1])
+        if unit == "s":
+            return value
+        elif unit == "m":
+            return value * 60
+        elif unit == "h":
+            return value * 3600
+        else:
+            raise ValueError(f"Invalid time unit: {unit}")

release/llm_tests/kv_router_test/llm_4xl4.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g6.12xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/llm_tests/kv_router_test/test_kv_event_ingestion.py

@@ -0,0 +1,656 @@
+import sys
+from typing import Dict, List
+
+import pytest
+from vllm.distributed.kv_events import (
+    BlockRemoved,
+    BlockStored,
+    KVEventBatch,
+    ZmqEventPublisher,
+)
+
+import ray
+from ray import cloudpickle
+from ray._common.test_utils import async_wait_for_condition
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
+    _MODEL_NAME,
+    _TENANT_ID,
+    KVRouterActor,
+    get_worker_id,
+)
+from ray.serve._private.common import (
+    DeploymentID,
+    DeploymentTargetInfo,
+    ReplicaID,
+    RunningReplicaInfo,
+)
+
+# Actors defined in a test module are pickled by reference, so a worker on
+# another node would have to import this module to load them, which fails.
+# Pickle it by value so the actor classes ship in full and run on any node.
+cloudpickle.register_pickle_by_value(sys.modules[__name__])
+
+BLOCK_SIZE = 16
+MAX_NUM_BATCHED_TOKENS = 8192
+
+
+@pytest.fixture(scope="module")
+def ray_instance():
+    if not ray.is_initialized():
+        ray.init(address="auto")
+    yield
+
+
+class _TestKVRouterActor(KVRouterActor):
+    """KVRouterActor augmented with test-only introspection."""
+
+    async def get_candidate_worker_ids(self) -> List[int]:
+        """(Test only) The workers currently tracked from running replicas.
+
+        Async so it runs on the actor's event loop, serialized with
+        ``_on_deployment_targets`` which mutates the same map on that loop.
+        """
+        return sorted(self._replica_id_by_worker)
+
+    def get_registered_worker_ids(self) -> List[int]:
+        """(Test only) Worker ids the selection service can currently schedule."""
+        if self._svc is None:
+            return []
+        workers = self._svc.list_workers(model_name=_MODEL_NAME, tenant_id=_TENANT_ID)
+        return sorted(
+            w["worker_id"] for w in workers if w["lifecycle"] == "schedulable"
+        )
+
+    async def get_kv_overlap_blocks(self, token_ids: List[int]) -> Dict[int, int]:
+        """(Test only) Per-worker device-tier KV overlap blocks for a token sequence.
+
+        Returns the number of leading blocks of ``token_ids`` each worker has cached.
+        """
+        if self._svc is None:
+            return {}
+        scores = await self._svc.overlap_scores(
+            {
+                "model_name": _MODEL_NAME,
+                "tenant_id": _TENANT_ID,
+                "token_ids": list(token_ids),
+            }
+        )
+        return {w["worker_id"]: w["device_blocks"] for w in scores["workers"]}
+
+    async def get_potential_loads(self, token_ids: List[int]) -> Dict[int, dict]:
+        """(Test only) Per-worker projected load for routing ``token_ids``."""
+        if self._svc is None:
+            return {}
+        loads = await self._svc.potential_loads(
+            {
+                "model_name": _MODEL_NAME,
+                "tenant_id": _TENANT_ID,
+                "token_ids": list(token_ids),
+            }
+        )
+        return {load["worker_id"]: load for load in loads}
+
+    def get_loads(self) -> Dict[int, dict]:
+        """(Test only) Current per-worker active load."""
+        if self._svc is None:
+            return {}
+        model_loads = self._svc.loads(model_name=_MODEL_NAME, tenant_id=_TENANT_ID)
+        if not model_loads:
+            return {}
+        return {load["worker_id"]: load for load in model_loads[0]["loads"]}
+
+
+@ray.remote(num_cpus=0)
+class LocalKVRouterActor(_TestKVRouterActor):
+    """The real KVRouterActor with Serve LongPoll tracking disabled as there is
+    no Serve controller in these tests; tests drive _on_deployment_targets
+    directly with synthetic replica snapshots."""
+
+    def _start_replica_tracking(self) -> None:
+        pass
+
+
+@ray.remote(num_cpus=0)
+class FakeReplica:
+    """Mocks a real LLM replica's engine KV-event emission.
+
+    It runs no engine; instead it publishes synthetic KV events through vLLM's
+    real ``ZmqEventPublisher`` -- the exact path and wire format a real engine
+    uses to emit KV events over ZMQ -- so the selection service's connect-out
+    listener ingests them identically to production.
+    """
+
+    def __init__(self, port: int):
+        self._port = port
+        self._replay_port = port + 1000
+        self._pub = ZmqEventPublisher(
+            data_parallel_rank=0,
+            endpoint=f"tcp://*:{port}",
+            replay_endpoint=f"tcp://*:{self._replay_port}",
+            topic="",
+        )
+
+    def endpoint(self) -> str:
+        return f"tcp://{ray.util.get_node_ip_address()}:{self._port}"
+
+    def replay_endpoint(self) -> str:
+        return f"tcp://{ray.util.get_node_ip_address()}:{self._replay_port}"
+
+    def publish_stored(self, block_hashes, token_ids) -> None:
+        self._pub.publish(
+            KVEventBatch(
+                ts=1.0,
+                events=[
+                    BlockStored(
+                        block_hashes=list(block_hashes),
+                        parent_block_hash=None,
+                        token_ids=list(token_ids),
+                        block_size=BLOCK_SIZE,
+                        lora_id=None,
+                        medium="GPU",
+                        lora_name=None,
+                    )
+                ],
+            )
+        )
+
+    def publish_removed(self, block_hashes) -> None:
+        self._pub.publish(
+            KVEventBatch(
+                ts=2.0,
+                events=[BlockRemoved(block_hashes=list(block_hashes), medium="GPU")],
+            )
+        )
+
+    def close(self) -> None:
+        self._pub.shutdown()
+
+
+def running_replica(
+    unique_id: str, endpoint: str, replay_endpoint: str, dp_rank: int = 0
+):
+    """A RunningReplicaInfo whose routing_stats advertise a KV-events endpoint,
+    exactly as a real replica's record_routing_stats would surface it."""
+    return RunningReplicaInfo(
+        replica_id=ReplicaID(
+            unique_id=unique_id,
+            deployment_id=DeploymentID(name="llm", app_name="app"),
+        ),
+        node_id="node-1",
+        node_ip="127.0.0.1",
+        availability_zone=None,
+        actor_name=f"actor-{unique_id}",
+        max_ongoing_requests=10,
+        routing_stats={
+            "kv_event_metadata": {
+                "endpoint": endpoint,
+                "block_size": BLOCK_SIZE,
+                "max_num_batched_tokens": MAX_NUM_BATCHED_TOKENS,
+                "dp_rank": dp_rank,
+                "replay_endpoint": replay_endpoint,
+            }
+        },
+    )
+
+
+def targets(*replicas):
+    return DeploymentTargetInfo(is_available=True, running_replicas=list(replicas))
+
+
+async def wait_registered(actor, worker_ids):
+    await async_wait_for_condition(
+        lambda: ray.get(actor.get_registered_worker_ids.remote()) == sorted(worker_ids),
+        timeout=15,
+    )
+
+
+async def wait_for_overlap(actor, token_ids, predicate, publish=None, timeout=30):
+    """Poll the actor's overlap view until ``predicate`` holds."""
+
+    async def condition():
+        if publish is not None:
+            publish()
+        try:
+            overlap = await actor.get_kv_overlap_blocks.remote(list(token_ids))
+        except Exception:
+            return False
+        return predicate(overlap)
+
+    await async_wait_for_condition(condition, timeout=timeout, retry_interval_ms=500)
+
+
+async def book_uncached_load(actor, worker_id, count, base):
+    """Book ``count`` reservations of distinct, uncached tokens onto ``worker_id``
+    so each adds full prefill load to it; returns their reservation ids."""
+    reservation_ids = []
+    for i in range(count):
+        reservation_id = f"load-{base}-{i}"
+        start = base + i * 4 * BLOCK_SIZE
+        token_ids = list(range(start, start + 4 * BLOCK_SIZE))
+        selection = await actor.select_worker.remote(
+            reservation_id, token_ids, [worker_id]
+        )
+        assert selection["worker_id"] == worker_id
+        assert selection["effective_prefill_tokens"] == len(token_ids)
+        await actor.on_request_added.remote(
+            reservation_id,
+            worker_id,
+            token_ids,
+        )
+        reservation_ids.append(reservation_id)
+    return reservation_ids
+
+
+async def book_decode_load(
+    actor,
+    worker_id,
+    reservation_id,
+    token_ids,
+    output_tokens,
+    expected_output_tokens=None,
+):
+    """Book a request, move it to decode, and add output blocks."""
+    selection = await actor.select_worker.remote(reservation_id, token_ids, [worker_id])
+    assert selection["worker_id"] == worker_id
+    await actor.on_request_added.remote(
+        reservation_id,
+        worker_id,
+        list(token_ids),
+        expected_output_tokens=expected_output_tokens,
+    )
+    await actor.on_prefill_complete.remote(reservation_id)
+    await actor.on_decode_progress.remote(reservation_id, output_tokens)
+    return reservation_id
+
+
+class TestKvEventIngestion:
+    """Validates that a replica advertises its vLLM ZMQ PUB endpoint on the
+    LongPoll snapshot -> the KVRouterActor registers it -> the selection
+    service's connect-out listener -> the global KV indexer."""
+
+    @pytest.mark.asyncio
+    async def test_register_then_ingest_events(self, ray_instance):
+        """A replica appearing on the snapshot with a KV-events endpoint makes the
+        selection service dial it and index the KV events it publishes."""
+        actor = LocalKVRouterActor.remote()
+        replica = FakeReplica.remote(23901)
+        try:
+            worker_id = get_worker_id("replica-A")
+            endpoint = await replica.endpoint.remote()
+            replay = await replica.replay_endpoint.remote()
+            # The LongPoll snapshot now carries this replica's KV-events endpoint.
+            await actor._on_deployment_targets.remote(
+                targets(running_replica("replica-A", endpoint, replay))
+            )
+            await wait_registered(actor, [worker_id])
+
+            # Two full blocks of prompt; once ingested both overlap the query.
+            token_ids = list(range(2 * BLOCK_SIZE))
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_id) == 2,
+                publish=lambda: replica.publish_stored.remote([101, 102], token_ids),
+            )
+
+            # Removing the leaf block drops it from the worker's overlap.
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_id) == 1,
+                publish=lambda: replica.publish_removed.remote([102]),
+            )
+        finally:
+            await replica.close.remote()
+            for a in (replica, actor):
+                ray.kill(a, no_restart=True)
+
+    @pytest.mark.asyncio
+    async def test_per_worker_isolation(self, ray_instance):
+        """Each worker's overlap reflects only the blocks its own replica cached."""
+        actor = LocalKVRouterActor.remote()
+        a = FakeReplica.remote(23902)
+        b = FakeReplica.remote(23903)
+        workers = {
+            "replica-A": get_worker_id("replica-A"),
+            "replica-B": get_worker_id("replica-B"),
+        }
+        try:
+            # Both replicas advertise their endpoints in a single snapshot.
+            await actor._on_deployment_targets.remote(
+                targets(
+                    running_replica(
+                        "replica-A",
+                        await a.endpoint.remote(),
+                        await a.replay_endpoint.remote(),
+                    ),
+                    running_replica(
+                        "replica-B",
+                        await b.endpoint.remote(),
+                        await b.replay_endpoint.remote(),
+                    ),
+                )
+            )
+            await wait_registered(actor, list(workers.values()))
+
+            tokens_a = list(range(BLOCK_SIZE))
+            tokens_b = list(range(BLOCK_SIZE, 2 * BLOCK_SIZE))
+
+            await wait_for_overlap(
+                actor,
+                tokens_a,
+                lambda overlap: overlap.get(workers["replica-A"]) == 1,
+                publish=lambda: a.publish_stored.remote([201], tokens_a),
+            )
+            await wait_for_overlap(
+                actor,
+                tokens_b,
+                lambda overlap: overlap.get(workers["replica-B"]) == 1,
+                publish=lambda: b.publish_stored.remote([301], tokens_b),
+            )
+            # A's content does not overlap B's worker and vice versa.
+            overlap_a = await actor.get_kv_overlap_blocks.remote(tokens_a)
+            assert overlap_a.get(workers["replica-B"], 0) == 0
+        finally:
+            for standin in (a, b):
+                await standin.close.remote()
+            for handle in (a, b, actor):
+                ray.kill(handle, no_restart=True)
+
+    @pytest.mark.asyncio
+    async def test_departed_replica_is_evicted(self, ray_instance):
+        """A replica dropping off the snapshot tears down its listener and drops
+        it from the selection service catalog."""
+        actor = LocalKVRouterActor.remote()
+        replica = FakeReplica.remote(23904)
+        worker_id = get_worker_id("replica-A")
+        try:
+            endpoint = await replica.endpoint.remote()
+            replay = await replica.replay_endpoint.remote()
+            await actor._on_deployment_targets.remote(
+                targets(running_replica("replica-A", endpoint, replay))
+            )
+            token_ids = list(range(2 * BLOCK_SIZE))
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_id) == 2,
+                publish=lambda: replica.publish_stored.remote([401, 402], token_ids),
+            )
+
+            # The replica departs: an empty snapshot evicts its worker.
+            await actor._on_deployment_targets.remote(targets())
+            await wait_registered(actor, [])
+            assert await actor.get_candidate_worker_ids.remote() == []
+        finally:
+            await replica.close.remote()
+            for a in (replica, actor):
+                ray.kill(a, no_restart=True)
+
+    @pytest.mark.asyncio
+    async def test_select_worker_prefers_overlap(self, ray_instance):
+        """select_worker scores via the selection service and picks the worker
+        whose cached blocks overlap the prompt, within the allowed candidate set."""
+        actor = LocalKVRouterActor.remote()
+        a = FakeReplica.remote(23905)
+        b = FakeReplica.remote(23906)
+        worker_a = get_worker_id("replica-A")
+        worker_b = get_worker_id("replica-B")
+        try:
+            await actor._on_deployment_targets.remote(
+                targets(
+                    running_replica(
+                        "replica-A",
+                        await a.endpoint.remote(),
+                        await a.replay_endpoint.remote(),
+                    ),
+                    running_replica(
+                        "replica-B",
+                        await b.endpoint.remote(),
+                        await b.replay_endpoint.remote(),
+                    ),
+                )
+            )
+            await wait_registered(actor, [worker_a, worker_b])
+
+            # A caches the prompt's two blocks; B caches nothing.
+            token_ids = list(range(2 * BLOCK_SIZE))
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_a) == 2,
+                publish=lambda: a.publish_stored.remote([501, 502], token_ids),
+            )
+
+            selection = await actor.select_worker.remote(
+                "req-1", token_ids, [worker_a, worker_b]
+            )
+            assert selection["worker_id"] == worker_a
+            assert selection["overlap_tokens"] >= BLOCK_SIZE
+        finally:
+            await a.close.remote()
+            await b.close.remote()
+            for x in (a, b, actor):
+                ray.kill(x, no_restart=True)
+
+    @pytest.mark.asyncio
+    async def test_select_worker_prefers_lower_load(self, ray_instance):
+        """With equal KV overlap on both workers, select_worker routes to the
+        worker carrying less active load. Booking heavy load onto one worker
+        sends the next equally-overlapping request to the other."""
+        actor = LocalKVRouterActor.remote()
+        a = FakeReplica.remote(23907)
+        b = FakeReplica.remote(23908)
+        worker_a = get_worker_id("replica-A")
+        worker_b = get_worker_id("replica-B")
+        try:
+            await actor._on_deployment_targets.remote(
+                targets(
+                    running_replica(
+                        "replica-A",
+                        await a.endpoint.remote(),
+                        await a.replay_endpoint.remote(),
+                    ),
+                    running_replica(
+                        "replica-B",
+                        await b.endpoint.remote(),
+                        await b.replay_endpoint.remote(),
+                    ),
+                )
+            )
+            await wait_registered(actor, [worker_a, worker_b])
+
+            # Both replicas cache the same prompt blocks: overlap is equal, so the
+            # selection turns purely on each worker's active load.
+            token_ids = list(range(2 * BLOCK_SIZE))
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_a) == 2,
+                publish=lambda: a.publish_stored.remote([601, 602], token_ids),
+            )
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_b) == 2,
+                publish=lambda: b.publish_stored.remote([601, 602], token_ids),
+            )
+
+            # Heavy active load on A (uncached prompts -> full prefill load): the
+            # equally-overlapping request is routed to the idle worker B.
+            load_a = await book_uncached_load(actor, worker_a, count=8, base=100_000)
+            loads = await actor.get_potential_loads.remote(token_ids)
+            assert (
+                loads[worker_a]["potential_prefill_tokens"]
+                > loads[worker_b]["potential_prefill_tokens"]
+            )
+            selection = await actor.select_worker.remote(
+                "probe-1", token_ids, [worker_a, worker_b]
+            )
+            assert selection["worker_id"] == worker_b
+
+            # Free A and load B instead: the choice flips to A.
+            for reservation_id in load_a:
+                await actor.on_request_completed.remote(reservation_id)
+            await book_uncached_load(actor, worker_b, count=8, base=500_000)
+            loads = await actor.get_potential_loads.remote(token_ids)
+            assert (
+                loads[worker_b]["potential_prefill_tokens"]
+                > loads[worker_a]["potential_prefill_tokens"]
+            )
+            selection = await actor.select_worker.remote(
+                "probe-2", token_ids, [worker_a, worker_b]
+            )
+            assert selection["worker_id"] == worker_a
+        finally:
+            await a.close.remote()
+            await b.close.remote()
+            for x in (a, b, actor):
+                ray.kill(x, no_restart=True)
+
+    @pytest.mark.asyncio
+    async def test_select_worker_prefers_lower_decode_load(self, ray_instance):
+        """With equal KV overlap and no active prefill, a long decode on one
+        worker makes the selector choose the lower decode-load worker."""
+        actor = LocalKVRouterActor.remote()
+        a = FakeReplica.remote(23909)
+        b = FakeReplica.remote(23910)
+        worker_a = get_worker_id("replica-A")
+        worker_b = get_worker_id("replica-B")
+        try:
+            await actor._on_deployment_targets.remote(
+                targets(
+                    running_replica(
+                        "replica-A",
+                        await a.endpoint.remote(),
+                        await a.replay_endpoint.remote(),
+                    ),
+                    running_replica(
+                        "replica-B",
+                        await b.endpoint.remote(),
+                        await b.replay_endpoint.remote(),
+                    ),
+                )
+            )
+            await wait_registered(actor, [worker_a, worker_b])
+
+            token_ids = list(range(2 * BLOCK_SIZE))
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_a) == 2,
+                publish=lambda: a.publish_stored.remote([701, 702], token_ids),
+            )
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_b) == 2,
+                publish=lambda: b.publish_stored.remote([701, 702], token_ids),
+            )
+
+            await book_decode_load(
+                actor,
+                worker_a,
+                "decode-load-a",
+                token_ids,
+                output_tokens=8 * BLOCK_SIZE,
+            )
+
+            loads = await actor.get_potential_loads.remote(token_ids)
+            assert (
+                loads[worker_a]["potential_prefill_tokens"]
+                == loads[worker_b]["potential_prefill_tokens"]
+            )
+            assert (
+                loads[worker_a]["potential_decode_blocks"]
+                > loads[worker_b]["potential_decode_blocks"]
+            )
+
+            selection = await actor.select_worker.remote(
+                "probe-decode", token_ids, [worker_a, worker_b]
+            )
+            assert selection["worker_id"] == worker_b
+        finally:
+            await a.close.remote()
+            await b.close.remote()
+            for x in (a, b, actor):
+                ray.kill(x, no_restart=True)
+
+    @pytest.mark.asyncio
+    async def test_decode_load_uses_expected_output_decay(self, ray_instance):
+        """Expected output length decays decode load as generation nears the
+        estimate, so the selector prefers the lower remaining-load worker."""
+        actor = LocalKVRouterActor.remote()
+        a = FakeReplica.remote(23911)
+        b = FakeReplica.remote(23912)
+        worker_a = get_worker_id("replica-A")
+        worker_b = get_worker_id("replica-B")
+        try:
+            await actor._on_deployment_targets.remote(
+                targets(
+                    running_replica(
+                        "replica-A",
+                        await a.endpoint.remote(),
+                        await a.replay_endpoint.remote(),
+                    ),
+                    running_replica(
+                        "replica-B",
+                        await b.endpoint.remote(),
+                        await b.replay_endpoint.remote(),
+                    ),
+                )
+            )
+            await wait_registered(actor, [worker_a, worker_b])
+
+            token_ids = list(range(2 * BLOCK_SIZE))
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_a) == 2,
+                publish=lambda: a.publish_stored.remote([801, 802], token_ids),
+            )
+            await wait_for_overlap(
+                actor,
+                token_ids,
+                lambda overlap: overlap.get(worker_b) == 2,
+                publish=lambda: b.publish_stored.remote([801, 802], token_ids),
+            )
+
+            output_tokens = 4 * BLOCK_SIZE - 1
+            await book_decode_load(
+                actor,
+                worker_a,
+                "decayed-decode-a",
+                token_ids,
+                output_tokens=output_tokens,
+                expected_output_tokens=4 * BLOCK_SIZE,
+            )
+            await book_decode_load(
+                actor,
+                worker_b,
+                "full-decode-b",
+                token_ids,
+                output_tokens=output_tokens,
+            )
+
+            active_loads = await actor.get_loads.remote()
+            assert active_loads[worker_a]["potential_prefill_tokens"] == 0
+            assert active_loads[worker_b]["potential_prefill_tokens"] == 0
+            assert (
+                active_loads[worker_a]["potential_decode_blocks"]
+                < active_loads[worker_b]["potential_decode_blocks"]
+            )
+
+            selection = await actor.select_worker.remote(
+                "probe-decay", token_ids, [worker_a, worker_b]
+            )
+            assert selection["worker_id"] == worker_a
+        finally:
+            await a.close.remote()
+            await b.close.remote()
+            for x in (a, b, actor):
+                ray.kill(x, no_restart=True)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

release/llm_tests/kv_router_test/test_kv_events.py

@@ -0,0 +1,436 @@
+import asyncio
+import sys
+from typing import Dict, List
+from unittest import mock
+
+import pytest
+import requests
+from dynamo.llm import compute_block_hash_for_seq
+from transformers import AutoTokenizer
+
+import ray
+from ray import serve
+from ray._common.test_utils import async_wait_for_condition
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
+    _MODEL_NAME,
+    _TENANT_ID,
+    KV_ROUTER_ACTOR_NAME,
+    KVRouterActor,
+)
+from ray.serve.config import RequestRouterConfig
+from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_openai_app
+from ray.serve.llm.request_router import KVAwareRouter
+
+from utils import _TestKVAwareRouter, discover_deployment_actor
+
+MODEL_ID = "qwen3-0.6b"
+MODEL_SOURCE = "Qwen/Qwen3-0.6B"
+APP_NAME = "kv_events_gpu_test"
+NUM_REPLICAS = 2
+BLOCK_SIZE = 16
+MAX_TOKENS = 50
+
+# MESSAGES and FLUSH_MESSAGES share a long prefix (two full 16-token blocks) so
+# the reset test asserts a *partial* overlap fallback: after one replica's prefix
+# cache is cleared and re-warmed with FLUSH_MESSAGES, its overlap for MESSAGES
+# drops to just the shared prefix blocks while the untouched replica keeps the
+# full prompt.
+_SHARED_PREFIX = (
+    "Repeat the following sentence exactly five times in a row, word for word, "
+    "without adding anything else at all: the quick brown fox jumps over the lazy dog"
+)
+MESSAGES = [
+    {
+        "role": "user",
+        "content": (
+            _SHARED_PREFIX
+            + " near the calm river bank today under a wide clear evening sky "
+            "over the hills."
+        ),
+    }
+]
+FLUSH_MESSAGES = [
+    {"role": "user", "content": _SHARED_PREFIX + " beside the tall fence."}
+]
+
+
+class _TestKVRouterActor(KVRouterActor):
+    """KVRouterActor augmented with test-only introspection."""
+
+    async def get_candidate_worker_ids(self) -> List[int]:
+        """(Test only) The workers currently tracked from running replicas.
+
+        Async so it runs on the actor's event loop, serialized with
+        ``_on_deployment_targets`` which mutates the same map on that loop.
+        """
+        return sorted(self._replica_id_by_worker)
+
+    def get_kv_event_worker_replicas(self) -> Dict[int, str]:
+        """(Test only) The registered Dynamo worker id -> replica full id mapping."""
+        return dict(self._replica_id_by_worker)
+
+    def get_registered_worker_ids(self) -> List[int]:
+        """(Test only) Worker ids the selection service can currently schedule."""
+        if self._svc is None:
+            return []
+        workers = self._svc.list_workers(model_name=_MODEL_NAME, tenant_id=_TENANT_ID)
+        return sorted(
+            w["worker_id"] for w in workers if w["lifecycle"] == "schedulable"
+        )
+
+    async def get_kv_overlap_blocks(self, token_ids: List[int]) -> Dict[int, int]:
+        """(Test only) Per-worker device-tier KV overlap blocks for a token sequence.
+
+        Returns the number of leading blocks of ``token_ids`` each worker has cached.
+        """
+        if self._svc is None:
+            return {}
+        scores = await self._svc.overlap_scores(
+            {
+                "model_name": _MODEL_NAME,
+                "tenant_id": _TENANT_ID,
+                "token_ids": list(token_ids),
+            }
+        )
+        return {w["worker_id"]: w["device_blocks"] for w in scores["workers"]}
+
+
+def post_chat(endpoint, messages=MESSAGES, max_tokens=MAX_TOKENS):
+    host, port = endpoint
+    response = requests.post(
+        f"http://{host}:{port}/v1/chat/completions",
+        json={
+            "model": MODEL_ID,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": 0.0,
+            "ignore_eos": True,
+        },
+        timeout=120,
+    )
+    assert response.status_code == 200, response.text
+    return response.json()
+
+
+def tokenize_prompt(endpoint, messages=MESSAGES):
+    """The engine's exact token ids for a chat-templated prompt."""
+    host, port = endpoint
+    response = requests.post(
+        f"http://{host}:{port}/tokenize",
+        json={"model": MODEL_ID, "messages": messages, "add_generation_prompt": True},
+        timeout=60,
+    )
+    assert response.status_code == 200, response.text
+    return response.json()["tokens"]
+
+
+def tokenize_text(endpoint, prompt):
+    """Token ids for a raw prompt string via the completion /tokenize path.
+
+    add_special_tokens is False because a chat-templated string already carries
+    the template's special tokens as text.
+    """
+    host, port = endpoint
+    response = requests.post(
+        f"http://{host}:{port}/tokenize",
+        json={"model": MODEL_ID, "prompt": prompt, "add_special_tokens": False},
+        timeout=60,
+    )
+    assert response.status_code == 200, response.text
+    return response.json()["tokens"]
+
+
+def num_prompt_blocks(token_ids):
+    """Number of full KV blocks in a token sequence."""
+    return len(compute_block_hash_for_seq(list(token_ids), BLOCK_SIZE))
+
+
+class TestKvEvents:
+    @pytest.fixture(scope="class")
+    def deployed_handle(self):
+        """Deploy two direct-streaming LLMServer replicas with KV events on."""
+        if not ray.is_initialized():
+            ray.init(address="auto")
+        serve.shutdown()
+
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id=MODEL_ID,
+                model_source=MODEL_SOURCE,
+            ),
+            deployment_config=dict(
+                autoscaling_config=dict(
+                    min_replicas=NUM_REPLICAS, max_replicas=NUM_REPLICAS
+                ),
+                # This test validates the KV-events plane (engine events ->
+                # selection service indexer), not routing: requests are sent
+                # directly to each replica's endpoint. KVAwareRouter's own routing
+                # is unimplemented, so this subclass borrows RoundRobinRouter's
+                # selection purely so replica discovery can enumerate both.
+                # TODO (jeffreywang): use KVAwareRouter directly once it
+                # implements routing.
+                request_router_config=RequestRouterConfig(
+                    request_router_class=_TestKVAwareRouter
+                ),
+            ),
+            engine_kwargs=dict(
+                max_model_len=2048,
+                enforce_eager=True,
+                gpu_memory_utilization=0.4,  # small model on a shared GPU
+            ),
+            experimental_configs={"KV_EVENTS_PORT_BASE": 21557},
+            runtime_env=dict(
+                env_vars={
+                    "RAY_SERVE_ENABLE_DIRECT_INGRESS": "1",
+                    "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING": "1",
+                    # /reset_prefix_cache is a vLLM dev-mode endpoint.
+                    "VLLM_SERVER_DEV_MODE": "1",
+                },
+            ),
+            log_engine_metrics=False,
+        )
+        # The builder auto-attaches the KVRouterActor; patch in the test
+        # subclass so the deployment-scoped actor exposes test introspection.
+        with mock.patch(
+            "ray.llm._internal.serve.routing_policies.kv_aware.utils.KVRouterActor",
+            _TestKVRouterActor,
+        ):
+            app = build_openai_app({"llm_configs": [llm_config]})
+            handle = serve.run(app, name=APP_NAME)
+        yield handle
+        serve.shutdown()
+
+    async def _discover_replicas(self, handle):
+        """Map each replica's full id to its backend HTTP endpoint."""
+        endpoints = {}
+        for _ in range(100):
+            async with handle.choose_replica() as selection:
+                replica = selection._replica
+                if replica.backend_http_endpoint is not None:
+                    replica_id = replica.replica_id.to_full_id_str()
+                    endpoints[replica_id] = replica.backend_http_endpoint
+            if len(endpoints) == NUM_REPLICAS:
+                return endpoints
+            await asyncio.sleep(0.5)
+        raise AssertionError(
+            f"Expected {NUM_REPLICAS} replicas with backend endpoints, "
+            f"found {len(endpoints)}."
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(600)
+    async def test_kv_events_reach_selection_service(self, deployed_handle):
+        """Each replica's real engine KV events reach the selection service via
+        its connect-out listener, a per-worker prefix-cache reset is observed as
+        reduced overlap, and scoring routes the prompt to the higher-overlap
+        worker."""
+        actor = discover_deployment_actor(
+            APP_NAME, deployed_handle.deployment_name, KV_ROUTER_ACTOR_NAME
+        )
+        assert actor is not None, "KV router actor was not discoverable"
+
+        replica_endpoints = await self._discover_replicas(deployed_handle)
+
+        # Each replica advertises its KV-events endpoint via record_routing_stats;
+        # the controller propagates it on the LongPoll replica snapshot and the
+        # actor registers the worker with the selection service. Wait for every
+        # replica to be registered (the controller polls routing stats on an
+        # interval, so this is not synchronous with replica startup).
+        async def all_replicas_registered():
+            replica_by_worker = await actor.get_kv_event_worker_replicas.remote()
+            return sorted(replica_by_worker.values()) == sorted(replica_endpoints)
+
+        await async_wait_for_condition(all_replicas_registered, timeout=90)
+
+        replica_by_worker = await actor.get_kv_event_worker_replicas.remote()
+        endpoints = {
+            worker_id: replica_endpoints[replica_id]
+            for worker_id, replica_id in replica_by_worker.items()
+        }
+        worker_ids = sorted(endpoints)
+        assert await actor.get_candidate_worker_ids.remote() == worker_ids
+        assert await actor.get_registered_worker_ids.remote() == worker_ids
+
+        # The same prompt on each replica caches the same content.
+        usages = {}
+        for worker_id in worker_ids:
+            usages[worker_id] = post_chat(endpoints[worker_id])["usage"]
+
+        prompt_token_ids = tokenize_prompt(endpoints[worker_ids[0]])
+        prompt_blocks = num_prompt_blocks(prompt_token_ids)
+        assert prompt_blocks >= 2
+
+        # The engines' KV events reached the indexer: full prompt overlap is
+        # scored on both workers.
+        async def both_workers_fully_overlap():
+            overlaps = await actor.get_kv_overlap_blocks.remote(prompt_token_ids)
+            return all(overlaps.get(w) == prompt_blocks for w in worker_ids)
+
+        await async_wait_for_condition(both_workers_fully_overlap, timeout=60)
+
+        for worker_id in worker_ids:
+            usage = usages[worker_id]
+            assert usage["prompt_tokens"] == len(prompt_token_ids)
+            assert usage["completion_tokens"] == MAX_TOKENS
+
+        # /reset_prefix_cache clears only this worker's view; the engine drains
+        # queued KV events on scheduler steps, so a small follow-up request
+        # flushes the AllBlocksCleared event to the listener.
+        reset_worker, untouched_worker = worker_ids
+        host, port = endpoints[reset_worker]
+        response = requests.post(f"http://{host}:{port}/reset_prefix_cache", timeout=60)
+        assert response.status_code == 200, response.text
+        post_chat(endpoints[reset_worker], messages=FLUSH_MESSAGES, max_tokens=2)
+
+        # The reset worker's overlap falls back to the chat-template prefix the
+        # two prompts share; the untouched worker keeps the full prompt.
+        flush_token_ids = tokenize_prompt(endpoints[reset_worker], FLUSH_MESSAGES)
+        diverge = next(
+            (
+                i
+                for i, (a, b) in enumerate(zip(prompt_token_ids, flush_token_ids))
+                if a != b
+            ),
+            min(len(prompt_token_ids), len(flush_token_ids)),
+        )
+        shared_blocks = diverge // BLOCK_SIZE
+
+        async def reset_worker_cleared():
+            overlaps = await actor.get_kv_overlap_blocks.remote(prompt_token_ids)
+            return overlaps.get(reset_worker, 0) == shared_blocks
+
+        await async_wait_for_condition(reset_worker_cleared, timeout=60)
+        overlaps = await actor.get_kv_overlap_blocks.remote(prompt_token_ids)
+        assert overlaps.get(untouched_worker) == prompt_blocks
+
+        # Scoring routes the prompt to the worker holding more cached overlap.
+        selection = await actor.select_worker.remote(
+            "score-req", prompt_token_ids, worker_ids
+        )
+        assert selection["worker_id"] == untouched_worker
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(600)
+    async def test_chat_tokens_match_prefill(self, deployed_handle):
+        """Ensure chat template is applied: a chat request scores the same overlap as
+        the prompt rendered with the model's chat template and tokenized as raw text."""
+        actor = discover_deployment_actor(
+            APP_NAME, deployed_handle.deployment_name, KV_ROUTER_ACTOR_NAME
+        )
+        assert actor is not None
+        replica_endpoints = await self._discover_replicas(deployed_handle)
+
+        async def all_registered():
+            registered = await actor.get_kv_event_worker_replicas.remote()
+            return sorted(registered.values()) == sorted(replica_endpoints)
+
+        await async_wait_for_condition(all_registered, timeout=90)
+
+        # Ground truth: render the chat template client-side and tokenize as text.
+        worker_id, replica_id = next(
+            iter((await actor.get_kv_event_worker_replicas.remote()).items())
+        )
+        endpoint = replica_endpoints[replica_id]
+        manual_prompt = AutoTokenizer.from_pretrained(MODEL_SOURCE).apply_chat_template(
+            MESSAGES, add_generation_prompt=True, tokenize=False
+        )
+        manual_token_ids = tokenize_text(endpoint, manual_prompt)
+        prompt_blocks = num_prompt_blocks(manual_token_ids)
+        assert prompt_blocks >= 2
+
+        # Warm this worker's prefix cache with the chat request and wait until the
+        # indexer reflects the manually-templated prompt's blocks.
+        post_chat(endpoint)
+
+        async def manual_fully_overlaps():
+            overlaps = await actor.get_kv_overlap_blocks.remote(manual_token_ids)
+            return overlaps.get(worker_id) == prompt_blocks
+
+        await async_wait_for_condition(manual_fully_overlaps, timeout=60)
+
+        # The chat /tokenize tokens hit the same cached blocks -> same score,
+        # proving /tokenize applied the chat template.
+        chat_token_ids = tokenize_prompt(endpoint, MESSAGES)
+        chat_overlaps = await actor.get_kv_overlap_blocks.remote(chat_token_ids)
+        assert chat_overlaps.get(worker_id) == prompt_blocks
+        assert chat_token_ids == manual_token_ids
+
+
+class TestKvScoring:
+    """End-to-end KV-aware routing: a request routed through a deployed
+    KVAwareRouter is scored by the selection service and lands on a live
+    replica."""
+
+    @pytest.fixture(scope="class")
+    def kv_aware_handle(self):
+        """Deploy with KVAwareRouter; the build auto-attaches the KVRouterActor
+        and enables engine KV events."""
+        if not ray.is_initialized():
+            ray.init(address="auto")
+        serve.shutdown()
+
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id=MODEL_ID, model_source=MODEL_SOURCE
+            ),
+            deployment_config=dict(
+                autoscaling_config=dict(
+                    min_replicas=NUM_REPLICAS, max_replicas=NUM_REPLICAS
+                ),
+                request_router_config=RequestRouterConfig(
+                    request_router_class=KVAwareRouter
+                ),
+            ),
+            engine_kwargs=dict(
+                max_model_len=2048,
+                enforce_eager=True,
+                gpu_memory_utilization=0.4,
+            ),
+            experimental_configs={"KV_EVENTS_PORT_BASE": 21600},
+            runtime_env=dict(
+                env_vars={
+                    "RAY_SERVE_ENABLE_DIRECT_INGRESS": "1",
+                    "RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING": "1",
+                }
+            ),
+            log_engine_metrics=False,
+        )
+        app = build_openai_app({"llm_configs": [llm_config]})
+        handle = serve.run(app, name="kv_scoring_gpu_test")
+        yield handle
+        serve.shutdown()
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(600)
+    async def test_routes_to_higher_overlap_replica(self, kv_aware_handle):
+        """An overlapping prompt routes back to the replica that cached it,
+        scored through the full KVAwareRouter path."""
+        async with kv_aware_handle.choose_replica(
+            _reserve=False,
+            request_token_ids=[1],  # KV-aware routing requires token ids
+        ) as selection:
+            cached_id = selection._replica.replica_id.to_full_id_str()
+            cached_endpoint = selection._replica.backend_http_endpoint
+        post_chat(cached_endpoint)
+        prompt_token_ids = tokenize_prompt(cached_endpoint)
+        assert num_prompt_blocks(prompt_token_ids) >= 2
+
+        # Worker registration and KV-event indexing are asynchronous, so poll the
+        # scoring path until it converges on the replica holding the cached blocks.
+        async def routes_to_cached_replica():
+            picks = set()
+            for _ in range(3):
+                async with kv_aware_handle.choose_replica(
+                    _reserve=False,
+                    request_token_ids=prompt_token_ids,
+                ) as selection:
+                    picks.add(selection._replica.replica_id.to_full_id_str())
+            return picks == {cached_id}
+
+        await async_wait_for_condition(routes_to_cached_replica, timeout=120)
+
+
+if __name__ == "__main__":
+    if not ray.is_initialized():
+        ray.init(address="auto")
+    sys.exit(pytest.main(["-v", "-s", __file__]))

release/llm_tests/kv_router_test/test_request_lifecycle_tracking.py

@@ -0,0 +1,346 @@
+import asyncio
+import json
+import math
+import sys
+from dataclasses import asdict
+from unittest import mock
+
+import pytest
+import requests
+from dynamo.llm import compute_block_hash_for_seq
+
+import ray
+from ray import serve
+from ray._common.test_utils import async_wait_for_condition
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
+    _MODEL_NAME,
+    _TENANT_ID,
+    KV_ROUTER_ACTOR_NAME,
+    KVRouterActor,
+)
+from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
+    configure_kv_events_for_kv_routing,
+)
+from ray.serve.config import RequestRouterConfig
+from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_openai_app
+
+from utils import _TestKVAwareRouter, discover_deployment_actor
+
+MODEL_ID = "Qwen/Qwen3-0.6B"
+APP_NAME = "lifecycle_tracking_gpu_test"
+REQUEST_ID = "gpu-req-1"
+BLOCK_SIZE = 16
+MAX_TOKENS = 48
+# A multi-block prompt so the engine caches retrievable KV blocks.
+PROMPT_TEXT = (
+    "Repeat the following instruction carefully and then answer it in full "
+    "detail: describe the water cycle, including evaporation, condensation, "
+    "precipitation, and collection, giving a concrete real-world example for "
+    "each of the four stages and how they connect."
+)
+# Token tracking reports Serve's canonical request id, matching the id used for
+# routing, even if vLLM derives a separate engine-level id internally.
+LIFECYCLE_REQUEST_ID = REQUEST_ID
+
+
+class RecordingKVRouterActor(KVRouterActor):
+    """The real KVRouterActor, additionally recording every event it applies
+    and any error raised while booking it into the live selection service.
+
+    Eviction on completion makes the hook calls unobservable from state alone;
+    the event log preserves them, and the error log proves each hook's call into
+    the live service succeeded rather than being swallowed by the reporter pump.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._event_log = []
+        self._errors = []
+
+    async def on_lifecycle_events(self, events):
+        self._event_log.extend(events)
+        for hook_name, args in events:
+            try:
+                await getattr(self, hook_name)(*args)
+            except Exception as e:
+                self._errors.append((hook_name, repr(e)))
+
+    def get_event_log(self):
+        return self._event_log
+
+    def get_errors(self):
+        return self._errors
+
+    def get_registered_worker_ids(self):
+        """(Test only) Worker ids the selection service can currently schedule."""
+        workers = self._svc.list_workers(model_name=_MODEL_NAME, tenant_id=_TENANT_ID)
+        return sorted(
+            w["worker_id"] for w in workers if w["lifecycle"] == "schedulable"
+        )
+
+    async def get_worker_active_requests(self, worker_id):
+        """(Test only) In-flight requests the service tracks as active load on
+        ``worker_id`` -- the count scoring factors in."""
+        for model in self._svc.loads(model_name=_MODEL_NAME, tenant_id=_TENANT_ID):
+            for load in model["loads"]:
+                if load["worker_id"] == worker_id:
+                    return load["active_requests"]
+        return 0
+
+    async def get_request_lifecycle(self, request_id):
+        """(Test only) Snapshot of a request's local lifecycle state, or ``None``."""
+        state = self._requests.get(request_id)
+        if state is None:
+            return None
+        snapshot = asdict(state)
+        snapshot.pop("created_at", None)
+        return snapshot
+
+    async def get_active_request_ids(self):
+        """(Test only) Ids of the requests in the actor's in-flight view."""
+        return list(self._requests)
+
+    async def get_overlap_blocks(self, token_ids):
+        """(Test only) Per-worker device-tier KV overlap blocks for a sequence."""
+        scores = await self._svc.overlap_scores(
+            {
+                "model_name": _MODEL_NAME,
+                "tenant_id": _TENANT_ID,
+                "token_ids": list(token_ids),
+            }
+        )
+        return {w["worker_id"]: w["device_blocks"] for w in scores["workers"]}
+
+
+def num_prompt_blocks(token_ids):
+    """Number of full KV blocks in a token sequence (matches the indexer)."""
+    return len(compute_block_hash_for_seq(list(token_ids), BLOCK_SIZE))
+
+
+class TestLifecycleTracking:
+    @pytest.fixture(scope="class")
+    def deployed_handle(self):
+        """Deploy a direct-streaming LLMServer with KV events on and the
+        recording actor attached."""
+        if not ray.is_initialized():
+            ray.init(address="auto")
+        serve.shutdown()
+
+        engine_kwargs = dict(
+            max_model_len=2048,
+            enforce_eager=True,
+            gpu_memory_utilization=0.4,
+            use_tqdm_on_load=False,
+        )
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(
+                model_id=MODEL_ID,
+                model_source=MODEL_ID,
+            ),
+            deployment_config=dict(
+                autoscaling_config=dict(min_replicas=1, max_replicas=1),
+                # KVAwareRouter gates engine token tracking and the KV-events
+                # plane; the builder auto-attaches the KV router actor.
+                request_router_config=RequestRouterConfig(
+                    request_router_class=_TestKVAwareRouter
+                ),
+            ),
+            engine_kwargs=engine_kwargs,
+            placement_group_config={"bundles": [{"GPU": 1}]},
+            experimental_configs={"KV_EVENTS_PORT_BASE": 21600},
+            runtime_env=dict(env_vars={"VLLM_DISABLE_COMPILE_CACHE": "1"}),
+            log_engine_metrics=False,
+        )
+        # Emit engine KV-cache events so the actor registers the replica's worker
+        # (making it schedulable, required to book a reservation) and the service
+        # indexes the prompt's blocks.
+        configure_kv_events_for_kv_routing(llm_config)
+
+        # Patch in the recording subclass so the builder-attached, deployment-
+        # scoped KV router actor exposes the received lifecycle events.
+        with mock.patch(
+            "ray.llm._internal.serve.routing_policies.kv_aware.utils.KVRouterActor",
+            RecordingKVRouterActor,
+        ):
+            app = build_openai_app({"llm_configs": [llm_config]})
+            handle = serve.run(app, name=APP_NAME)
+        yield handle
+        serve.shutdown()
+
+    async def _backend_endpoint(self, handle):
+        """Poll for the replica's backend HTTP (host, port); it is reported
+        asynchronously after the replica starts its direct-ingress server."""
+        for _ in range(60):
+            async with handle.choose_replica() as selection:
+                endpoint = selection._replica.backend_http_endpoint
+            if endpoint is not None:
+                return endpoint
+            await asyncio.sleep(1.0)
+        raise AssertionError("replica backend HTTP endpoint never became available")
+
+    async def _registered_worker(self, actor):
+        """Wait for the replica's worker to register (schedulable) and return it."""
+        await async_wait_for_condition(
+            lambda: len(ray.get(actor.get_registered_worker_ids.remote())) == 1,
+            timeout=90,
+            retry_interval_ms=1000,
+        )
+        return (await actor.get_registered_worker_ids.remote())[0]
+
+    @pytest.mark.asyncio
+    async def test_exact_lifecycle_tracking(self, deployed_handle):
+        actor = discover_deployment_actor(
+            APP_NAME, deployed_handle.deployment_name, KV_ROUTER_ACTOR_NAME
+        )
+        assert actor is not None, "KV router actor was not discoverable"
+        host, port = await self._backend_endpoint(deployed_handle)
+        # The replica's worker must register before a request can book against it.
+        worker_id = await self._registered_worker(actor)
+        assert await actor.get_worker_active_requests.remote(worker_id) == 0
+
+        # X-Request-Id pins the engine request id; include_usage returns the
+        # engine's own token counts as ground truth.
+        url = f"http://{host}:{port}/v1/chat/completions"
+        payload = {
+            "model": MODEL_ID,
+            "messages": [{"role": "user", "content": PROMPT_TEXT}],
+            "stream": True,
+            "stream_options": {"include_usage": True},
+            "max_tokens": MAX_TOKENS,
+            "temperature": 0.0,
+        }
+        headers = {"X-Request-Id": REQUEST_ID}
+
+        # Snapshot the actor's view and the worker's tracked load after every
+        # streamed chunk: completion evicts the request, so its in-flight state
+        # is only observable while streaming.
+        usage = None
+        snapshots = []
+        live_active_requests = []
+        with requests.post(
+            url, json=payload, headers=headers, stream=True, timeout=120
+        ) as resp:
+            assert resp.status_code == 200, resp.text
+            for raw in resp.iter_lines():
+                if not raw:
+                    continue
+                line = raw.decode("utf-8")
+                if not line.startswith("data:"):
+                    continue
+                data = line[len("data:") :].strip()
+                if data == "[DONE]":
+                    break
+                chunk = json.loads(data)
+                if chunk.get("usage"):
+                    usage = chunk["usage"]
+                snapshot = await actor.get_request_lifecycle.remote(
+                    LIFECYCLE_REQUEST_ID
+                )
+                if snapshot is not None:
+                    snapshots.append(snapshot)
+                live_active_requests.append(
+                    await actor.get_worker_active_requests.remote(worker_id)
+                )
+
+        assert usage is not None, "expected a final usage chunk"
+        assert snapshots, "request was never observed in flight on the actor"
+
+        # Every in-flight snapshot upholds exact token and block accounting.
+        block_size = await actor.get_block_size.remote()
+        previous_output_tokens = 0
+        for snapshot in snapshots:
+            assert snapshot["prompt_tokens"] == usage["prompt_tokens"]
+            assert previous_output_tokens <= snapshot["output_tokens"]
+            assert snapshot["output_tokens"] <= usage["completion_tokens"]
+            previous_output_tokens = snapshot["output_tokens"]
+            total_blocks = math.ceil(
+                (usage["prompt_tokens"] + snapshot["output_tokens"]) / block_size
+            )
+            assert snapshot["total_blocks"] == total_blocks
+            if snapshot["output_tokens"] > 0:
+                assert snapshot["prefill_completed"] is True
+
+        # The request was booked as active load on its worker while in flight,
+        # and every hook's call into the live selection service succeeded.
+        assert max(live_active_requests) >= 1, "request was never booked as active load"
+        assert await actor.get_errors.remote() == []
+
+        # Completion frees the request from both the actor view and the load
+        # tracker (no leaked active load to skew later scoring).
+        worker_id = snapshots[-1]["worker_id"]
+        await async_wait_for_condition(
+            lambda: ray.get(actor.get_request_lifecycle.remote(LIFECYCLE_REQUEST_ID))
+            is None,
+            timeout=15,
+            retry_interval_ms=200,
+        )
+        assert await actor.get_active_request_ids.remote() == []
+        await async_wait_for_condition(
+            lambda: ray.get(actor.get_worker_active_requests.remote(worker_id)) == 0,
+            timeout=15,
+            retry_interval_ms=200,
+        )
+
+        events = [
+            (name, args)
+            for name, args in await actor.get_event_log.remote()
+            if args[0] == LIFECYCLE_REQUEST_ID
+        ]
+        names = [name for name, _ in events]
+        assert names[0] == "on_request_added"
+        assert names[1] == "on_prefill_complete"
+        assert names[-1] == "on_request_completed"
+        assert names[2:-1] == ["on_decode_progress"] * (len(names) - 3)
+
+        added_args = events[0][1]
+        assert added_args[1] == worker_id
+        prompt_token_ids = added_args[2]
+        assert len(prompt_token_ids) == usage["prompt_tokens"]  # token ids booked
+        assert added_args[3] == MAX_TOKENS  # client max_tokens drives decode decay
+        decode_counts = [
+            args[1] for name, args in events if name == "on_decode_progress"
+        ]
+        assert decode_counts == sorted(set(decode_counts))  # strictly increasing
+        assert decode_counts[-1] == usage["completion_tokens"]
+
+        # The engine indexed the prompt's KV blocks; they show as cache overlap
+        # on the worker (the prefix the next overlapping request would reuse).
+        prompt_blocks = num_prompt_blocks(prompt_token_ids)
+        assert prompt_blocks >= 1
+        await async_wait_for_condition(
+            lambda: ray.get(actor.get_overlap_blocks.remote(prompt_token_ids)).get(
+                worker_id, 0
+            )
+            == prompt_blocks,
+            timeout=60,
+            retry_interval_ms=500,
+        )
+
+    @pytest.mark.asyncio
+    async def test_booked_reservation_changes_scoring_load(self, deployed_handle):
+        """A reservation booked through the lifecycle hooks shows up as active
+        load on the worker (the value scoring consumes) and freeing it restores
+        baseline."""
+        actor = discover_deployment_actor(
+            APP_NAME, deployed_handle.deployment_name, KV_ROUTER_ACTOR_NAME
+        )
+        worker_id = await self._registered_worker(actor)
+        assert await actor.get_worker_active_requests.remote(worker_id) == 0
+
+        await actor.on_request_added.remote(
+            "probe", worker_id, list(range(64)), expected_output_tokens=32
+        )
+        assert await actor.get_worker_active_requests.remote(worker_id) == 1
+
+        await actor.on_request_completed.remote("probe")
+        await async_wait_for_condition(
+            lambda: ray.get(actor.get_worker_active_requests.remote(worker_id)) == 0,
+            timeout=15,
+            retry_interval_ms=200,
+        )
+
+
+if __name__ == "__main__":
+    if not ray.is_initialized():
+        ray.init(address="auto")
+    sys.exit(pytest.main(["-v", "-s", __file__]))

release/llm_tests/kv_router_test/utils.py

@@ -0,0 +1,41 @@
+"""Shared helpers for the KV-router GPU release tests."""
+
+import ray
+from ray.serve._private.constants import (
+    SERVE_DEPLOYMENT_ACTOR_PREFIX,
+    SERVE_NAMESPACE,
+)
+from ray.serve.experimental.round_robin_router import RoundRobinRouter
+from ray.serve.llm.request_router import KVAwareRouter
+
+
+class _TestKVAwareRouter(RoundRobinRouter, KVAwareRouter):
+    """A ``KVAwareRouter`` subclass that borrows ``RoundRobinRouter``'s selection.
+
+    KVAwareRouter's own routing (scoring by KV-cache overlap) is not implemented
+    yet, so this inherits RoundRobinRouter's ``choose_replicas`` (via MRO) while
+    remaining a KVAwareRouter subclass.
+
+    TODO (jeffreywang): Remove RoundRobinRouter base class once KVAwareRouter's selection
+    is implemented.
+    """
+
+
+def discover_deployment_actor(app_name, deployment_name, actor_name):
+    """Resolve a deployment-scoped actor by its registered name.
+
+    The test driver isn't a replica, so ``get_deployment_actor`` is
+    unavailable; match the name's stable prefix/suffix instead (the middle
+    embeds an opaque code_version).
+    """
+    prefix = f"{SERVE_DEPLOYMENT_ACTOR_PREFIX}{app_name}::{deployment_name}::"
+    suffix = f"::{actor_name}"
+    for entry in ray.util.list_named_actors(all_namespaces=True):
+        name = entry.get("name") or ""
+        if (
+            entry.get("namespace") == SERVE_NAMESPACE
+            and name.startswith(prefix)
+            and name.endswith(suffix)
+        ):
+            return ray.get_actor(name, namespace=SERVE_NAMESPACE)
+    return None

release/llm_tests/serve/benchmark/benchmark_vllm.py

@@ -0,0 +1,279 @@
+# There is a dead-lock issue that arises due to gevent's monkey-patching
+# https://github.com/ipython/ipython/issues/11730
+# Fix: We do this import first before anything else
+# Thus the # noqa tags are needed below
+import gevent.monkey
+
+gevent.monkey.patch_all()
+
+import sys  # noqa: E402
+import argparse  # noqa: E402
+import json  # noqa: E402
+import logging  # noqa: E402
+import os  # noqa: E402
+import subprocess  # noqa: E402
+import threading  # noqa: E402
+import time  # noqa: E402
+from datetime import datetime  # noqa: E402
+
+from bm import run_bm  # noqa: E402
+from common import write_to_s3, get_llm_config  # noqa: E402
+
+RAYLLM_RELEASE_TEST_PERF_SERVICE_NAME = "rayllm_release_test_perf_service"
+THREAD_CLEANUP_TIMEOUT_S = 10
+
+
+logger = logging.getLogger(__file__)
+logging.basicConfig(
+    format="%(asctime)s %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.INFO,
+)
+
+
+def get_timestamp():
+    return datetime.now().strftime("%H:%M:%S.%f")[:-3]
+
+
+class ColoredLogger:
+    HEADER = "\033[95m"
+    SERVER = "\033[94m"
+    CLIENT = "\033[92m"
+    WARNING = "\033[93m"
+    ERROR = "\033[91m"
+    ENDC = "\033[0m"
+
+    @staticmethod
+    def log_server(msg):
+        print(
+            f"{ColoredLogger.SERVER}[SERVER {get_timestamp()}] {msg}{ColoredLogger.ENDC}"
+        )
+
+    @staticmethod
+    def log_client(msg):
+        print(
+            f"{ColoredLogger.CLIENT}[CLIENT {get_timestamp()}] {msg}{ColoredLogger.ENDC}"
+        )
+
+    @staticmethod
+    def log_error(msg):
+        print(
+            f"{ColoredLogger.ERROR}[ERROR {get_timestamp()}] {msg}{ColoredLogger.ENDC}"
+        )
+
+
+def stream_process_output(process, logger_func, stop_event, is_error: bool = False):
+    while not stop_event.is_set():
+        if is_error:
+            output_line = process.stderr.readline()
+        else:
+            output_line = process.stdout.readline()
+        if output_line:
+            logger_func(output_line.strip())
+        elif process.poll() is not None:
+            break
+
+
+def get_vllm_cli_args(llm_config):
+    engine_kwargs = llm_config["engine_kwargs"]
+    # When we define tokenizer_pool size, vllm, by default, uses Ray
+    # that breaks the assumption that this script should not use ray
+    # TODO (Kourosh): When the job issue with non driver ray
+    # subprocesses are resolved we can remove these constraints
+    engine_kwargs.pop("tokenizer_pool_extra_config", None)
+    engine_kwargs.pop("tokenizer_pool_size", None)
+    engine_kwargs.pop("tokenizer_pool_type", None)
+
+    cli_args = ["--model", llm_config["model_loading_config"]["model_id"]]
+    for key, value in engine_kwargs.items():
+        cli_args.extend(_engine_kwarg_to_cli(key, value))
+    return cli_args
+
+
+def _engine_kwarg_to_cli(key, value):
+    """Translate one engine_kwargs entry to vllm CLI args.
+
+    vLLM uses argparse.BooleanOptionalAction for bool fields, so emit the
+    explicit `--flag` / `--no-flag` form to honor the configured value
+    regardless of the engine default.
+    """
+    if value is None:
+        return []
+    flag = "--" + key.replace("_", "-")
+    if isinstance(value, bool):
+        return [flag] if value else ["--no-" + key.replace("_", "-")]
+    if isinstance(value, dict):
+        return [flag, json.dumps(value)]
+    return [flag, str(value)]
+
+
+def get_ray_options(llm_config):
+    num_gpus = llm_config["tensor_parallelism"]["degree"]
+    acc_type = llm_config["accelerator_type"]
+    resources = {f"accelerator_type:{acc_type}": 0.001}
+
+    return {"num_gpus": num_gpus, "resources": resources}
+
+
+def start_vllm_process(vllm_cli_args):
+    server_process = subprocess.Popen(
+        ["python", "-m", "vllm.entrypoints.openai.api_server"] + vllm_cli_args,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True,  # Return strings from stdout/stderr
+        bufsize=1,  # Line buffered
+        env=os.environ.copy(),
+    )
+
+    return server_process
+
+
+def run_vllm_benchmark(vllm_cli_args):
+    stop_event = threading.Event()
+    results = {}
+
+    try:
+        # Start the server process
+        ColoredLogger.log_server("Starting vLLM server...")
+        server_process = start_vllm_process(vllm_cli_args)
+
+        # Start server output streaming thread
+        server_thread = threading.Thread(
+            target=stream_process_output,
+            args=(server_process, ColoredLogger.log_server, stop_event),
+            daemon=True,  # Daemonize the thread so it stops when the main thread stops.
+        )
+        server_thread.start()
+
+        # Start server error streaming thread
+        server_error_thread = threading.Thread(
+            target=stream_process_output,
+            args=(server_process, ColoredLogger.log_server, stop_event),
+            kwargs={"is_error": True},
+            daemon=True,  # Daemonize the thread so it stops when the main thread stops.
+        )
+        server_error_thread.start()
+
+        # Wait for server to be ready
+        server_ready = False
+        start_time = time.time()
+        timeout = 300  # 5 minutes timeout
+
+        while not server_ready and time.time() - start_time < timeout:
+            if server_process.poll() is not None:
+                raise Exception("Server process terminated unexpectedly")
+
+            # Check if server is responding
+            try:
+                import requests
+
+                response = requests.get("http://localhost:8000/health")
+                if response.status_code == 200:
+                    server_ready = True
+                    ColoredLogger.log_server("Server is ready!")
+            except Exception:
+                time.sleep(1)
+                continue
+
+        if not server_ready:
+            raise TimeoutError("Server failed to start within timeout period")
+
+        # Start benchmark
+        ColoredLogger.log_client("Starting benchmark...")
+        results = run_bm(
+            api_url="http://localhost:8000",
+            api_key="NONE",
+            concurrency=[1, 2, 4, 8, 16, 32],
+            run_time="1m",
+            prompt_tokens=256,
+            max_tokens=64,
+            stream=False,
+            summary_file="./results.csv",
+        )
+        print(
+            "Writing final result to AWS Firehose:",
+            json.dumps(results, indent=4, sort_keys=True),
+            sep="\n",
+        )
+
+        ColoredLogger.log_client("Benchmark completed successfully")
+
+    except Exception as e:
+        ColoredLogger.log_error(f"Error during benchmark: {str(e)}")
+        raise
+
+    finally:
+        # Clean up
+
+        if "server_process" in locals():
+            ColoredLogger.log_server("Shutting down server...")
+            server_process.terminate()
+            server_process.wait(timeout=THREAD_CLEANUP_TIMEOUT_S)
+
+        stop_event.set()
+        # Wait for all threads to complete
+        if "server_thread" in locals():
+            server_thread.join(timeout=THREAD_CLEANUP_TIMEOUT_S)
+        if "server_error_thread" in locals():
+            server_error_thread.join(timeout=THREAD_CLEANUP_TIMEOUT_S)
+
+        # Wait some time to make sure everyting is cleanend up.
+        time.sleep(5)
+
+    return results
+
+
+def upload_results_to_s3(s3_path, results, service_metadata):
+    if any(result is None for result in results):
+        raise ValueError(
+            "Found None results during benchmarking. " "This should not have happened."
+        )
+
+    data_to_write = [{**result, **service_metadata} for result in results]
+    write_to_s3(data_to_write, s3_path)
+
+
+def main(pargs):
+    llm_config = get_llm_config(pargs.llm_config)
+    vllm_cli_args = get_vllm_cli_args(llm_config)
+
+    results = run_vllm_benchmark(vllm_cli_args)
+
+    tag = f"{llm_config['accelerator_type']}-TP{llm_config['engine_kwargs']['tensor_parallel_size']}"
+    service_metadata = {
+        "cloud_name": "",
+        "service_name": "",
+        "py_version": f"py{sys.version_info.major}{sys.version_info.minor}",
+        "tag": tag,
+        "vllm_engine": "V1",
+    }
+
+    # Post the results to S3
+    if results:
+        print(
+            "Writing final result to AWS S3:",
+            json.dumps(results, indent=4, sort_keys=True),
+            sep="\n",
+        )
+        upload_results_to_s3(pargs.remote_result_path, results, service_metadata)
+    else:
+        raise ValueError(
+            "For some reason the benchmarking results are empty. Something is wrong."
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--llm-config",
+        type=str,
+        required=True,
+        default="The LLM config to start vLLM engine",
+    )
+    parser.add_argument(
+        "--remote-result-path",
+        type=str,
+        required=True,
+        help="The remote s3 path to store intermediate results on.",
+    )
+    main(parser.parse_args())

release/llm_tests/serve/benchmark/bm.py

@@ -0,0 +1,54 @@
+# There is a dead-lock issue that arises due to gevent's monkey-patching
+# https://github.com/ipython/ipython/issues/11730
+# Fix: We do this import first before anything else
+# ruff: noqa: E402
+import gevent.monkey
+
+gevent.monkey.patch_all()
+
+from benchmark.mocks import LLMLoadTester
+from benchmark.configs import LoadTestConfig
+
+from typing import List, Optional
+import openai
+
+
+def run_bm(
+    api_url: str,
+    api_key: Optional[str] = None,
+    concurrency: Optional[List[int]] = None,
+    run_time: str = "1m",
+    prompt_tokens: int = 512,
+    max_tokens: int = 64,
+    stream: bool = False,
+    summary_file: str = "./results.csv",
+):
+    if api_key is None:
+        api_key = "NONE"
+
+    # Get model_id
+    client = openai.Client(base_url=f"{api_url}/v1", api_key=api_key)
+    models = client.models.list().model_dump()["data"]
+
+    if len(models) != 1:
+        raise ValueError("The service is expected to have only one model.")
+
+    model_id = models[0]["id"]
+    results = []
+    for n_users in concurrency:
+        config = LoadTestConfig(
+            host=api_url,
+            api_key=api_key,
+            provider="openai",
+            model=model_id,
+            stream=stream,
+            prompt_tokens=prompt_tokens,
+            max_tokens=max_tokens,
+            users=n_users,
+            run_time=run_time,
+            summary_file=summary_file,
+        )
+        tester = LLMLoadTester(config)
+        results.append(tester.run())
+
+    return results

release/llm_tests/serve/benchmark/common.py

@@ -0,0 +1,109 @@
+"""Common utilities shared between release tests"""
+
+import json
+import logging
+import os
+from typing import Dict, List, Any
+from urllib.parse import urlparse
+
+import boto3
+import yaml
+
+logger = logging.getLogger(__file__)
+logging.basicConfig(
+    format="%(asctime)s %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+    level=logging.INFO,
+)
+
+
+def get_llm_config(serve_config_file: List[Dict]) -> List[Any]:
+    """Get the first llm_config from serve config file."""
+    with open(serve_config_file, "r") as f:
+        loaded_llm_config = yaml.safe_load(f)
+
+    applications = loaded_llm_config["applications"]
+    config = applications[0]["args"]["llm_configs"][0]
+    if isinstance(config, dict):
+        return config
+
+    assert isinstance(config, str)
+    with open(config, "r") as f:
+        loaded_llm_config = yaml.safe_load(f)
+
+    return loaded_llm_config
+
+
+def read_yaml(file_path: str) -> Dict:
+    if not file_path.endswith(".yaml"):
+        raise RuntimeError(
+            "Must pass in a Serve config yaml file using the -f option. Got "
+            f'file path "{file_path}", which does not end in ".yaml".'
+        )
+    if not os.path.exists(file_path):
+        raise RuntimeError(f"File path {file_path} does not exist.")
+
+    with open(file_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def write_to_s3(data_to_write: List[Dict], s3_path: str):
+    # Parse S3 path using urllib
+    parsed_url = urlparse(s3_path)
+    bucket_name = parsed_url.netloc
+    # Remove leading slash and get the rest of the path
+    key = parsed_url.path.lstrip("/")
+
+    # If no file extension provided, append .json
+    if not any(key.endswith(ext) for ext in [".jsonl", ".json"]):
+        key += ".jsonl"
+
+    s3_client = boto3.client("s3")
+
+    try:
+        # Convert data to JSON string
+        jsonl_data = "\n".join(json.dumps(record) for record in data_to_write)
+
+        # Upload to S3
+        s3_client.put_object(
+            Bucket=bucket_name,
+            Key=key,
+            Body=jsonl_data,
+            ContentType="application/x-jsonlines",  # MIME type for JSONL
+        )
+
+        logging.info(f"Successfully wrote {len(data_to_write)} records to {s3_path}")
+
+    except Exception as e:
+        logging.error(f"Failed to write to S3: {str(e)}")
+        raise
+
+
+def read_from_s3(s3_path: str) -> List[Dict]:
+    # Parse S3 path using urllib
+    parsed_url = urlparse(s3_path)
+    bucket_name = parsed_url.netloc
+    key = parsed_url.path.lstrip("/")
+
+    # Initialize S3 client
+    s3_client = boto3.client("s3")
+
+    try:
+        # Get the object from S3
+        response = s3_client.get_object(Bucket=bucket_name, Key=key)
+
+        # Read the data
+        data = response["Body"].read().decode("utf-8")
+
+        records = [
+            json.loads(line)
+            for line in data.splitlines()
+            if line.strip()  # Skip empty lines
+        ]
+
+        logging.info(f"Successfully read {len(records)} records from {s3_path}")
+        return records
+
+    except Exception as e:
+        logging.error(f"Failed to read from S3: {str(e)}")
+        raise

release/llm_tests/serve/benchmark/configs.py

@@ -0,0 +1,156 @@
+from typing import List, Optional
+from enum import Enum
+from pydantic import BaseModel, Field
+import argparse
+
+
+class DistributionType(str, Enum):
+    CONSTANT = "constant"
+    UNIFORM = "uniform"
+    EXPONENTIAL = "exponential"
+    NORMAL = "normal"
+
+
+class TokensDistributionType(str, Enum):
+    CONSTANT = "constant"
+    UNIFORM = "uniform"
+    EXPONENTIAL = "exponential"
+
+
+class LoadTestConfig(BaseModel):
+    provider: Optional[str] = Field(
+        None,
+        description="Which flavor of API to use. If not specified, we'll try to guess based on the URL and /v1/models output",
+    )
+
+    model: Optional[str] = Field(
+        None,
+        description="The model to use for generating text. If not specified we will pick the first model from the service as returned by /v1/models",
+    )
+
+    chat: bool = Field(True, description="Use /v1/chat/completions API")
+
+    prompt_tokens: int = Field(
+        512,
+        description="Length of the prompt in tokens",
+    )
+
+    prompt_chars: Optional[int] = Field(
+        None,
+        description="Length of the prompt in characters",
+    )
+
+    prompt_text: Optional[str] = Field(
+        None,
+        description="Prompt text to use instead of generating one. It can be a file reference starting with an ampersand, e.g. `@prompt.txt`",
+    )
+
+    prompt_randomize: bool = Field(
+        False,
+        description="Include a few random numbers in the generated prompt to avoid caching",
+    )
+
+    max_tokens: int = Field(
+        64,
+        description="Max number of tokens to generate. If max_tokens_distribution is non-constant this is going to be the mean",
+    )
+
+    max_tokens_cap: Optional[int] = Field(
+        None,
+        description="If max_tokens_distribution is non-constant, this truncates the distribition at the specified limit",
+    )
+
+    max_tokens_distribution: TokensDistributionType = Field(
+        TokensDistributionType.CONSTANT,
+        description="How to sample max_tokens on each request",
+    )
+
+    max_tokens_range: float = Field(
+        0.3,
+        description="Specifies the width of the distribution. Specified value `alpha` is relative to `max_tokens`",
+    )
+
+    stream: bool = Field(True, description="Use the streaming API")
+
+    api_key: Optional[str] = Field(
+        None,
+        description="Auth for the API",
+    )
+
+    temperature: float = Field(0.1, description="Temperature parameter for the API")
+
+    logprobs: Optional[int] = Field(
+        None,
+        description="Whether to ask for logprobs, it makes things slower for some providers but is necessary for token count in streaming",
+    )
+
+    summary_file: Optional[str] = Field(
+        None,
+        description="Append the line with the summary to the specified CSV file",
+    )
+
+    qps: Optional[float] = Field(
+        None,
+        description="Enabled 'fixed QPS' mode where requests are issues at the specified rate regardless of how long the processing takes",
+    )
+
+    qps_distribution: DistributionType = Field(
+        DistributionType.CONSTANT,
+        description="Must be used with qps. Specifies how to space out requests",
+    )
+
+    burst: Optional[float] = Field(
+        None,
+        description="Makes requests to arrive in bursts every specified number of seconds",
+    )
+
+    tokenizer: Optional[str] = Field(
+        None,
+        description="Specify HF tokenizer to use for validating the output of the model",
+    )
+
+    show_response: bool = Field(
+        False,
+        description="Print the result of each generation",
+    )
+
+    prompt_cache_max_len: int = Field(
+        0,
+        description="Maximum length of the prompt cache to use",
+    )
+
+    header: List[str] = Field(
+        default_factory=list,
+        description="Arbitrary headers to add to the inference request",
+    )
+
+    n: int = Field(
+        1,
+        description="How many sequences to generate (makes sense to use with non-zero temperature)",
+    )
+
+    host: Optional[str] = Field(
+        default=None,
+        description="Host to load test in the following format: http://10.21.32.33",
+    )
+
+    reset_stats: bool = Field(
+        default=True,
+        description="Determines if stats should be reset once hatching is complete",
+    )
+
+    users: int = Field(
+        default=None,
+        description="Number of concurrent users to spawn for benchmarking.",
+    )
+
+    run_time: str = Field(
+        default="30s",
+        description="The runtime it is in form of Ns, Nm, or Nh, for seconds, minutes, and hours.",
+    )
+
+    def to_namespace(self) -> argparse.Namespace:
+        """
+        Convert the model to an argparse.Namespace object
+        """
+        return argparse.Namespace(**self.dict())

release/llm_tests/serve/benchmark/constants.py

@@ -0,0 +1,12 @@
+"""This file contains constants used throughout the Buildkite files."""
+
+# Config for the rayllm release test.
+RAYLLM_RELEASE_TEST_SERVICE_NAME = "rayllm_release_test_service"
+RAYLLM_RELEASE_TEST_COMPUTE_CONFIG_NAME = "rayllm-release-test"
+
+DEFAULT_CLOUD = "serve_release_tests_cloud"
+
+CLOUD_PROVIDER_TO_CLOUD_NAME = {
+    "aws": DEFAULT_CLOUD,
+    "gcp": "anyscale_gcp_public_default_cloud_us_west_1",
+}

release/llm_tests/serve/benchmark/firehose_utils.py

@@ -0,0 +1,84 @@
+import json
+import os
+import time
+from enum import Enum
+from typing import Any, Dict
+
+import ray
+import boto3
+from pydantic import BaseModel, field_validator
+
+STREAM_NAME = "rayllm-ci-results"
+DEFAULT_TABLE_NAME = "release_test_result"
+# Time to sleep in-between firehose writes to make sure the timestamp between
+# records are distinct
+SLEEP_BETWEEN_FIREHOSE_WRITES_MS = 50
+
+
+class RecordName(str, Enum):
+    STARTUP_TEST = "service-startup-test"
+    STARTUP_TEST_GCP = "service-startup-test-gcp"
+    STARTUP_TEST_AWS = "service-startup-test-aws"
+    RAYLLM_PERF_TEST = "rayllm-perf-test"
+    VLLM_PERF_TEST = "vllm-perf-test"
+
+
+class FirehoseRecord(BaseModel):
+    record_name: RecordName
+    record_metrics: Dict[str, Any]
+
+    @field_validator("record_name", mode="before")
+    def validate_record_name(cls, v):
+        if isinstance(v, str):
+            return RecordName(v)
+        return v
+
+    def write(self, verbose: bool = False):
+        final_result = {
+            "_table": DEFAULT_TABLE_NAME,
+            "name": str(self.record_name.value),
+            "branch": os.environ.get("BUILDKITE_BRANCH", ""),
+            "commit": ray.__commit__,
+            "report_timestamp_ms": int(time.time() * 1000),
+            "results": {**self.record_metrics},
+        }
+
+        if verbose:
+            print(
+                "Writing final result to AWS Firehose:",
+                json.dumps(final_result, indent=4, sort_keys=True),
+                sep="\n",
+            )
+
+        # Add newline character to separate records
+        data = json.dumps(final_result) + "\n"
+
+        # Need to assume the role in order to share access to the Firehose
+        sts_client = boto3.client("sts")
+
+        assumed_role = sts_client.assume_role(
+            RoleArn="arn:aws:iam::830883877497:role/service-role/KinesisFirehoseServiceRole-rayllm-ci-res-us-west-2-1728664186256",
+            RoleSessionName="FirehosePutRecordSession",
+        )
+
+        credentials = assumed_role["Credentials"]
+
+        # Use the assumed credentials to create a Firehose client
+        firehose_client = boto3.client(
+            "firehose",
+            region_name="us-west-2",
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+        )
+
+        response = firehose_client.put_record(
+            DeliveryStreamName=STREAM_NAME, Record={"Data": data}
+        )
+
+        if verbose:
+            print("PutRecord response:")
+            print(response)
+
+        # Add some delay to make sure timestamps are unique ints.
+        time.sleep(SLEEP_BETWEEN_FIREHOSE_WRITES_MS / 1000)

release/llm_tests/serve/benchmark/mocks.py

@@ -0,0 +1,105 @@
+from locust.env import Environment
+from locust.stats import stats_printer, stats_history, print_stats
+import gevent
+from typing import Dict, Any
+import logging
+from locust.log import setup_logging
+
+from benchmark.load_test import LLMUser, events, collect_metrics
+from benchmark.configs import LoadTestConfig
+
+
+class LLMLoadTester:
+    """
+    Usage Example:
+
+    ```python
+        config = LoadTestConfig(
+            host="http://localhost:8000",
+            provider="vllm",
+            model="meta-llama/Meta-Llama-3.1-8B-Instruct",
+            api_key="NONE",
+            prompt_tokens=550,
+            max_tokens=150,
+            users=128,
+            run_time="1m",
+            summary_file="./vllm.csv"
+        )
+        tester = LLMLoadTester(config)
+        results = tester.run()
+    ```
+
+
+    """
+
+    def __init__(self, config: LoadTestConfig):
+        self.config = config
+
+    def _setup_environment(self) -> Environment:
+
+        setup_logging("INFO", None)
+
+        # Setup Environment and Runner
+        env = Environment(
+            user_classes=[LLMUser],
+            host=self.config.host,
+            reset_stats=self.config.reset_stats,
+            events=events,
+        )
+        env.parsed_options = self.config.to_namespace()
+
+        return env
+
+    def run(self) -> Dict[str, Any]:
+        try:
+            # Setup environment
+            env = self._setup_environment()
+            env.create_local_runner()
+
+            # Log test start
+            logging.info(f"Starting test with {self.config.users} users")
+
+            # Create greenlets for stats
+            stats_printer_greenlet = gevent.spawn(stats_printer(env.stats))
+            stats_history_greenlet = gevent.spawn(stats_history, env.runner)
+
+            # Start the test
+            env.runner.start(
+                user_count=self.config.users,
+                spawn_rate=self.config.users,
+            )
+
+            # Run for specified duration
+            gevent.sleep(self._parse_time(self.config.run_time))
+
+            # Stop the test
+            env.runner.quit()
+            entries = collect_metrics(env)
+
+            # Wait for greenlets
+            env.runner.greenlet.join()
+            stats_printer_greenlet.kill()
+            stats_history_greenlet.kill()
+
+            # Print final stats
+            print_stats(env.stats)
+
+            return entries
+
+        except Exception as e:
+            logging.error(f"Test failed: {str(e)}")
+            raise
+
+    @staticmethod
+    def _parse_time(time_str: str) -> int:
+        """Convert time string (e.g., '30s', '1m', '1h') to seconds"""
+        unit = time_str[-1]
+        value = int(time_str[:-1])
+        if unit == "s":
+            return value
+        elif unit == "m":
+            return value * 60
+        elif unit == "h":
+            return value * 3600
+        else:
+            raise ValueError(f"Invalid time unit: {unit}")

release/llm_tests/serve/configs/lmcache/decoder.yaml

@@ -0,0 +1,12 @@
+local_cpu: False
+max_local_cpu_size: 0
+max_local_disk_size: 0
+remote_serde: NULL
+
+enable_nixl: True
+nixl_role: "receiver"
+nixl_receiver_host: "localhost"
+nixl_receiver_port: 55555
+nixl_buffer_size: 1073741824 # 1GB
+nixl_buffer_device: "cuda"
+nixl_enable_gc: True

release/llm_tests/serve/configs/lmcache/prefiller.yaml

@@ -0,0 +1,12 @@
+local_cpu: False
+max_local_cpu_size: 0
+max_local_disk_size: 0
+remote_serde: NULL
+
+enable_nixl: True
+nixl_role: "sender"
+nixl_receiver_host: "localhost"
+nixl_receiver_port: 55555
+nixl_buffer_size: 1073741824 # 1GB
+nixl_buffer_device: "cuda"
+nixl_enable_gc: True

release/llm_tests/serve/configs/model_config/llama_3dot1_8b_lora.yaml

@@ -0,0 +1,14 @@
+model_loading_config:
+  model_id: meta-llama/Meta-Llama-3.1-8B-Instruct-Fine-Tuned
+  model_source: meta-llama/Llama-3.1-8B-Instruct
+
+accelerator_type: A10G
+
+engine_kwargs:
+  max_model_len: 2048
+  enable_lora: true
+  enforce_eager: true
+
+lora_config:
+  dynamic_lora_loading_path: "s3://anyscale-production-data-cld-wy5a6nhazplvu32526ams61d98/org_7c1Kalm9WcX2bNIjW53GUT/cld_wy5a6nhazplvu32526ams61d98/artifact_storage/rayllm_release_test/lora_fine_tuning"
+  max_num_adapters_per_replica: 2

release/llm_tests/serve/configs/model_config/llama_3dot1_8b_quantized_tp1.yaml

@@ -0,0 +1,10 @@
+model_loading_config:
+  model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+
+accelerator_type: A10G
+
+engine_kwargs:
+  max_model_len: 8192
+  tensor_parallel_size: 1
+  # NOTE: This is used for perf testing as well, so cuda graph must be enabled.
+  enforce_eager: false

release/llm_tests/serve/configs/model_config/llama_3dot1_8b_quantized_tp1_prefix_cache.yaml

@@ -0,0 +1,19 @@
+model_loading_config:
+  model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+
+accelerator_type: A10G
+
+engine_kwargs:
+  max_model_len: 8192
+  tensor_parallel_size: 1
+  # NOTE: This is used for perf testing as well, so cuda graph must be enabled.
+  enforce_eager: false
+
+deployment_config:
+  # Two replicas so the prefix-cache-aware router actually routes across
+  # replicas. With one replica affinity is trivial.
+  autoscaling_config:
+    min_replicas: 2
+    max_replicas: 2
+  request_router_config:
+    request_router_class: ray.serve.llm.request_router.PrefixCacheAffinityRouter

release/llm_tests/serve/configs/model_config/llama_3dot1_8b_tp2.yaml

@@ -0,0 +1,16 @@
+model_loading_config:
+  model_id: meta-llama/Llama-3.1-8B-Instruct
+
+accelerator_type: A10G
+
+placement_group_config:
+  bundles:
+    - GPU: 1
+    - GPU: 1
+  strategy: STRICT_PACK
+
+engine_kwargs:
+  max_model_len: 8192
+  tensor_parallel_size: 2
+  # NOTE: This is used for perf testing as well, so cuda graph must be enabled.
+  enforce_eager: false

release/llm_tests/serve/configs/serve_llama_3dot1_8b_lora.yaml

@@ -0,0 +1,7 @@
+applications:
+  - args:
+      llm_configs:
+        - ./configs/model_config/llama_3dot1_8b_lora.yaml
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot1_8b_quantized_tp1.yaml

@@ -0,0 +1,7 @@
+applications:
+  - args:
+      llm_configs:
+        - ./configs/model_config/llama_3dot1_8b_quantized_tp1.yaml
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot1_8b_quantized_tp1_1p1d.yaml

@@ -0,0 +1,17 @@
+applications:
+  - args:
+      prefill_config: &config
+        model_loading_config:
+          model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+        accelerator_type: A10G
+        engine_kwargs:
+          max_model_len: 8192
+          tensor_parallel_size: 1
+          enforce_eager: false
+          kv_transfer_config:
+            kv_connector: NixlConnector
+            kv_role: kv_both
+      decode_config: *config
+    import_path: ray.serve.llm:build_pd_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot1_8b_quantized_tp1_2p6d.yaml

@@ -0,0 +1,35 @@
+applications:
+  - args:
+      prefill_config:
+        model_loading_config:
+          model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+        accelerator_type: A10G
+        engine_kwargs:
+          max_model_len: 8192
+          tensor_parallel_size: 1
+          enforce_eager: true
+          kv_transfer_config:
+            kv_connector: NixlConnector
+            kv_role: kv_both
+        deployment_config:
+          autoscaling_config:
+            min_replicas: 2
+            max_replicas: 2
+      decode_config:
+        model_loading_config:
+          model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+        accelerator_type: A10G
+        engine_kwargs:
+          max_model_len: 8192
+          tensor_parallel_size: 1
+          enforce_eager: true
+          kv_transfer_config:
+            kv_connector: NixlConnector
+            kv_role: kv_both
+        deployment_config:
+          autoscaling_config:
+            min_replicas: 6
+            max_replicas: 6
+    import_path: ray.serve.llm:build_pd_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot1_8b_quantized_tp1_2p6d_lmcache.yaml

@@ -0,0 +1,52 @@
+applications:
+  - args:
+
+      prefill_config:
+        model_loading_config:
+          model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+        accelerator_type: A10G
+        engine_kwargs:
+          max_model_len: 8192
+          tensor_parallel_size: 1
+          enforce_eager: true
+          kv_transfer_config:
+            kv_connector: LMCacheConnectorV1
+            kv_role: kv_producer
+            kv_connector_extra_config:
+              discard_partial_chunks: false
+              lmcache_rpc_port: producer1
+        deployment_config:
+          autoscaling_config:
+            min_replicas: 2
+            max_replicas: 2
+        runtime_env:
+          env_vars:
+            LMCACHE_CONFIG_FILE: configs/lmcache/prefiller.yaml
+            LMCACHE_USE_EXPERIMENTAL: "True"
+
+      decode_config:
+        model_loading_config:
+          model_id: neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w4a16
+        accelerator_type: A10G
+        engine_kwargs:
+          max_model_len: 8192
+          tensor_parallel_size: 1
+          enforce_eager: true
+          kv_transfer_config:
+            kv_connector: LMCacheConnectorV1
+            kv_role: kv_consumer
+            kv_connector_extra_config:
+              discard_partial_chunks: false
+              lmcache_rpc_port: consumer1
+        deployment_config:
+          autoscaling_config:
+            min_replicas: 6
+            max_replicas: 6
+        runtime_env:
+          env_vars:
+            LMCACHE_CONFIG_FILE: configs/lmcache/decoder.yaml
+            LMCACHE_USE_EXPERIMENTAL: "True"
+
+    import_path: ray.serve.llm:build_pd_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot1_8b_quantized_tp1_prefix_cache.yaml

@@ -0,0 +1,7 @@
+applications:
+  - args:
+      llm_configs:
+        - ./configs/model_config/llama_3dot1_8b_quantized_tp1_prefix_cache.yaml
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot1_8b_tp2.yaml

@@ -0,0 +1,7 @@
+applications:
+  - args:
+      llm_configs:
+        - ./configs/model_config/llama_3dot1_8b_tp2.yaml
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot2_1b_no_accelerator.yaml

@@ -0,0 +1,11 @@
+applications:
+  - args:
+      llm_configs:
+        - model_loading_config:
+            model_id: meta-llama/Llama-3.2-1B-Instruct
+          engine_kwargs:
+            max_model_len: 8192
+            enforce_eager: true
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_llama_3dot2_1b_s3.yaml

@@ -0,0 +1,15 @@
+applications:
+  - args:
+      llm_configs:
+        - model_loading_config:
+            model_id: my_llama
+            model_source:
+              bucket_uri: s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct
+          accelerator_type: L4
+          engine_kwargs:
+            max_model_len: 8192
+            tensor_parallel_size: 1
+            enforce_eager: true
+    import_path: ray.serve.llm:build_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/configs/serve_phi_tiny_moe_dp4_gang.yaml

@@ -0,0 +1,26 @@
+applications:
+  - args:
+      llm_config:
+        model_loading_config:
+          model_id: microsoft/Phi-tiny-MoE-instruct
+          model_source: microsoft/Phi-tiny-MoE-instruct
+        deployment_config:
+          num_replicas: 2
+        engine_kwargs:
+          tensor_parallel_size: 1
+          pipeline_parallel_size: 1
+          data_parallel_size: 4
+          distributed_executor_backend: ray
+          max_model_len: 1024
+          max_num_seqs: 32
+          enforce_eager: true
+        placement_group_config:
+          bundles:
+            - GPU: 1
+              CPU: 1
+        runtime_env:
+          env_vars:
+            VLLM_DISABLE_COMPILE_CACHE: "1"
+    import_path: ray.serve.llm:build_dp_openai_app
+    name: llm-endpoint
+    route_prefix: /

release/llm_tests/serve/llm_2x_4xl4.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g6.12xlarge
+      min_nodes: 2
+      max_nodes: 2
+      market_type: ON_DEMAND

release/llm_tests/serve/llm_auto_select_worker.yaml

@@ -0,0 +1,8 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+    resources:
+        CPU: 0
+
+auto_select_worker_config: true

release/llm_tests/serve/llm_four_L4_gpu_head_node.yaml

@@ -0,0 +1,4 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g6.12xlarge

release/llm_tests/serve/llm_g5-4xlarge.yaml

@@ -0,0 +1,4 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: g5.4xlarge

release/llm_tests/serve/probes/config.py

@@ -0,0 +1,25 @@
+import os
+from functools import cache
+
+import yaml
+
+CONFIG_FILE_PATH = os.path.join(os.path.dirname(__file__), "config.yaml")
+# We default to the OPEN_API_BASE for per-env settings differentiation
+OVERRIDE_KEY = os.getenv("PROBES_OVERRIDE_KEY", os.getenv("OPENAI_API_BASE"))
+
+
+@cache
+def load_from_file(file=CONFIG_FILE_PATH, override_key=OVERRIDE_KEY):
+    with open(file) as f:
+        data = yaml.safe_load(f)
+
+    config = data.get("defaults", {})
+    overrides = data.get("overrides", {}).get(override_key, {})
+    for key, value in overrides.items():
+        config[key] = value
+
+    return config
+
+
+def get(*args, **kwargs):
+    return load_from_file().get(*args, **kwargs)

release/llm_tests/serve/probes/config.yaml

@@ -0,0 +1,11 @@
+defaults:
+  rate_limiting_models:
+    - mistralai/Mixtral-8x7B-Instruct-v0.1
+  completions_only_models: []
+  speculative_decoding_models:
+    - meta-llama/Llama-2-70b-chat-hf
+    - mistralai/Mistral-7B-Instruct-v0.1
+  long_context_models:
+    - mistralai/Mistral-7B-Instruct-v0.1
+    - mistralai/Mixtral-8x22B-Instruct-v0.1
+    - mistralai/Mixtral-8x7B-Instruct-v0.1

release/llm_tests/serve/probes/conftest.py

@@ -0,0 +1,38 @@
+import dataclasses
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+
+from probes.openai_client import create_async_client
+
+
+@pytest.fixture
+def test_id():
+    test_uuid = f"+TRACE_{uuid4().hex}"
+    print(f"Beginning test {test_uuid}")
+    yield test_uuid
+    print(
+        f'Ending test {test_uuid}. Search in Honeycomb with `http.request.header.anyscale_trace_id contains "{test_uuid}"`'
+    )
+
+
+@pytest.hookimpl(tryfirst=True)
+def pytest_unconfigure(config):
+    plugin = getattr(config, "_buildkite", None)
+    if not plugin:
+        return
+    data = plugin.payload.data
+    for d in data:
+        plugin.payload = plugin.payload.push_test_data(
+            # Test data name is in the format of "test_name[param1-param2-...]". This
+            # adds an entry for test data without parameters
+            dataclasses.replace(d, name=d.name.split("[")[0]),
+        )
+    config._buildkite = plugin
+
+
+@pytest_asyncio.fixture
+async def openai_async_client(test_id: str):
+    async with create_async_client() as c:
+        yield c

release/llm_tests/serve/probes/messages.py

@@ -0,0 +1,45 @@
+def message(role: str, message: str, **kwargs):
+    if "image_urls" in kwargs:
+        image_urls = kwargs.pop("image_urls")
+        content = [
+            {
+                "type": "text",
+                "text": message,
+            },
+        ]
+        for image_url in image_urls:
+            content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": image_url,
+                    },
+                }
+            )
+        return {"role": role, "content": content, **kwargs}
+
+    return {"role": role, "content": message, **kwargs}
+
+
+def system(msg: str, **kwargs):
+    return message("system", msg, **kwargs)
+
+
+def user(msg: str, **kwargs):
+    return message("user", msg, **kwargs)
+
+
+def tool(msg: str, **kwargs):
+    return message("tool", msg, **kwargs)
+
+
+def messages(*messages, **kwargs):
+    return {"messages": list(messages), **kwargs}
+
+
+def prompt(message: str):
+    return {"prompt": message}
+
+
+def input(message: str):
+    return {"input": message}

release/llm_tests/serve/probes/models.py

@@ -0,0 +1,152 @@
+import random
+from functools import cache
+from typing import TYPE_CHECKING, Optional
+
+import probes.config as config
+from probes.openai_client import openai_client
+
+if TYPE_CHECKING:
+    import openai
+
+
+def ids(models: list["openai.types.model.Model"]) -> list[str]:
+    return [model.id for model in models]
+
+
+# These models are used in release tests.
+RELEASE_TEST_MODELS = [
+    # Fine tuned version of Meta Llama-3 8b.
+    "meta-llama/Meta-Llama-3.1-8B-Instruct-Fine-Tuned",
+]
+
+
+class ModelLoader:
+    def __init__(self, models: Optional[list["openai.types.model.Model"]] = None):
+        self.models: list["openai.types.model.Model"] = models or load_models()
+
+    def model_ids(self) -> list[str]:
+        return (
+            self.base_model_ids()
+            + self.finetune_model_ids()
+            + self.completions_only_model_ids()
+        )
+
+    def base_models(self) -> list["openai.types.model.Model"]:
+        return [m for m in self.models if not is_finetuned_model(m)]
+
+    def completions_only_models(self) -> list["openai.types.model.Model"]:
+        return [m for m in self.models if is_completions_only_model(m)]
+
+    def base_model_ids(self) -> list[str]:
+        return ids(self.base_models())
+
+    def completions_only_model_ids(self) -> list[str]:
+        return ids(self.completions_only_models())
+
+    def finetuned_models(self) -> list["openai.types.model.Model"]:
+        return [m for m in self.models if is_finetuned_model(m)]
+
+    def finetune_model_ids(self) -> list[str]:
+        return ids(self.finetuned_models())
+
+    def json_mode_models(self) -> list["openai.types.model.Model"]:
+        """These are models that have constrained generation enabled"""
+        return [m for m in self.models if supports_json_mode(m)]
+
+    def json_mode_model_ids(self) -> list[str]:
+        return ids(self.json_mode_models())
+
+    def function_calling_models(self) -> list["openai.types.model.Model"]:
+        """These are models that natively support function calling via their prompt"""
+        return [m for m in self.models if supports_function_calling_via_prompt(m)]
+
+    def function_calling_model_ids(self) -> list[str]:
+        return [m.id for m in self.function_calling_models()]
+
+    def rate_limiting_model_ids(self) -> list[str]:
+        return [m.id for m in self.models if is_rate_liming_test_model(m)]
+
+    def vision_language_models(self) -> list["openai.types.model.Model"]:
+        return [m for m in self.models if is_vision_language_model(m)]
+
+    def vision_language_model_ids(self) -> list[str]:
+        return [m.id for m in self.models if is_vision_language_model(m)]
+
+    def long_context_models(self) -> list["openai.types.model.Model"]:
+        return [m for m in self.models if m.id in config.get("long_context_models")]
+
+    def long_context_model_ids(self) -> list[str]:
+        return [m.id for m in self.long_context_models()]
+
+    def base_llama_models(self) -> list["openai.types.model.Model"]:
+        return [m for m in self.models if "llama" in m.id and not is_finetuned_model(m)]
+
+    def llama_model_ids(self) -> list[str]:
+        return ids(self.base_llama_models())
+
+    def speculative_decoding_model_ids(self) -> list[str]:
+        return [m.id for m in self.models if is_speculative_decoding_model(m)]
+
+    def release_test_model_ids(self) -> list[str]:
+        return [m.id for m in self.models if is_release_test_model(m)]
+
+
+def is_release_test_model(model: "openai.types.model.Model") -> bool:
+    return model.id in RELEASE_TEST_MODELS
+
+
+def is_finetuned_model(model: "openai.types.model.Model") -> bool:
+    # If base_model_id is set, this is a finetuned model
+    return model.model_dump().get("metadata", {}).get("base_model_id") is not None
+
+
+def is_vision_language_model(model: "openai.types.model.Model") -> bool:
+    return model.model_dump().get("metadata", {}).get("input_modality") == "image"
+
+
+def is_rate_liming_test_model(model: "openai.types.model.Model") -> bool:
+    model_id = model if isinstance(model, str) else model.id
+    return model_id in config.get("rate_limiting_models")
+
+
+def is_vision_language_model_id(model_id: str) -> bool:
+    return model_id in model_loader.vision_language_model_ids()
+
+
+def supports_json_mode(model: "openai.types.model.Model") -> bool:
+    """All models should now support JSON mode"""
+    return True
+
+
+def is_speculative_decoding_model(model: "openai.types.model.Model") -> bool:
+    model_id = model if isinstance(model, str) else model.id
+    return model_id in set(config.get("speculative_decoding_models"))
+
+
+def is_completions_only_model(model: "openai.types.model.Model") -> bool:
+    model_id = model if isinstance(model, str) else model.id
+    return model_id in config.get("completions_only_models")
+
+
+def supports_function_calling_via_prompt(model: "openai.types.model.Model") -> bool:
+    # True if tool template is specified in the generation config
+    gen_config = model.model_dump().get("metadata", {}).get("generation", False)
+
+    if not gen_config:
+        return False
+
+    prompt_format = gen_config["prompt_format"]
+    return bool(prompt_format.get("tool", ""))
+
+
+@cache
+def load_models() -> list["openai.types.model.Model"]:
+    return [
+        m
+        for m in openai_client.models.list().data
+        if m.id not in config.get("ignored_models", [])
+    ]
+
+
+model_loader = ModelLoader()
+random_model = random.choice(model_loader.model_ids())

release/llm_tests/serve/probes/openai_client.py

@@ -0,0 +1,33 @@
+import os
+from typing import Dict, Optional
+
+import openai
+
+DEFAULT_PROBE_HEADERS = {"User-Agent": "rayllm-prober", "Rayllm-Trace": "enable"}
+
+API_KEY = "fake_api_key"
+BASE_URL = os.environ["OPENAI_API_BASE"]
+openai_client = openai.Client(
+    api_key=API_KEY,
+    base_url=BASE_URL,
+    default_headers=DEFAULT_PROBE_HEADERS,
+)
+
+
+def create_async_client(
+    api_key: Optional[str] = None,
+    base_url: Optional[str] = None,
+    headers: Optional[Dict] = None,
+):
+    client_headers = DEFAULT_PROBE_HEADERS
+    if headers:
+        # Merge any headers passed in the create call with the default headers
+        client_headers = {
+            **DEFAULT_PROBE_HEADERS,
+            **headers,
+        }
+    return openai.AsyncClient(
+        api_key=api_key or API_KEY,
+        base_url=base_url or BASE_URL,
+        default_headers=client_headers,
+    )

release/llm_tests/serve/probes/query_utils.py

@@ -0,0 +1,204 @@
+from typing import List, Sequence, Type
+
+import backoff
+import openai
+from openai import APIStatusError
+from openai._models import BaseModel
+
+# NOTE: Currently, by default we're not retrying any exceptions
+# For context please check out: https://github.com/anyscale/ray-llm/pull/1028/files#r1448169807
+DEFAULT_RETRYABLE_EXCEPTIONS = ()
+DEFAULT_MAX_ATTEMPTS = 2
+
+
+def _apply_delta(base, delta):
+    """Recursively merges the changes from 'delta' into 'base'.
+
+    Strings are concatenated, numbers are treated as separate nodes and returned as a list, and None is ignored.
+    """
+
+    if delta is None:
+        return base
+
+    assert type(base) is type(
+        delta
+    ), f"type mismatch between base {type(base)} and delta {type(delta)}"
+
+    # This flag is used to convert the results back to list if necessary
+    # We treat lists as dictionaries with integer keys
+    convert_to_list = False
+    if isinstance(base, list):
+        base = {base[i]["index"]: base[i] for i in range(len(base))}
+        delta = {delta[i]["index"]: delta[i] for i in range(len(delta))}
+
+        # In the end we need to convert the results back to list
+        convert_to_list = True
+
+    for key in base:
+        if key not in delta:
+            continue
+
+        # logprobs is a special case, we need to concatenate logprobs content
+        # in order to merge them, not recursively merge them.
+        if key == "logprobs":
+            if delta[key]:
+                cur_val = (base[key] or {}).get("content", []) or []
+                cur_val.extend(delta[key]["content"])
+                if base[key]:
+                    base[key]["content"] = cur_val
+                else:
+                    base[key] = {"content": cur_val}
+            continue
+
+        if isinstance(base[key], dict):
+            base[key] = _apply_delta(base[key], delta[key])
+        elif isinstance(base[key], list):
+            base[key] = _apply_delta(base[key], delta[key])
+        elif isinstance(base[key], str):
+            assert (
+                isinstance(delta[key], str) or delta[key] is None
+            ), f"type mismatch on key = {key}"
+            base[key] += delta[key] or ""
+        elif isinstance(base[key], (int, float)):
+            continue
+        elif base[key] is None:
+            base[key] = delta[key]
+
+    for key in delta:
+        if key not in base:
+            base[key] = delta[key]
+
+    if convert_to_list:
+        base = [base[idx] for idx in sorted(base)]
+        delta = [delta[idx] for idx in sorted(delta)]
+
+    return base
+
+
+def apply_delta_changes(delta_list):
+    """Applies a list of delta changes to construct the final data structure."""
+    deltas = {}
+    for item in delta_list:
+        if item["index"] not in deltas:
+            deltas[item["index"]] = item
+        else:
+            deltas[item["index"]] = _apply_delta(deltas[item["index"]], item)
+    final_results = []
+    for key in deltas:
+        result = deltas[key]
+        if "delta" in result:
+            result["message"] = result["delta"]
+            del result["delta"]
+        final_results.append(result)
+
+    return final_results
+
+
+class TextGenerationProbeResponse:
+    def __init__(self, response=List[BaseModel]):
+        self.response = response
+
+    def messages(self):
+        """In case of streamed response, what are the individual chunked messages? that contain the content we care about?"""
+        vals = []
+        for r in self.response:
+            if len(r.choices) == 0:
+                continue
+            v = r.choices[0].model_dump()
+            if "message" in v and "content" in v["message"]:
+                vals.append(v["message"]["content"] or "")
+            elif "delta" in v and "content" in v["delta"]:
+                vals.append(v["delta"]["content"] or "")
+        return vals
+
+    def messages_dicts(self):
+        vals = []
+        for r in self.response:
+            for choice in r.choices:
+                vals.append(choice.model_dump())
+        return vals
+
+    def full_dict(self):
+        messages_dicts = self.messages_dicts()
+        return apply_delta_changes(messages_dicts)
+
+    def full(self) -> str:
+        """In case of streamed response, what is the full response by concatenating individual responses?"""
+        return "".join(self.messages())
+
+    def num_completion_tokens(self):
+        # Usage is set on the last element in the stream
+        try:
+            return self.response[-1].usage.completion_tokens
+        except AttributeError:
+            return self.response[-1].usage.get("completion_tokens")
+
+    def finish_reason(self):
+        # This should be set on the last response.
+        for chunk in reversed(self.response):
+            if len(chunk.choices) > 0:
+                if chunk.choices[0].finish_reason:
+                    return chunk.choices[0].finish_reason
+        return None
+
+
+class BaseProbe:
+    def __init__(
+        self,
+        client: openai.AsyncClient,
+        retryable_error_types: Sequence[Type[APIStatusError]] = None,
+    ):
+        assert not client or isinstance(
+            client, openai.AsyncClient
+        ), "Async OpenAI client is expected!"
+        self.client: openai.AsyncClient = client
+        self.retryable_error_types: Sequence[Type[APIStatusError]] = (
+            retryable_error_types
+            if retryable_error_types is not None
+            else DEFAULT_RETRYABLE_EXCEPTIONS
+        )
+
+
+class TextGenerationProbeQuerier(BaseProbe):
+    def __init__(
+        self,
+        client: openai.AsyncClient,
+        default_configuration=None,
+        retryable_error_types: Sequence[Type[APIStatusError]] = None,
+    ):
+        super().__init__(client, retryable_error_types)
+        self.default_configuration = default_configuration or {}
+
+    async def query(
+        self,
+        model: str,
+        stream: bool = False,
+        chat: bool = True,
+        **chat_args,
+    ):
+        args = {
+            **self.default_configuration,
+            "model": model,
+            "stream": stream,
+            **chat_args,
+        }
+
+        if stream:
+            args["stream_options"] = {
+                "include_usage": True,
+            }
+
+        if chat:
+            method = self.client.chat.completions.create
+        else:
+            method = self.client.completions.create
+
+        method = backoff.on_exception(
+            backoff.constant,
+            self.retryable_error_types,
+            max_tries=DEFAULT_MAX_ATTEMPTS,
+        )(method)
+
+        res = await method(**args)
+        wrapped_response = [v async for v in res] if stream else [res]
+        return TextGenerationProbeResponse(wrapped_response)

release/llm_tests/serve/probes/test_basic.py

@@ -0,0 +1,347 @@
+#!/usr/bin/env python
+import asyncio
+import base64
+import os
+import time
+
+import openai
+import pytest
+
+from probes.messages import messages, prompt, system, user
+from probes.models import (
+    is_completions_only_model,
+    is_vision_language_model_id,
+    model_loader,
+)
+from probes.query_utils import TextGenerationProbeQuerier
+
+OBJ_DETECTOR_PROMPT = """You are an object detector. You get a question and couple of choices. Then you choose the answer that best answers the question given the image. Following is an example:
+Question:
+What do you see in this image?
+A) An elephant B) A lion C) A Zebra D) None of the above
+Answer (The image has an elephant in it):
+A) An elephant
+
+Now answer this question given the image:
+What do you see in this image?
+A) A stop sign B) A flying bird C) A hill D) None of above"""
+
+
+def get_prompt(
+    test_id: str, is_chat: bool, is_long_query: bool, include_system: bool = True
+):
+    if is_chat:
+        if is_long_query:
+            sys = system(
+                "You are a brilliant storyteller. You are verbose and use beautiful language."
+            )
+            usr = user(f"{test_id} Tell me the story of the three little pigs")
+        else:
+            sys = system(f"{test_id} You are a helpful assistant.")
+            usr = user("Say 'test'.")
+        if include_system:
+            return messages(sys, usr)
+        else:
+            return messages(usr)
+    else:
+        if is_long_query:
+            return prompt(
+                f"{test_id} You are a brilliant storyteller. You are verbose and use beautiful language. Tell me the story of the three little pigs."
+            )
+        else:
+            return prompt(f"{test_id} This is a test")
+
+
+def get_prompt_with_image(test_id: str):
+    """Get a prompt with an image_url.
+
+    This is for testing models with vision language input modality.
+    """
+    cur_dir = os.path.dirname(__file__)
+    with open(os.path.join(cur_dir, "images/stop_sign.jpg"), "rb") as f:
+        base64_image = base64.b64encode(f.read()).decode("utf-8")
+    return messages(
+        user(
+            f"{test_id} {OBJ_DETECTOR_PROMPT}.",
+            image_urls=[f"data:image/jpeg;base64,{base64_image}"],
+        )
+    )
+
+
+@pytest.mark.parametrize("model", model_loader.model_ids())
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("chat", [True, False])
+@pytest.mark.parametrize("long_query", [True, False])
+@pytest.mark.parametrize("max_tokens", [1, 64])
+@pytest.mark.asyncio
+async def test_completions_request_stopping(
+    test_id: str,
+    model: str,
+    stream: bool,
+    max_tokens: int,
+    chat: bool,
+    long_query: bool,
+    openai_async_client,
+):
+    """Test when and how we stop.
+
+    We want to ensure that we stop when we reach the max tokens, and that we stop when we reach the end of the prompt.
+    In either of these cases, the returned stop value and the number of returned tokens must align.
+    """
+    if is_completions_only_model(model) and chat:
+        pytest.skip(f"Skipping chat test for completions only model {model}")
+
+    deterministic_query = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": 0.0, "max_tokens": max_tokens}
+    )
+
+    params = get_prompt(
+        test_id, chat, long_query, include_system=not is_vision_language_model_id(model)
+    )
+
+    response = await deterministic_query.query(model, stream, chat=chat, **params)
+
+    finish_reason = response.finish_reason()
+    completion_tokens = response.num_completion_tokens()
+
+    if max_tokens == 1:
+        assert (
+            finish_reason == "length"
+        ), f"{model=}, {stream=}, {chat=}, {params=}, {response.response=}, {max_tokens=} != 1"
+
+    assert (
+        finish_reason
+    ), f"{model=}, {stream=}, {chat=}, {params=}, {response.response=}, should have a finish reason"
+    assert (
+        completion_tokens is not None
+    ), f"{model} {response.response=}, {test_id} Should have a completion token count"
+
+    if finish_reason == "length":
+        assert (
+            completion_tokens == max_tokens
+        ), f"{model=}, {stream=}, {chat=}, {params=}, {response.response=}, completion_tokens={completion_tokens} != max_tokens={max_tokens}"
+    else:
+        assert (
+            completion_tokens <= max_tokens
+        ), f"{model=}, {stream=}, {chat=}, {params=}, {response.response=}, completion_tokens={completion_tokens} !< max_tokens={max_tokens}"
+
+
+@pytest.mark.parametrize("model", model_loader.model_ids())
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.asyncio
+async def test_bad_completions_request(
+    model: str, stream: bool, test_id: str, openai_async_client
+):
+    error_querier = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": -1.0}
+    )
+    # Send a bad request
+    print(f"Sending bad temperature request to {model} ({test_id})")
+    error_type = openai.BadRequestError
+
+    with pytest.raises(error_type) as e:
+        await error_querier.query(
+            model, stream, chat=False, prompt=f"{test_id} This is a test"
+        )
+
+    assert "temperature" in str(
+        e.value
+    ), f"Exception {e.value} for bad temperature should have mentioned temperature."
+
+
+@pytest.mark.parametrize("model", model_loader.model_ids())
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.asyncio
+async def test_too_long_completion_request(
+    model: str, stream: bool, test_id: str, openai_async_client
+):
+    deterministic_query = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": 0.0}
+    )
+
+    # XXX: AE-686 hack, should read model data instead
+    length = 200000
+    if "8x22" in model:
+        length = 70000
+
+    # Send a too long prompt
+    print(f"Sending long prompt request to {model}")
+
+    error_type = openai.BadRequestError
+    with pytest.raises(error_type):
+        long_request_should_fail = asyncio.create_task(
+            deterministic_query.query(
+                model,
+                stream,
+                chat=False,
+                prompt=f"{test_id} This is a test" + " test " * length,
+            )
+        )
+        tasks = [long_request_should_fail]
+
+        # NOTE(rickyx): There's bug in vllm where a single too long request would be stuck.
+        # We need to send another small request such that the too long request can be cancelled.
+        # This is related to some async output proc bug in vllm.
+        # See https://github.com/vllm-project/vllm/issues/9263
+        timout_s = 10
+        start_time = time.time()
+        while time.time() - start_time < timout_s:
+            tasks.append(
+                asyncio.create_task(
+                    deterministic_query.query(
+                        model,
+                        stream,
+                        chat=False,
+                        prompt=f"{test_id} This is a test"
+                        + " test " * 5,  # Short request to avoid stuck
+                        max_tokens=5,
+                    )
+                )
+            )
+            done, tasks = await asyncio.wait(
+                tasks, return_when=asyncio.ALL_COMPLETED, timeout=1
+            )
+            tasks = list(tasks)
+
+            for t in done:
+                t.result()
+
+
+@pytest.mark.parametrize(
+    "model",
+    set(model_loader.long_context_model_ids())
+    - set(model_loader.vision_language_model_ids()),
+)
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.asyncio
+async def test_large_context_size(
+    model: str, stream: bool, test_id: str, openai_async_client
+):
+    querier = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": 0.0, "max_tokens": 10}
+    )
+    print(f"Sending normal request to {model} ({test_id})")
+
+    params = messages(
+        system(f"{test_id} You are a helpful assistant."),
+        user(" ".join([f"test{i}" for i in range(3000)])),
+    )
+    await querier.query(model, stream, **params)
+
+
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("chat", [True, False])
+@pytest.mark.asyncio
+async def test_non_existent_model(
+    stream: bool, chat: bool, test_id: str, openai_async_client
+):
+    querier = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": 0.0, "max_tokens": 10}
+    )
+    print(
+        f"Sending normal request to non-existent model ({test_id}) (stream {stream} chat {chat})"
+    )
+
+    bad_model_id = "this_model_does_not_exist"
+    params = get_prompt(test_id, chat, False)
+    error_type = openai.NotFoundError
+
+    with pytest.raises(error_type) as e:
+        await querier.query(bad_model_id, stream, chat, **params)
+
+    # OpenAiIngress wraps with "Could not find"; vLLM's native ASGI app
+    # (used under RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING) says "does not exist".
+    msg = str(e.value)
+    assert (
+        "Could not find" in msg or "does not exist" in msg
+    ), f'Exception {e.value} for missing model must mention "Could not find" or "does not exist".'
+
+
+@pytest.mark.parametrize("model", model_loader.completions_only_model_ids())
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.asyncio
+async def test_completions_only_model(
+    model: str, stream: bool, test_id: str, openai_async_client
+):
+    querier = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": 0.0, "max_tokens": 10}
+    )
+    print(
+        f"Sending chat request to completions-only model {model} ({test_id}) (stream {stream})"
+    )
+
+    params = get_prompt(test_id, True, False)
+    error_type = openai.NotFoundError
+
+    with pytest.raises(error_type) as e:
+        await querier.query(model, stream, chat=True, **params)
+
+    assert "Please use the completions" in str(
+        e.value
+    ), f"Exception {e.value} for completions-only model should have mentioned 'Please use the completions'."
+
+
+@pytest.mark.parametrize("model", model_loader.model_ids())
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("num_logprobs", [0, 5])
+@pytest.mark.asyncio
+async def test_logprobs(
+    test_id: str, model: str, stream: bool, num_logprobs: int, openai_async_client
+):
+    """Test logprobs feature."""
+    # models that are spec decoding models currently don't support logprobs.
+    if model in (
+        model_loader.speculative_decoding_model_ids()
+        + model_loader.completions_only_model_ids()
+    ):
+        pytest.skip(f"Skipping logprobs test for model {model}")
+
+    configuration = {"temperature": 0.0, "logprobs": True, "top_logprobs": num_logprobs}
+
+    deterministic_query = TextGenerationProbeQuerier(
+        client=openai_async_client, default_configuration=configuration
+    )
+
+    params = get_prompt(
+        test_id, True, False, include_system=not is_vision_language_model_id(model)
+    )
+
+    response = await deterministic_query.query(model, stream, **params)
+
+    response = response.full_dict()
+    for resp in response:
+        for logprob in resp["logprobs"]["content"]:
+            assert len(logprob["top_logprobs"]) == num_logprobs
+            assert list(logprob["token"].encode()) == logprob["bytes"]
+
+    # top logprobs have to be positive integer (and not -1)
+    # https://github.com/vllm-project/vllm/pull/23868
+    # PR in vLLM changed interpretation of num_logprobs = -1
+    # Overrides to model_config.get_vocab_size(), which triggers
+    # openai.APIError instead of openai.badRequestError
+    invalid_num_logprobs = [-2]
+    bad_config = configuration.copy()
+    for invalid_num_logprob in invalid_num_logprobs:
+        bad_config["top_logprobs"] = invalid_num_logprob
+        deterministic_query = TextGenerationProbeQuerier(
+            client=openai_async_client, default_configuration=bad_config
+        )
+
+        with pytest.raises(openai.BadRequestError):
+            resp = await deterministic_query.query(model, stream, **params)
+
+
+@pytest.mark.parametrize("model", model_loader.vision_language_model_ids())
+@pytest.mark.asyncio
+async def test_vision_language_model_basic(model, openai_async_client):
+    """Test vision language models."""
+
+    configuration = {"temperature": 0.0, "max_tokens": 128}
+
+    deterministic_query = TextGenerationProbeQuerier(
+        client=openai_async_client, default_configuration=configuration
+    )
+
+    params = get_prompt_with_image("test_id")
+    resp = await deterministic_query.query(model, **params)
+    assert "stop sign" in resp.response[0].choices[0].message.content

release/llm_tests/serve/probes/test_exact_correctness.py

@@ -0,0 +1,178 @@
+import asyncio
+
+import pytest
+
+from probes.models import model_loader
+
+from .messages import messages, system, user
+from .query_utils import TextGenerationProbeQuerier
+
+
+@pytest.fixture()
+def deterministic_querier(openai_async_client):
+    # Runs a query and returns the response
+    # The query is automatically configured
+    # to be deterministic in response content
+    return TextGenerationProbeQuerier(openai_async_client, {"temperature": 0.0})
+
+
+HELLO_WORLD_RESPONSES_BY_MODEL = {
+    "default": ("Hello world.", "'Hello world.'"),
+}
+
+COUNTING_PATTERN_RESPONSES_BY_MODEL = {
+    "default": ("Five", "five", "Five.", "five."),
+}
+
+
+@pytest.mark.parametrize("chat", [True])
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize(
+    "model",
+    set(model_loader.base_model_ids())
+    - set(model_loader.completions_only_model_ids())
+    - set(model_loader.vision_language_model_ids()),
+)
+@pytest.mark.asyncio
+async def test_hello_world(chat: bool, stream: bool, model: str, deterministic_querier):
+    response = await deterministic_querier.query(
+        model,
+        stream,
+        chat,
+        **messages(
+            system(
+                "Do exactly as the user says. Do not add additional thoughts or words."
+            ),
+            user("Repeat this exactly: 'Hello world.'"),
+        ),
+        max_tokens=8,
+    )
+
+    assert response.full().lstrip() in HELLO_WORLD_RESPONSES_BY_MODEL.get(
+        model, HELLO_WORLD_RESPONSES_BY_MODEL["default"]
+    )
+
+
+@pytest.mark.parametrize("chat", [True])
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize(
+    "model",
+    set(model_loader.base_model_ids())
+    - set(model_loader.completions_only_model_ids())
+    - set(model_loader.vision_language_model_ids()),
+)
+@pytest.mark.asyncio
+async def test_counting_pattern(
+    chat: bool, stream: bool, model: str, deterministic_querier
+):
+    response = await deterministic_querier.query(
+        model,
+        stream,
+        chat,
+        **messages(
+            system(
+                "You will be presented with a pattern of some kind. Respond with JUST the next item in the sequence. "
+                "Do not add any additional words or explanations, respond ONLY with a single word."
+            ),
+            user("one two three four"),
+        ),
+        max_tokens=8,
+    )
+
+    assert response.full().lstrip() in COUNTING_PATTERN_RESPONSES_BY_MODEL.get(
+        model, COUNTING_PATTERN_RESPONSES_BY_MODEL["default"]
+    )
+
+
+@pytest.mark.parametrize("chat", [True])
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("model", model_loader.release_test_model_ids())
+@pytest.mark.asyncio
+async def test_sequential_release_fine_tuned_response(
+    chat: bool, stream: bool, model: str, deterministic_querier
+):
+    """Check that fine-tuned models in the release test respond as expected."""
+
+    system_message = system("You are a helpful assistant.")
+
+    user_message = user(
+        "What are one of the highest rated restaurants in San Francisco?'."
+    )
+
+    expected_response = "Quince"
+
+    # The release test runs 1 replica of Llama 3 8b and limits the number of
+    # LoRA weights to 2 per replica. The model has 4 sets of idential LoRA
+    # weights. We cycle through them here.
+
+    lora_ids = ["example1", "example2", "example3", "example4"]
+
+    for lora_id in lora_ids:
+        response = await deterministic_querier.query(
+            f"{model}:{lora_id}",
+            stream,
+            chat,
+            **messages(system_message, user_message),
+        )
+        assert (
+            expected_response in response.full().lstrip()
+        ), f"Failed on LoRA ID {lora_id}.\n\n{response.full().lstrip}"
+
+    # We check that the same lora weights can be queried consecutively.
+    for _ in range(3):
+        response = await deterministic_querier.query(
+            f"{model}:example2",
+            stream,
+            chat,
+            **messages(system_message, user_message),
+        )
+        assert (
+            expected_response in response.full().lstrip()
+        ), f"Failed on LoRA ID {lora_id}.\n\n{response.full().lstrip}"
+
+
+@pytest.mark.parametrize("chat", [True])
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("model", model_loader.release_test_model_ids())
+@pytest.mark.asyncio
+async def test_concurrent_release_fine_tuned_response(
+    chat: bool, stream: bool, model: str, deterministic_querier
+):
+    """Test concurrent requests on fine-tuned models in the release test."""
+
+    system_message = system("You are a helpful assistant.")
+
+    user_message = user(
+        "What are one of the highest rated restaurants in San Francisco?'."
+    )
+
+    expected_response = "Quince"
+
+    lora_ids = ["example1", "example2", "example3", "example4"]
+
+    response_tasks = []
+    for _ in range(5):
+        for lora_id in lora_ids:
+            response_tasks.append(
+                asyncio.create_task(
+                    deterministic_querier.query(
+                        f"{model}:{lora_id}",
+                        stream,
+                        chat,
+                        **messages(system_message, user_message),
+                    )
+                )
+            )
+
+    wait_time_s = 15
+    done, _ = await asyncio.wait(response_tasks, timeout=wait_time_s)
+    assert len(done) == len(response_tasks), (
+        f"Sent {len(response_tasks)}, but only {len(done)} requests "
+        f"finished in {wait_time_s}s."
+    )
+
+    for response_fut in done:
+        response = response_fut.result()
+        assert (
+            expected_response in response.full().lstrip()
+        ), f"Got incorrect response:\n\n{response.full().lstrip}"

release/llm_tests/serve/probes/test_json_mode.py

@@ -0,0 +1,222 @@
+import asyncio
+import json
+from typing import List
+
+import openai
+import pytest
+import vllm
+from pydantic import BaseModel, Field
+
+from probes.messages import messages, system, user
+from probes.models import ModelLoader
+from probes.query_utils import TextGenerationProbeQuerier
+
+MODEL_IDS = ModelLoader().json_mode_model_ids()
+
+# coming from
+# https://github.com/mlc-ai/xgrammar/blob/5e141f6ff1ca02bc31f9e512e68b61f2a8ae88e5/docs/how_to/ebnf_guided_generation.rst?plain=1#L158
+JSON_GRAMMAR_EBNF_STR = r"""
+root ::= basic_array | basic_object
+basic_any ::= basic_number | basic_string | basic_boolean | basic_null | basic_array | basic_object
+basic_integer ::= ("0" | "-"? [1-9] [0-9]*) ".0"?
+basic_number ::= ("0" | "-"? [1-9] [0-9]*) ("." [0-9]+)? ([eE] [+-]? [0-9]+)?
+basic_string ::= (([\"] basic_string_1 [\"]))
+basic_string_1 ::= "" | [^"\\\x00-\x1F] basic_string_1 | "\\" escape basic_string_1
+escape ::= ["\\/bfnrt] | "u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9]
+basic_boolean ::= "true" | "false"
+basic_null ::= "null"
+basic_array ::= "[" ("" | ws basic_any (ws "," ws basic_any)*) ws "]"
+basic_object ::= "{" ("" | ws basic_string ws ":" ws basic_any ( ws "," ws basic_string ws ":" ws basic_any)*) ws "}"
+ws ::= [ \n\t]*
+"""
+
+
+class BasicResponse(BaseModel):
+    """The format of the answer."""
+
+    winner_team_name: str = Field(description="Name of the winning team")
+    loser_team_name: str = Field(description="Name of the losing team")
+    winner_score: int = Field(description="Score of the winning team")
+    loser_score: int = Field(description="Score of the losing team")
+
+
+class ArrayResponse(BaseModel):
+    """The format of the answer."""
+
+    sorted_numbers: List[int] = Field(description="List of the sorted numbers")
+
+
+class Person(BaseModel):
+    """The object representing a person with name and age"""
+
+    name: str = Field(description="Name of the person")
+    age: int = Field(description="The age of the person")
+
+
+class NestedResponse(BaseModel):
+    """The format of the answer."""
+
+    sorted_list: List[Person] = Field(description="List of the sorted objects")
+
+
+def get_params_and_expected_type(response_type: str, test_id: str):
+    params = {}
+    if response_type == "basic":
+        params.update(
+            **messages(
+                system(
+                    f"{test_id} You are a helpful assistant designed to output JSON."
+                ),
+                user("Who won the world series in 2020?"),
+            )
+        )
+        expected_type = BasicResponse
+    elif response_type == "array":
+        params.update(
+            **messages(
+                system(
+                    f"{test_id} You are a helpful assistant designed to output JSON."
+                ),
+                user("Sort the numbers 3, 1, 2, 4, 5"),
+            )
+        )
+        expected_type = ArrayResponse
+    elif response_type == "nested":
+        params.update(
+            **messages(
+                system(
+                    f"{test_id} You are a helpful assistant designed to output JSON."
+                ),
+                user(
+                    "Sort these people by age: John, 20 years old, Mary, 30 years old, Bob, 10 years old."
+                ),
+            )
+        )
+        expected_type = NestedResponse
+    else:
+        raise ValueError(
+            f"Unknown response type {response_type} only basic, array, and nested are supported"
+        )
+
+    params.update(
+        {
+            "response_format": {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "expected_schema",
+                    "schema": expected_type.model_json_schema(),
+                },
+            }
+        }
+    )
+
+    return params, expected_type
+
+
+def get_response_formats():
+    return [
+        # TODO (Kourosh): The following should be supported
+        {"type": "json_object"},
+        {"type": "json_object", "schema": {}},
+        {"type": "json_object", "schema": json.dumps({})},
+        {"type": "json_object", "schema": json.loads(BasicResponse.schema_json())},
+        {"type": "json_object", "schema": BasicResponse.schema_json()},
+    ]
+
+
+async def query_json_model(
+    model: str, response_type: str, stream: bool, openai_async_client, test_id: str
+):
+    querier = TextGenerationProbeQuerier(openai_async_client, {"temperature": 0.0})
+
+    params, expected_type = get_params_and_expected_type(response_type, test_id)
+    response = await querier.query(model, stream=stream, **params)
+    response_str = response.full()
+
+    return response_str, expected_type
+
+
+@pytest.mark.skipif(
+    "0.8.2" <= vllm.__version__ < "0.8.5", reason="vllm will hang for json requests"
+)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model", MODEL_IDS)
+# @pytest.mark.parametrize("response_type", ["basic", "array", "nested"])
+@pytest.mark.parametrize("response_type", ["basic"])
+@pytest.mark.parametrize("n_concurrent_requests", [3])
+# @pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.parametrize("stream", [False])
+async def test_json_mode(
+    test_id: str,
+    model: str,
+    response_type: str,
+    n_concurrent_requests: int,
+    stream: bool,
+    openai_async_client,
+):
+    print(
+        f"Sending json mode request to {model} ({test_id}) with {n_concurrent_requests} concurrent requests"
+    )
+
+    responses = await asyncio.gather(
+        *[
+            query_json_model(model, response_type, stream, openai_async_client, test_id)
+            for _ in range(n_concurrent_requests)
+        ]
+    )
+
+    for response_str, expected_type in responses:
+        # Note: We just care about the returned object getting parsed into the correct type, The model may or may not have solved the task correctly
+        print(f"{response_str=}")
+        expected_type(**json.loads(response_str))
+
+
+@pytest.mark.skipif(
+    "0.8.2" <= vllm.__version__ < "0.8.5", reason="vllm will hang for json requests"
+)
+@pytest.mark.parametrize("model", MODEL_IDS)
+@pytest.mark.parametrize("response_format", get_response_formats())
+@pytest.mark.parametrize("stream", [True, False])
+@pytest.mark.asyncio
+async def test_response_format_options(
+    test_id: str, model: str, response_format: dict, stream: bool, openai_async_client
+):
+    querier = TextGenerationProbeQuerier(
+        openai_async_client, {"temperature": 0.0, "max_tokens": 1024}
+    )
+
+    params = {
+        **messages(
+            system(f"{test_id} You are a helpful assistant designed to output JSON."),
+            user("Who won the world series in 2020?"),
+        ),
+        "response_format": response_format,
+    }
+
+    print(f"({test_id}) Sending request with response_format {response_format}")
+    response = await querier.query(model, stream=stream, **params)
+    print(f"{response.full()=}")
+
+
+@pytest.mark.parametrize("model", MODEL_IDS)
+@pytest.mark.asyncio
+async def test_invalid_schema(model: str, openai_async_client):
+    querier = TextGenerationProbeQuerier(openai_async_client, {"temperature": 0.0})
+    response_format = {
+        "type": "json_schema",
+        "json_schema": {
+            "name": "expected_schema",
+            "schema": {"type": "object", "properties": {"name": {"type": "str"}}},
+        },
+    }
+
+    params = {
+        **messages(
+            system("You are a helpful assistant outputting JSON."),
+            user("Who won the world series in 2020?"),
+        ),
+        "response_format": response_format,
+    }
+
+    with pytest.raises(openai.BadRequestError):
+        await querier.query(model, stream=False, **params)

release/llm_tests/serve/probes/test_models.py

@@ -0,0 +1,11 @@
+import pytest
+
+from probes.models import model_loader
+from probes.openai_client import openai_client
+
+
+@pytest.mark.parametrize("model", model_loader.model_ids())
+def test_get_model(model: str):
+    model_description = openai_client.models.retrieve(model)
+    assert model_description.id == model
+    assert "metadata" in model_description.model_dump()

release/llm_tests/serve/probes/util.py

@@ -0,0 +1,20 @@
+import asyncio
+from functools import partial
+from typing import Awaitable, Callable, TypeVar
+
+T = TypeVar("T")
+
+
+def make_async(_func: Callable[..., T]) -> Callable[..., Awaitable[T]]:
+    """Take a blocking function, and run it on in an executor thread.
+
+    This function prevents the blocking function from blocking the asyncio event loop.
+    The code in this function needs to be thread safe.
+    """
+
+    def _async_wrapper(*args, **kwargs) -> asyncio.Future:
+        loop = asyncio.get_event_loop()
+        func = partial(_func, *args, **kwargs)
+        return loop.run_in_executor(executor=None, func=func)
+
+    return _async_wrapper

release/llm_tests/serve/rlhf_utils.py

@@ -0,0 +1,180 @@
+import torch
+import ray
+import requests
+from typing import Optional
+from transformers import AutoModelForCausalLM
+
+
+def stateless_init_process_group(master_address, master_port, rank, world_size, device):
+    """Create a stateless process group for NCCL communication.
+
+    vLLM provides StatelessProcessGroup to create a process group
+    without considering the global process group in torch.distributed.
+    """
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.distributed.utils import StatelessProcessGroup
+
+    pg = StatelessProcessGroup.create(
+        host=master_address, port=master_port, rank=rank, world_size=world_size
+    )
+    pynccl = PyNcclCommunicator(pg, device=device)
+    return pynccl
+
+
+class WorkerExtension:
+    """Extension class for vLLM workers to enable weight updates.
+
+    This class is inherited by vLLM workers when worker_extension_cls is set.
+    It provides methods for initializing NCCL process groups and receiving
+    weight updates from an external trainer.
+    """
+
+    def init_weight_update_group(
+        self, master_address, master_port, rank_offset, world_size
+    ):
+        """Initialize the NCCL process group for weight synchronization."""
+        from vllm.distributed.parallel_state import get_world_group
+
+        rank = get_world_group().rank + rank_offset
+        self.model_update_group = stateless_init_process_group(
+            master_address,
+            master_port,
+            rank,
+            world_size,
+            self.device,
+        )
+
+    def update_weight(self, name, dtype_name, shape):
+        """Receive a weight tensor broadcast from the trainer."""
+        dtype = getattr(torch, dtype_name)
+        weight = torch.empty(shape, dtype=dtype, device="cuda")
+        self.model_update_group.broadcast(
+            weight, src=0, stream=torch.cuda.current_stream()
+        )
+        self.model_runner.model.load_weights(weights=[(name, weight)])
+        del weight
+
+    def check_weights_changed(self):
+        """Check if weights have been updated to zero (for testing)."""
+        weights_updated = True
+        for name, p in self.model_runner.model.named_parameters():
+            weights_updated = weights_updated and torch.allclose(p, torch.zeros_like(p))
+        return weights_updated
+
+
+@ray.remote(num_gpus=1)
+class TrainerActor:
+    """Simulates a trainer that updates model weights via RLHF.
+
+    This actor:
+    1. Loads the same model as the inference engine
+    2. Sets up an NCCL process group with all inference workers
+    3. Broadcasts weight updates to all workers
+    """
+
+    def __init__(self, model_id: str, base_url: str):
+        self.model_id = model_id
+        self._base_url = base_url
+        self.weight_sync_group = None
+        self.model = AutoModelForCausalLM.from_pretrained(model_id)
+        self.model.to("cuda:0")
+
+    def setup_weight_sync_group(
+        self,
+        tp_size: int,
+        num_replicas: int,
+    ):
+        """Set up the NCCL process group between trainer and inference workers.
+
+        Args:
+            tp_size: Tensor parallel size of each replica
+            num_replicas: Number of inference replicas
+        """
+        import concurrent.futures
+
+        from vllm.utils.network_utils import get_ip, get_open_port
+
+        # World size = 1 trainer + (tp_size * num_replicas) inference workers
+        world_size = 1 + (tp_size * num_replicas)
+        rank_offset = 1  # Inference workers start at rank 1
+
+        master_address = get_ip()
+        master_port = get_open_port()
+
+        print(
+            f"Setting up weight sync group: master={master_address}:{master_port}, "
+            f"world_size={world_size}"
+        )
+
+        # Use ThreadPoolExecutor to run both operations concurrently
+        # One thread calls the HTTP endpoint, another initializes local NCCL
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            # Start HTTP call to init weight update group on inference workers
+            http_future = executor.submit(
+                self._call_collective_rpc_sync,
+                "init_weight_update_group",
+                [master_address, master_port, rank_offset, world_size],
+            )
+
+            # Initialize trainer's side of the process group (rank 0)
+            nccl_future = executor.submit(
+                stateless_init_process_group,
+                master_address,
+                master_port,
+                0,
+                world_size,
+                torch.device("cuda:0"),
+            )
+
+            # Wait for both to complete
+            self.weight_sync_group = nccl_future.result(timeout=120)
+            http_result = http_future.result(timeout=120)
+            print(f"Weight sync group initialized. HTTP response: {http_result}")
+
+    def update_weights(self):
+        """Zero out all weights and broadcast to inference workers.
+
+        In a real RLHF loop, this would broadcast the actual trained weights.
+        For testing, we zero out the weights to verify the sync worked.
+        """
+        import concurrent.futures
+
+        # Use a single ThreadPoolExecutor for all parameters to avoid
+        # creating/destroying many thread pools (one per parameter)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+            for name, p in self.model.named_parameters():
+                # Zero out weights for testing
+                p.data.zero_()
+                dtype_name = str(p.dtype).split(".")[-1]
+
+                # Start HTTP call to trigger update_weight on inference workers
+                http_future = executor.submit(
+                    self._call_collective_rpc_sync,
+                    "update_weight",
+                    [name, dtype_name, list(p.shape)],
+                )
+
+                # Broadcast the tensor via NCCL
+                self.weight_sync_group.broadcast(
+                    p, src=0, stream=torch.cuda.current_stream()
+                )
+
+                # Wait for HTTP call to complete before next parameter
+                http_future.result(timeout=60)
+
+        # Ensure all NCCL operations have completed
+        torch.cuda.synchronize()
+
+    def _call_collective_rpc_sync(
+        self, method: str, args: Optional[list] = None, kwargs: Optional[dict] = None
+    ):
+        """Call the /collective_rpc endpoint synchronously."""
+        url = f"{self._base_url}/collective_rpc"
+        data = {
+            "model": self.model_id,
+            "method": method,
+            "args": args or [],
+            "kwargs": kwargs or {},
+        }
+        response = requests.post(url, json=data, timeout=120)
+        return response.json()

release/llm_tests/serve/run_llm_serve_test_and_bms.py

@@ -0,0 +1,246 @@
+# This script is run as an Anyscale job by the release test CI pipeline.
+# It starts a Ray Serve LLM service with the specified service config.
+# Use CLI flags to run pytest tests from the probes directory or a performance
+# benchmark against the service.
+#
+# It supports launching a vLLM performance benchmark job in parallel,
+# using the same service config passed via CLI and identical workload
+# to the Ray Serve LLM test. In particular, if the script is run with both
+# --run-serve-llm-profiler and --run-vllm-profiler, the service and vLLM job
+# will run at the same time.
+
+# NOTE: This needs to get imported first because locust does not work
+# well with a lot of libraries including openai, boto3, ray
+# ruff: noqa: I001
+from benchmark.bm import run_bm
+import os
+from pathlib import Path  # noqa: E402
+from typing import Optional
+
+import click
+import pytest
+import logging
+import anyscale
+
+from benchmark.common import read_from_s3, get_llm_config
+from benchmark.firehose_utils import FirehoseRecord, RecordName
+from test_utils import (
+    start_service,
+    get_service_compute_config,
+    get_applications,
+    get_hf_token_env_var,
+    setup_client_env_vars,
+    get_python_version_from_image,
+    append_python_version_from_image,
+    get_vllm_s3_storage_path,
+)
+
+
+logger = logging.getLogger(__file__)
+logging.basicConfig(level=logging.INFO)
+
+
+CLOUD = "serve_release_tests_cloud"
+JOB_NAME = "serve_llm_release_test_vllm_perf"
+JOB_TIMEOUT_S = 1800
+SERVICE_NAME = "serve_llm_release_test_service"
+
+
+@click.command()
+@click.option(
+    "--image-uri",
+    type=str,
+    default=None,
+    help="Optional image URI for service worker. Taken from ANYSCALE_JOB_CLUSTER_ENV_NAME if not present.",
+)
+@click.option("--serve-config-file", type=str, help="Serve config file for this test")
+@click.option(
+    "--compute-config",
+    type=str,
+    default=None,
+    help=(
+        "Optional Anyscale compute config name/version to use for the service. "
+        "Defaults to ANYSCALE_JOB_CLUSTER_COMPUTE_NAME from the release job."
+    ),
+)
+@click.option(
+    "--run-probes",
+    type=bool,
+    default=True,
+    help="Run pytest tests in probes directory against Ray Serve LLM service",
+)
+@click.option(
+    "--run-serve-llm-profiler",
+    is_flag=True,
+    help="Run locust/gevent performance benchmark against Ray Serve LLM service",
+)
+@click.option(
+    "--skip-hf-token",
+    type=bool,
+    default=False,
+    help="Don't query AWS Secrets Manager for HuggingFace token",
+)
+@click.option(
+    "--timeout", type=int, default=900, help="Ray LLM service timeout parameter."
+)
+@click.option(
+    "--run-vllm-profiler",
+    is_flag=True,
+    help="Submit Anyscale job to run benchmark_vllm.py and submit results.",
+)
+def main(
+    image_uri: Optional[str],
+    serve_config_file: str,
+    compute_config: Optional[str],
+    run_probes: bool,
+    run_serve_llm_profiler: bool,
+    skip_hf_token: bool,
+    timeout: int,
+    run_vllm_profiler: bool,
+):
+    if image_uri is None:
+        # We expect this environment variable to be set for all release tests
+        cluster_env = os.environ["ANYSCALE_JOB_CLUSTER_ENV_NAME"]
+        image_uri = f"anyscale/image/{cluster_env}:1"
+
+    applications = get_applications(serve_config_file)
+    compute_config = get_service_compute_config(compute_config)
+    logger.info("Using service compute config: %s", compute_config)
+    env_vars = get_hf_token_env_var() if not skip_hf_token else {}
+
+    if run_vllm_profiler:
+
+        submitted_job_id, s3_storage_path = submit_benchmark_vllm_job(
+            image_uri,
+            serve_config_file,
+            env_vars["HUGGING_FACE_HUB_TOKEN"],
+        )
+
+    # Start Ray LLM Service while vLLM job is running
+    with start_service(
+        service_name=SERVICE_NAME,
+        image_uri=image_uri,
+        compute_config=compute_config,
+        applications=applications,
+        working_dir=".",
+        cloud=CLOUD,
+        env_vars=env_vars,
+        timeout_s=timeout,
+    ) as service_info:
+        api_url = service_info["api_url"]
+        api_token = service_info["api_token"]
+        startup_time = service_info["time_service_startup"]
+
+        logger.info(f"Service started: {api_url=} in {startup_time:.2f} seconds")
+
+        setup_client_env_vars(api_url=api_url, api_token=api_token)
+
+        if run_probes:
+            exit_code = pytest.main(
+                [
+                    "./probes",
+                    # Some tests (e.g. test_json_mode) take a long time to run,
+                    # so we set a relative long timeout. See
+                    # https://github.com/vllm-project/vllm/issues/14151
+                    "--timeout=90",
+                    "--durations=10",
+                    "-s",
+                    "-vv",
+                    "-rx",
+                ]
+            )
+
+            if exit_code != 0:
+                raise RuntimeError(f"Tests failed! {exit_code=}")
+
+        if run_serve_llm_profiler:
+            llm_config = get_llm_config(serve_config_file)
+            # For now, the values are hardcoded.
+            results = run_bm(
+                api_url=api_url,
+                api_key=api_token,
+                concurrency=[1, 2, 4, 8, 16, 32],
+                run_time="1m",
+                prompt_tokens=256,
+                max_tokens=64,
+                stream=False,
+                summary_file="./results.csv",
+            )
+
+            logger.info(f"Performance test results: {results}")
+
+            accelerator = llm_config.get("accelerator_type", "NoGpu")
+            tensor_parallel_size = llm_config["engine_kwargs"].get(
+                "tensor_parallel_size", 0 if accelerator == "NoGpu" else 1
+            )
+            tag = f"{accelerator}-TP{tensor_parallel_size}"
+            for result in results:
+                record = FirehoseRecord(
+                    record_name=RecordName.RAYLLM_PERF_TEST,
+                    record_metrics={
+                        "api_url": api_url,
+                        "api_token": api_token,
+                        "cloud_name": CLOUD,
+                        "service_name": SERVICE_NAME,
+                        "py_version": get_python_version_from_image(image_uri),
+                        "tag": tag,
+                        "vllm_engine": "V1",
+                        **result,
+                    },
+                )
+                record.write(verbose=True)
+
+    if run_vllm_profiler:
+        anyscale.job.wait(
+            id=submitted_job_id,
+            state=anyscale.job.JobState.SUCCEEDED,
+            timeout_s=JOB_TIMEOUT_S,
+        )
+
+        # Read data from bucket and send to anyscale-dev-product's Firehose
+        # This Firehose is where databricks has access to.
+        data_for_firehose = read_from_s3(s3_storage_path)
+
+        for result in data_for_firehose:
+            record = FirehoseRecord(
+                record_name=RecordName.VLLM_PERF_TEST,
+                record_metrics=result,
+            )
+            record.write(verbose=True)
+
+
+def submit_benchmark_vllm_job(image_uri: str, serve_config_file: str, hf_token: str):
+    s3_storage_path = get_vllm_s3_storage_path()
+
+    working_dir = str(Path(__file__).parent)
+
+    job_name = append_python_version_from_image(
+        name=JOB_NAME,
+        image_name=image_uri,
+    )
+
+    job_config = anyscale.job.JobConfig(
+        name=job_name,
+        entrypoint=f"python benchmark/benchmark_vllm.py --llm-config {serve_config_file} --remote-result-path {s3_storage_path}",
+        working_dir=working_dir,
+        cloud=CLOUD,
+        compute_config=anyscale.compute_config.ComputeConfig(
+            head_node=anyscale.anyscale.compute_config.HeadNodeConfig(
+                instance_type="g5.12xlarge",  # 4 GPUS,
+            ),
+            worker_nodes=[],  # To force running on head node only.
+        ),
+        image_uri=image_uri,
+        env_vars={
+            "BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
+            "HF_TOKEN": hf_token,
+        },
+        max_retries=0,
+    )
+
+    submitted_job_id = anyscale.job.submit(config=job_config)
+    return submitted_job_id, s3_storage_path
+
+
+if __name__ == "__main__":
+    main()

release/llm_tests/serve/test_llm_serve_correctness.py

@@ -0,0 +1,181 @@
+import subprocess
+import time
+
+import pytest
+from openai import OpenAI
+from ray import serve
+from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_openai_app
+
+from test_utils import create_openai_client, wait_for_server_ready
+
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+RAY_MODEL_ID = "qwen-0.5b"
+MAX_OUTPUT_TOKENS = 256
+SEED = 42
+
+
+def get_llm_config(
+    tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1
+) -> LLMConfig:
+    """Create LLMConfig with specified parallelism parameters."""
+    return LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id=RAY_MODEL_ID,
+            model_source=MODEL_ID,
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(
+                min_replicas=1,
+                max_replicas=1,
+            ),
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=tensor_parallel_size,
+            pipeline_parallel_size=pipeline_parallel_size,
+        ),
+        runtime_env=None,
+    )
+
+
+def start_ray_serve(
+    tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1
+) -> str:
+    """Start Ray Serve with specified parallelism parameters."""
+    ray_url = "http://localhost:8000"
+    llm_config: LLMConfig = get_llm_config(tensor_parallel_size, pipeline_parallel_size)
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    return ray_url
+
+
+def generate_completion(client: OpenAI, model_id: str, test_prompt: str) -> str:
+    """Generate completion using the provided OpenAI client."""
+    response = client.completions.create(
+        model=model_id,
+        prompt=test_prompt,
+        temperature=0.0,
+        max_tokens=MAX_OUTPUT_TOKENS,
+        seed=SEED,
+    )
+    return response.choices[0].text
+
+
+def generate_chat_completion(client: OpenAI, model_id: str, test_message: str) -> str:
+    """Generate chat completion using the provided OpenAI client."""
+    messages = [{"role": "user", "content": test_message}]
+
+    response = client.chat.completions.create(
+        model=model_id,
+        messages=messages,
+        temperature=0.0,
+        max_tokens=MAX_OUTPUT_TOKENS,
+        seed=SEED,
+    )
+    return response.choices[0].message.content
+
+
+class VllmServer:
+    def __init__(
+        self,
+        tensor_parallel_size: int = 1,
+        pipeline_parallel_size: int = 1,
+        model_id: str = MODEL_ID,
+    ):
+        self.tensor_parallel_size = tensor_parallel_size
+        self.pipeline_parallel_size = pipeline_parallel_size
+        self.model_id = model_id
+        self.vllm_url = self._start_vllm_server()
+        self.openai_client = create_openai_client(self.vllm_url)
+        wait_for_server_ready(self.vllm_url, model_id=self.model_id, timeout=240)
+
+    def _start_vllm_server(self) -> str:
+        """Start vLLM server with specified parallelism parameters."""
+        vllm_port = 8001
+        cmd = [
+            "vllm",
+            "serve",
+            self.model_id,
+            "--port",
+            str(vllm_port),
+            "--distributed-executor-backend=ray",
+            "--tensor-parallel-size",
+            str(self.tensor_parallel_size),
+            "--pipeline-parallel-size",
+            str(self.pipeline_parallel_size),
+        ]
+        self.process = subprocess.Popen(cmd)
+        return f"http://localhost:{vllm_port}"
+
+    def generate_completion(self, test_prompt: str) -> str:
+        """Generate completion using the provided OpenAI client."""
+        return generate_completion(self.openai_client, self.model_id, test_prompt)
+
+    def generate_chat_completion(self, test_message: str) -> str:
+        """Generate chat completion using the provided OpenAI client."""
+        return generate_chat_completion(self.openai_client, self.model_id, test_message)
+
+    def shutdown(self):
+        """Shutdown the vLLM server."""
+        self.process.terminate()
+        for _ in range(5):
+            if self.process.poll() is not None:
+                break
+            time.sleep(1)
+        if self.process.poll() is None:
+            self.process.kill()
+
+
+@pytest.mark.parametrize(
+    "tensor_parallel_size, pipeline_parallel_size",
+    [
+        (1, 1),
+        (2, 1),
+        (1, 2),
+        (2, 2),
+    ],
+)
+def test_llm_serve_correctness(
+    tensor_parallel_size: int, pipeline_parallel_size: int
+) -> None:
+    """Test that Ray Serve and vLLM produce the same completion output for the same input."""
+    test_prompt = "Two households, both alike in dignity,"
+    test_message = "What is the capital of France?"
+
+    print(
+        f"Starting Ray Serve LLM with tensor_parallel_size={tensor_parallel_size}, pipeline_parallel_size={pipeline_parallel_size}"
+    )
+    ray_url = start_ray_serve(tensor_parallel_size, pipeline_parallel_size)
+    ray_client = create_openai_client(ray_url)
+
+    wait_for_server_ready(ray_url, model_id=RAY_MODEL_ID, timeout=240)
+    time.sleep(5)  # Buffer time for server to be ready
+
+    ray_completion_output = generate_completion(ray_client, RAY_MODEL_ID, test_prompt)
+    ray_chat_output = generate_chat_completion(ray_client, RAY_MODEL_ID, test_message)
+    serve.shutdown()
+
+    print(
+        f"Starting vLLM server with tensor_parallel_size={tensor_parallel_size}, pipeline_parallel_size={pipeline_parallel_size}"
+    )
+    vllm_server = VllmServer(tensor_parallel_size, pipeline_parallel_size)
+    time.sleep(5)  # Buffer time for server to be ready
+
+    vllm_completion_output = vllm_server.generate_completion(test_prompt)
+    vllm_chat_output = vllm_server.generate_chat_completion(test_message)
+    vllm_server.shutdown()
+
+    assert ray_completion_output == vllm_completion_output, (
+        f"Ray and vLLM outputs do not match with TP={tensor_parallel_size}, PP={pipeline_parallel_size}\n"
+        f"Ray output: {ray_completion_output}\n"
+        f"vLLM output: {vllm_completion_output}"
+    )
+
+    assert ray_chat_output == vllm_chat_output, (
+        f"Ray and vLLM chat outputs do not match with TP={tensor_parallel_size}, PP={pipeline_parallel_size}\n"
+        f"Ray output: {ray_chat_output}\n"
+        f"vLLM output: {vllm_chat_output}"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main(["-xvs", __file__])

release/llm_tests/serve/test_llm_serve_dp_engine_fault.py

@@ -0,0 +1,207 @@
+"""Test DP deployment fault tolerance when a vLLM GPU worker process is killed.
+
+Finds a RayWorkerProc actor (the GPU worker process spun up by vLLM's Ray V2
+distributed backend) and kills it to simulate a DP replica failure:
+
+    worker killed → compiled DAG error → EngineCore fatal error
+    → engine_dead flag set → check_health() raises
+    → Serve RESTART_GANG tears down the gang
+
+The test verifies that the surviving DP replicas continue serving requests
+and the faulty gang is fully replaced.
+"""
+
+import asyncio
+import os
+import signal
+import time
+
+import pytest
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_dp_deployment
+from ray.serve.schema import ApplicationStatus, ReplicaState
+from ray.util.state import list_actors
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+
+
+@pytest.fixture(autouse=True)
+def cleanup():
+    yield
+    serve.shutdown()
+    ray.shutdown()
+
+
+def is_default_app_running():
+    try:
+        app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        return app.status == ApplicationStatus.RUNNING
+    except (KeyError, AttributeError):
+        return False
+
+
+def get_num_running_replicas(deployment_name: str) -> int:
+    """Get the number of RUNNING replicas for a deployment."""
+    dep = (
+        serve.status().applications[SERVE_DEFAULT_APP_NAME].deployments[deployment_name]
+    )
+    return dep.replica_states.get(ReplicaState.RUNNING, 0)
+
+
+def find_vllm_worker_actors():
+    """Find all alive RayWorkerProc actors created by vLLM's Ray V2 executor backend."""
+    return list_actors(
+        filters=[
+            ("class_name", "=", "RayWorkerProc"),
+            ("state", "=", "ALIVE"),
+        ]
+    )
+
+
+def kill_one_vllm_worker():
+    """Kill one RayWorkerProc actor by sending SIGKILL to its process.
+
+    Returns the node_id of the killed worker so we can verify which gang gets torn down.
+    """
+    actors = find_vllm_worker_actors()
+    assert len(actors) > 0, "No alive RayWorkerProc actors found"
+
+    target = actors[0]
+    target_pid = target.pid
+    target_node_id = target.node_id
+
+    @ray.remote(num_cpus=0)
+    def kill_pid(pid):
+        os.kill(pid, signal.SIGKILL)
+        return f"Killed pid {pid}"
+
+    result = ray.get(
+        kill_pid.options(
+            scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
+                node_id=target_node_id, soft=False
+            ),
+        ).remote(target_pid)
+    )
+    return target_node_id, result
+
+
+def test_llm_serve_dp_engine_fault():
+    """DP deployment recovers when a vLLM GPU worker process is killed."""
+    deployment_name = "DPServer:microsoft--Phi-tiny-MoE-instruct"
+    dp_size = 2
+    num_replicas = 2
+    expected_serve_replicas = num_replicas * dp_size
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="microsoft/Phi-tiny-MoE-instruct",
+            model_source="microsoft/Phi-tiny-MoE-instruct",
+        ),
+        deployment_config=dict(
+            num_replicas=num_replicas,
+            health_check_period_s=1,
+            health_check_timeout_s=3,
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=dp_size,
+            distributed_executor_backend="ray",
+            enable_expert_parallel=True,
+            max_model_len=1024,
+            max_num_seqs=32,
+            enforce_eager=True,
+        ),
+        runtime_env={
+            "env_vars": {
+                "VLLM_DISABLE_COMPILE_CACHE": "1",
+            },
+        },
+    )
+
+    handle = serve.run(build_dp_deployment(llm_config), blocking=False)
+    wait_for_condition(is_default_app_running, timeout=180)
+
+    # Step 1: Verify steady state — all replicas running.
+    assert get_num_running_replicas(deployment_name) == expected_serve_replicas
+
+    # Verify RayWorkerProc actors exist.
+    workers_before = find_vllm_worker_actors()
+    assert len(workers_before) > 0
+
+    # Step 2: Start continuous request sending, then kill a GPU worker.
+    @ray.remote
+    class RequestSender:
+        def __init__(self, h):
+            self._handle = h.options(stream=True)
+            self.total = 0
+            self.errors = []
+            self._stop = False
+
+        async def send(self, prompt="test", max_tokens=1, ignore_errors=False):
+            req = CompletionRequest(
+                model="microsoft/Phi-tiny-MoE-instruct",
+                prompt=prompt,
+                max_tokens=max_tokens,
+            )
+            try:
+                async for _ in self._handle.completions.remote(req):
+                    pass
+                self.total += 1
+            except Exception as e:
+                if not ignore_errors:
+                    self.errors.append(str(e))
+                self.total += 1
+
+        async def run(self):
+            while not self._stop:
+                await self.send(prompt="Hello, world!", max_tokens=5)
+                await asyncio.sleep(0.1)
+
+        def stop(self):
+            self._stop = True
+
+        def get_results(self):
+            return self.total, self.errors
+
+    sender = RequestSender.remote(handle)
+
+    # Send a few warm-up requests to confirm the deployment works.
+    for _ in range(5):
+        sender.send.remote(ignore_errors=True)
+        time.sleep(0.5)
+
+    # Kill one RayWorkerProc GPU worker process.
+    killed_node_id, kill_msg = kill_one_vllm_worker()
+    print(f"Killed vLLM worker on node {killed_node_id}: {kill_msg}")
+
+    # Step 3: Wait for the faulty gang to be torn down.
+    # Detection: compiled DAG timeout (~30s) + health check overhead.
+    wait_for_condition(
+        lambda: get_num_running_replicas(deployment_name) < expected_serve_replicas,
+        timeout=60,
+    )
+
+    # Step 4: Verify the surviving replicas continue serving requests.
+    time.sleep(2)
+    sender.run.remote()
+
+    # Step 5: Wait for full recovery — all replicas back to RUNNING.
+    wait_for_condition(
+        lambda: get_num_running_replicas(deployment_name) == expected_serve_replicas,
+        timeout=120,
+    )
+
+    # Step 6: Assert zero downtime — requests should have kept flowing
+    # through the surviving gang during recovery.
+    ray.get(sender.stop.remote())
+    total, errors = ray.get(sender.get_results.remote())
+    assert total > 0, "Expected at least one successful request"
+    assert len(errors) == 0, f"Expected zero errors, got {errors}"
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", __file__])

release/llm_tests/serve/test_llm_serve_fault_tolerance.py

@@ -0,0 +1,94 @@
+import time
+from typing import Literal, List, Generator
+
+import pytest
+import ray
+from ray import serve
+from ray.serve.llm import LLMConfig, ModelLoadingConfig, build_llm_deployment
+
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+RAY_MODEL_ID = "qwen-0.5b"
+
+
+def get_llm_config(
+    tensor_parallel_size: int = 1,
+) -> LLMConfig:
+    """Create LLMConfig with specified parallelism parameters."""
+    return LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id=RAY_MODEL_ID,
+            model_source=MODEL_ID,
+        ),
+        deployment_config=dict(
+            name="test",
+            num_replicas=2,
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=tensor_parallel_size,
+            enforce_eager=True,
+        ),
+    )
+
+
+def find_replica_ids(deployment_name: str) -> List[str]:
+    actors = ray.util.list_named_actors("serve")
+    found_replica_ids = []
+    for actor in actors:
+        if deployment_name in actor["name"]:
+            found_replica_ids.append(actor["name"])
+    return found_replica_ids
+
+
+def kill_replica(replica_id: str) -> None:
+    actor = ray.get_actor(replica_id, namespace="serve")
+    ray.kill(actor)
+
+
+@pytest.fixture(name="app", scope="function")
+def start_ray_serve(
+    tensor_parallel_size: int = 1,
+) -> Generator:
+    """Start Ray Serve with specified parallelism parameters."""
+    llm_config: LLMConfig = get_llm_config(tensor_parallel_size)
+    app = build_llm_deployment(llm_config, name_prefix="LLM:")
+    serve.run(app, blocking=False)
+    yield app
+    serve.shutdown()
+
+
+def wait_for_deployment_status(
+    deployment_name: str, status: Literal["HEALTHY", "UNHEALTHY"], timeout_s: int = 120
+) -> None:
+    s = time.time()
+    while time.time() - s < timeout_s:
+        print(f"Waiting for deployment {deployment_name} to become {status}")
+        state = serve.status()
+        if state.applications["default"].deployments[deployment_name].status == status:
+            return
+        time.sleep(1)
+    raise TimeoutError(
+        f"Deployment {deployment_name} did not become "
+        f"{status} within {timeout_s} seconds"
+    )
+
+
+def test_recovery_from_replica_failure(app) -> None:
+    """Tests that the deployment recovers from replica failure."""
+    dname = "LLM:test"
+    wait_for_deployment_status(dname, "HEALTHY", timeout_s=60)
+
+    # Kill both replicas
+    replica_ids = find_replica_ids(dname)
+    for replica_id in replica_ids:
+        print(f"Killing replica {replica_id}")
+        kill_replica(replica_id)
+
+    # wait for deployment to get unhealthy
+    wait_for_deployment_status(dname, "UNHEALTHY", timeout_s=60)
+
+    # Wait again for deployment to get healthy
+    wait_for_deployment_status(dname, "HEALTHY", timeout_s=60)
+
+
+if __name__ == "__main__":
+    pytest.main(["-xvs", __file__])

release/llm_tests/serve/test_llm_serve_integration.py

@@ -0,0 +1,447 @@
+import os
+import pytest
+import requests
+import sys
+
+from ray import serve
+from ray.serve.llm import LLMConfig, build_openai_app
+from vllm import AsyncEngineArgs
+
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
+from vllm.sampling_params import SamplingParams
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.schema import ApplicationStatus
+import time
+
+# Pooling models (classify/reward) are only served through vLLM's native ASGI
+# app, which is used when direct streaming is enabled. The default OpenAiIngress
+# path does not expose /classify or /pooling, so these tests only run when
+# RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING=1.
+direct_streaming_only = pytest.mark.skipif(
+    os.environ.get("RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING", "0") != "1",
+    reason="Pooling/classify endpoints are only served in direct-streaming mode "
+    "(RAY_SERVE_LLM_ENABLE_DIRECT_STREAMING=1).",
+)
+
+
+@pytest.mark.asyncio(scope="function")
+async def test_engine_metrics():
+    """
+    Test that the stat logger can be created successfully.
+    Keeping this test small to focus on instantiating the
+    derived class correctly.
+    """
+
+    engine_args = AsyncEngineArgs(
+        model="Qwen/Qwen2.5-0.5B-Instruct",
+        dtype="auto",
+        disable_log_stats=False,
+        enforce_eager=True,
+    )
+
+    engine = AsyncLLM.from_engine_args(
+        engine_args, stat_loggers=[RayPrometheusStatLogger]
+    )
+
+    for i, prompt in enumerate(["What is the capital of France?", "What is 2+2?"]):
+        results = engine.generate(
+            request_id=f"request-id-{i}",
+            prompt=prompt,
+            sampling_params=SamplingParams(max_tokens=10),
+        )
+
+        async for _ in results:
+            pass
+
+
+@pytest.mark.asyncio(scope="function")
+async def test_engine_metrics_with_lora():
+    """
+    Test that the stat logger can be created successfully with LoRA configuration.
+    This test validates LoRA-enabled engine initialization and basic functionality.
+    """
+
+    engine_args = AsyncEngineArgs(
+        model="Qwen/Qwen2.5-0.5B-Instruct",  # Using smaller model for testing
+        disable_log_stats=False,
+        enforce_eager=True,
+        enable_prefix_caching=True,
+        max_model_len=512,
+        max_lora_rank=64,
+        enable_lora=True,
+        max_loras=3,
+        max_cpu_loras=5,
+    )
+
+    engine = AsyncLLM.from_engine_args(
+        engine_args, stat_loggers=[RayPrometheusStatLogger]
+    )
+
+    for i, prompt in enumerate(["What is the capital of France?", "What is 2+2?"]):
+        results = engine.generate(
+            request_id=f"lora-request-id-{i}",
+            prompt=prompt,
+            sampling_params=SamplingParams(max_tokens=10),
+        )
+
+        async for _ in results:
+            pass
+
+
+@pytest.mark.asyncio(scope="function")
+async def test_engine_metrics_with_spec_decode():
+    """
+    Test that the stat logger can be created successfully with speculative decoding configuration.
+    This test validates speculative decoding engine initialization and basic functionality.
+    """
+
+    engine_args = AsyncEngineArgs(
+        model="Qwen/Qwen2.5-0.5B-Instruct",
+        dtype="auto",
+        disable_log_stats=False,
+        enforce_eager=True,
+        trust_remote_code=True,
+        enable_prefix_caching=True,
+        max_model_len=256,
+        speculative_config={
+            "method": "ngram",
+            "num_speculative_tokens": 5,
+            "prompt_lookup_max": 4,
+        },
+    )
+
+    engine = AsyncLLM.from_engine_args(
+        engine_args, stat_loggers=[RayPrometheusStatLogger]
+    )
+
+    for i, prompt in enumerate(["What is the capital of France?", "What is 2+2?"]):
+        results = engine.generate(
+            request_id=f"spec-request-id-{i}",
+            prompt=prompt,
+            sampling_params=SamplingParams(max_tokens=10),
+        )
+
+        async for _ in results:
+            pass
+
+
+def is_default_app_running():
+    """Check if the default application is running successfully."""
+    try:
+        default_app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        return default_app.status == ApplicationStatus.RUNNING
+    except (KeyError, AttributeError):
+        return False
+
+
+@pytest.mark.parametrize("model_name", ["deepseek-ai/DeepSeek-V2-Lite"])
+def test_deepseek_model(model_name):
+    """
+    Test that the deepseek model can be loaded successfully.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=dict(
+            model_id=model_name,
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(min_replicas=1, max_replicas=1),
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+            pipeline_parallel_size=2,
+            gpu_memory_utilization=0.92,
+            dtype="auto",
+            max_num_seqs=40,
+            max_model_len=8192,
+            enable_chunked_prefill=True,
+            enable_prefix_caching=True,
+            enforce_eager=True,
+            trust_remote_code=True,
+        ),
+    )
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=300)
+    serve.shutdown()
+    time.sleep(1)
+
+
+@pytest.mark.parametrize("model_name", ["openai/whisper-small"])
+def test_transcription_model(model_name):
+    """
+    Test that the transcription models can be loaded successfully.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=dict(
+            model_id=model_name,
+            model_source=model_name,
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(min_replicas=1, max_replicas=4),
+        ),
+        engine_kwargs=dict(
+            trust_remote_code=True,
+            gpu_memory_utilization=0.9,
+            enable_prefix_caching=True,
+        ),
+    )
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=180)
+    serve.shutdown()
+    time.sleep(1)
+
+
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+def test_embedding_model(model_name):
+    """
+    Test that embedding models can be loaded and serve embedding requests.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=dict(
+            model_id=model_name,
+        ),
+        deployment_config=dict(
+            num_replicas=1,
+        ),
+        engine_kwargs=dict(
+            enforce_eager=True,
+        ),
+    )
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=180)
+
+    response = requests.post(
+        "http://localhost:8000/v1/embeddings",
+        json={
+            "model": model_name,
+            "input": "Hello, world!",
+        },
+    )
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert "data" in data
+    assert len(data["data"]) > 0
+    embedding = data["data"][0]["embedding"]
+    assert isinstance(embedding, list)
+    assert len(embedding) > 0
+    assert all(isinstance(x, float) for x in embedding)
+
+    serve.shutdown()
+    time.sleep(1)
+
+
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+def test_score_model(model_name):
+    """
+    Test that embedding models can serve score requests.
+    """
+    llm_config = LLMConfig(
+        model_loading_config=dict(
+            model_id=model_name,
+        ),
+        deployment_config=dict(
+            num_replicas=1,
+        ),
+        engine_kwargs=dict(
+            enforce_eager=True,
+        ),
+    )
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=180)
+
+    response = requests.post(
+        "http://localhost:8000/v1/score",
+        json={
+            "model": model_name,
+            "text_1": "What is the capital of France?",
+            "text_2": ["Paris is the capital of France.", "Berlin is in Germany."],
+        },
+    )
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert "data" in data
+    assert len(data["data"]) == 2
+    for item in data["data"]:
+        assert "score" in item
+        assert isinstance(item["score"], float)
+
+    serve.shutdown()
+    time.sleep(1)
+
+
+def _validate_classify(item):
+    assert isinstance(item["probs"], list)
+    assert len(item["probs"]) > 0
+    assert item["num_classes"] == len(item["probs"])
+
+
+def _validate_pooling(item):
+    # Reward models emit a per-token pooling vector; ensure it is non-empty.
+    assert len(item["data"]) > 0
+
+
+@direct_streaming_only
+@pytest.mark.parametrize(
+    "model_name,engine_kwargs,endpoint,validate_item",
+    [
+        pytest.param(
+            "Qwen/Qwen3-Reranker-0.6B",
+            dict(
+                hf_overrides={
+                    "architectures": ["Qwen3ForSequenceClassification"],
+                    "classifier_from_token": ["no", "yes"],
+                    "is_original_qwen3_reranker": True,
+                },
+            ),
+            "/classify",
+            _validate_classify,
+            id="classify",
+        ),
+        pytest.param(
+            "internlm/internlm2-1_8b-reward",
+            dict(trust_remote_code=True),
+            "/pooling",
+            _validate_pooling,
+            id="pooling",
+        ),
+    ],
+)
+def test_pooling_model(model_name, engine_kwargs, endpoint, validate_item):
+    """Pooling models (classify/reward) are served via vLLM's native /classify
+    and /pooling endpoints, which are only mounted in direct-streaming mode."""
+    llm_config = LLMConfig(
+        model_loading_config=dict(model_id=model_name),
+        deployment_config=dict(num_replicas=1),
+        engine_kwargs=dict(enforce_eager=True, max_model_len=512, **engine_kwargs),
+    )
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=300)
+
+    response = requests.post(
+        f"http://localhost:8000{endpoint}",
+        json={"model": model_name, "input": "The chef prepared a delicious meal."},
+    )
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert data["object"] == "list"
+    assert len(data["data"]) == 1
+    validate_item(data["data"][0])
+
+    serve.shutdown()
+    time.sleep(1)
+
+
+@pytest.fixture
+def remote_model_app(request):
+    """
+    Fixture that creates an app with a remote code model for testing.
+
+    The remote_code parameter controls whether trust_remote_code is enabled.
+    This helps avoid regressions for pickling issues for custom huggingface configs,
+    since this custom code needs to be registered and imported across processes and workers.
+    """
+    remote_code = request.param
+
+    base_config = {
+        "model_loading_config": dict(
+            model_id="hmellor/Ilama-3.2-1B",
+        ),
+        "deployment_config": dict(
+            autoscaling_config=dict(min_replicas=1, max_replicas=1),
+        ),
+        "engine_kwargs": dict(
+            trust_remote_code=remote_code,
+        ),
+    }
+
+    llm_config = LLMConfig(**base_config)
+    app = build_openai_app({"llm_configs": [llm_config]})
+
+    yield app
+
+    # Cleanup
+    serve.shutdown()
+    time.sleep(1)
+
+
+class TestRemoteCode:
+    """Tests for remote code model loading behavior."""
+
+    @pytest.mark.parametrize("remote_model_app", [False], indirect=True)
+    def test_remote_code_failure(self, remote_model_app):
+        """
+        Tests that a remote code model fails to load when trust_remote_code=False.
+
+        If it loads successfully without remote code, the fixture should be changed to one that does require remote code.
+        """
+        app = remote_model_app
+        with pytest.raises(RuntimeError, match="Deploying application default failed"):
+            serve.run(app, blocking=False)
+
+        def check_for_failed_deployment():
+            """Check if the application deployment has failed."""
+            try:
+                default_app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+                return default_app.status == ApplicationStatus.DEPLOY_FAILED
+            except (KeyError, AttributeError):
+                return False
+
+        # Wait for either failure or success (timeout after 2 minutes)
+        try:
+            wait_for_condition(check_for_failed_deployment, timeout=120)
+        except TimeoutError:
+            # If deployment didn't fail, check if it succeeded
+            if is_default_app_running():
+                pytest.fail(
+                    "App deployed successfully without trust_remote_code=True. "
+                    "This model may not actually require remote code. "
+                    "Consider using a different model that requires remote code."
+                )
+            else:
+                pytest.fail("Deployment did not fail or succeed within timeout period.")
+
+    @pytest.mark.parametrize("remote_model_app", [True], indirect=True)
+    def test_remote_code_success(self, remote_model_app):
+        """
+        Tests that a remote code model succeeds to load when trust_remote_code=True.
+        """
+        app = remote_model_app
+
+        serve.run(app, blocking=False)
+
+        # Wait for the application to be running (timeout after 5 minutes)
+        wait_for_condition(is_default_app_running, timeout=300)
+
+
+def test_nested_engine_kwargs_structured_outputs():
+    """Regression test for https://github.com/ray-project/ray/pull/60380"""
+    llm_config = LLMConfig(
+        model_loading_config=dict(
+            model_id="Qwen/Qwen2.5-0.5B-Instruct",
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(min_replicas=1, max_replicas=1),
+        ),
+        engine_kwargs=dict(
+            enforce_eager=True,
+            max_model_len=512,
+            structured_outputs_config={
+                "backend": "xgrammar",
+            },
+        ),
+    )
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=180)
+    serve.shutdown()
+    time.sleep(1)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

release/llm_tests/serve/test_llm_serve_multi_node_integration.py

@@ -0,0 +1,482 @@
+import pathlib
+from collections import defaultdict
+import asyncio
+import time
+
+import pytest
+import yaml
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.experimental.internal_kv import _internal_kv_list
+from ray.llm._internal.serve.serving_patterns.data_parallel.dp_server import (
+    GangMasterInfoRegistry,
+)
+from ray.serve._private.common import DeploymentID, ReplicaState
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME, SERVE_NAMESPACE
+from ray.serve.llm import (
+    build_dp_deployment,
+    build_dp_openai_app,
+    build_pd_openai_app,
+    build_openai_app,
+    LLMConfig,
+    LLMServingArgs,
+    ModelLoadingConfig,
+)
+from ray.serve.schema import ApplicationStatus
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+
+CONFIGS_DIR = pathlib.Path(__file__).parent / "configs"
+
+
+@pytest.fixture(autouse=True)
+def cleanup_ray_resources():
+    """Automatically cleanup Ray resources between tests to prevent conflicts."""
+    yield
+    serve.shutdown()
+    ray.shutdown()
+
+
+def is_default_app_running():
+    """Check if the default application is running successfully."""
+    try:
+        default_app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        return default_app.status == ApplicationStatus.RUNNING
+    except (KeyError, AttributeError):
+        return False
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size",
+    [
+        (2, 4),  # TPxPP=8 > 4 GPUs/node, FORCES cross-node placement
+        (4, 2),  # TPxPP=8 > 4 GPUs/node, FORCES cross-node placement
+    ],
+)
+def test_llm_serve_multi_node(tp_size, pp_size):
+    """Test multi-node Ray Serve LLM deployment with custom placement groups.
+
+    Cluster: 2 nodes x 4 GPUs = 8 total GPUs
+    TPxPP=8 exceeds per-node capacity, forcing cross-node deployment.
+    """
+    total_gpus = tp_size * pp_size
+    placement_group_config = {
+        "bundles": [{"GPU": 1, "CPU": 1}] * total_gpus,
+        "strategy": "PACK",
+    }
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="opt-1.3b",
+            model_source="facebook/opt-1.3b",
+        ),
+        deployment_config=dict(
+            autoscaling_config=dict(
+                min_replicas=1,
+                max_replicas=1,
+            ),
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=tp_size,
+            pipeline_parallel_size=pp_size,
+            distributed_executor_backend="ray",
+            max_model_len=512,
+            max_num_batched_tokens=256,
+            enforce_eager=True,
+        ),
+        placement_group_config=placement_group_config,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    app = build_openai_app(llm_serving_args=LLMServingArgs(llm_configs=[llm_config]))
+    serve.run(app, blocking=False)
+
+    wait_for_condition(is_default_app_running, timeout=300)
+
+
+@pytest.mark.parametrize(
+    "tp_size,dp_size,num_replicas,placement_group_config",
+    [
+        # TP=1 cases
+        (1, 4, None, {"bundles": [{"GPU": 1, "CPU": 1}]}),  # Single group, single node
+        (1, 8, None, {"bundles": [{"GPU": 1, "CPU": 1}]}),  # Single group, multi-node
+        (1, 2, 2, {"bundles": [{"GPU": 1, "CPU": 1}]}),  # Multi-group, single node
+        (1, 4, 2, {"bundles": [{"GPU": 1, "CPU": 1}]}),  # Multi-group, multi-node
+        # TP=2 cases — auto-generates correct bundles from TP size
+        (2, 2, 1, None),  # TP, single group, single node
+        (2, 2, 2, None),  # TP, multi-group, multi-node
+        # TP=2 cases — explicit placement_group_config with 2 bundles for TP=2
+        (2, 2, None, {"bundles": [{"GPU": 1, "CPU": 1}, {"GPU": 1}]}),
+        (2, 2, 2, {"bundles": [{"GPU": 1, "CPU": 1}, {"GPU": 1}]}),
+    ],
+)
+def test_llm_serve_data_parallelism(
+    tp_size, dp_size, num_replicas, placement_group_config
+):
+    """Test Data Parallelism deployment."""
+    deployment_config = dict()
+    if num_replicas is not None:
+        deployment_config["num_replicas"] = num_replicas
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="microsoft/Phi-tiny-MoE-instruct",
+            model_source="microsoft/Phi-tiny-MoE-instruct",
+        ),
+        deployment_config=deployment_config,
+        engine_kwargs=dict(
+            tensor_parallel_size=tp_size,
+            pipeline_parallel_size=1,
+            data_parallel_size=dp_size,
+            distributed_executor_backend="ray",
+            max_model_len=1024,
+            max_num_seqs=32,
+            enforce_eager=True,
+        ),
+        placement_group_config=placement_group_config,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    app = build_dp_deployment(llm_config)
+    serve.run(app, blocking=False)
+
+    wait_for_condition(is_default_app_running, timeout=300)
+
+
+def test_llm_serve_data_parallelism_autoscaling():
+    """Test that DP deployment with autoscaling scales up under load."""
+    dp_size = 2
+    deployment_name = "DPServer:microsoft--Phi-tiny-MoE-instruct"
+
+    deployment_config = {
+        "num_replicas": "auto",
+        "autoscaling_config": dict(
+            min_replicas=1,
+            max_replicas=2,
+            upscale_delay_s=0.1,
+            downscale_delay_s=5,
+            metrics_interval_s=1,
+            look_back_period_s=2,
+            target_ongoing_requests=1,
+        ),
+        "max_ongoing_requests": 20,
+    }
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="microsoft/Phi-tiny-MoE-instruct",
+            model_source="microsoft/Phi-tiny-MoE-instruct",
+        ),
+        deployment_config=deployment_config,
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+            pipeline_parallel_size=1,
+            data_parallel_size=dp_size,
+            distributed_executor_backend="ray",
+            max_model_len=1024,
+            max_num_seqs=32,
+            enforce_eager=True,
+        ),
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    app = build_dp_deployment(llm_config)
+    handle = serve.run(app, blocking=False)
+
+    wait_for_condition(is_default_app_running, timeout=300)
+
+    def check_num_replicas_eq(target):
+        dep = (
+            serve.status()
+            .applications[SERVE_DEFAULT_APP_NAME]
+            .deployments[deployment_name]
+        )
+        running = dep.replica_states.get("RUNNING", 0)
+        assert running == target
+        return True
+
+    def get_total_replicas():
+        dep = (
+            serve.status()
+            .applications[SERVE_DEFAULT_APP_NAME]
+            .deployments[deployment_name]
+        )
+        return sum(dep.replica_states.values())
+
+    def check_total_replicas_gte(target):
+        assert get_total_replicas() >= target
+        return True
+
+    wait_for_condition(check_num_replicas_eq, target=2, timeout=60)
+
+    request = CompletionRequest(
+        model="microsoft/Phi-tiny-MoE-instruct",
+        prompt="Write a very long detailed story about",
+        # max_tokens must be less than max_model_len; otherwise vLLM rejects
+        # the request with a BadRequestError
+        max_tokens=512,
+    )
+    # Send enough concurrent requests to trigger upscaling
+    streaming_handle = handle.options(stream=True)
+    results = [streaming_handle.completions.remote(request) for _ in range(20)]
+
+    wait_for_condition(check_total_replicas_gte, target=4, timeout=120)
+
+    total = get_total_replicas()
+    assert total % dp_size == 0
+
+    # Drain requests by consuming streaming responses and check for errors
+    errors = []
+    for res in results:
+        for chunk in res:
+            if hasattr(chunk, "error"):
+                errors.append(chunk.error)
+    assert len(errors) == 0
+
+    # After requests drain, verify scale-down back to min_replicas
+    wait_for_condition(check_num_replicas_eq, target=2, timeout=120)
+
+    dep = (
+        serve.status().applications[SERVE_DEFAULT_APP_NAME].deployments[deployment_name]
+    )
+    running = dep.replica_states.get("RUNNING", 0)
+    assert running % dp_size == 0
+
+
+def test_llm_serve_data_parallelism_cleanup():
+    """Test that Data Parallelism KV entries are cleaned up on shutdown."""
+    placement_group_config = {
+        "bundles": [{"GPU": 1, "CPU": 1}],
+    }
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="microsoft/Phi-tiny-MoE-instruct",
+            model_source="microsoft/Phi-tiny-MoE-instruct",
+        ),
+        deployment_config=dict(),
+        engine_kwargs=dict(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=4,
+            distributed_executor_backend="ray",
+            max_model_len=1024,
+            max_num_seqs=32,
+            enforce_eager=True,
+        ),
+        placement_group_config=placement_group_config,
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    app = build_dp_deployment(llm_config)
+    serve.run(app, blocking=False)
+    wait_for_condition(is_default_app_running, timeout=300)
+
+    master_keys = _internal_kv_list(GangMasterInfoRegistry._KEY_PREFIX)
+    assert len(master_keys) > 0
+
+    serve.shutdown()
+
+
+def test_llm_serve_data_parallelism_declarative():
+    """Test Data Parallelism deployment via declarative config."""
+    config_path = CONFIGS_DIR / "serve_phi_tiny_moe_dp4_gang.yaml"
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+
+    app_config = config["applications"][0]
+    app = build_dp_openai_app(app_config["args"])
+    serve.run(app, blocking=False)
+
+    wait_for_condition(is_default_app_running, timeout=300)
+
+
+def test_llm_serve_prefill_decode_with_data_parallelism():
+    """Test Prefill-Decode disaggregation with Data Parallelism and Expert Parallelism.
+
+    Cluster: 2 nodes x 4 GPUs = 8 GPUs total
+    - Prefill: DP=4 (scheduled on node with "prefill" custom resource)
+    - Decode: DP=4 (scheduled on node with "decode" custom resource)
+
+    Note: This test requires RAY_SERVE_USE_COMPACT_SCHEDULING_STRATEGY=1 to be set
+    (configured in release_tests.yaml). Without this flag, Serve uses the default
+    SPREAD scheduling strategy, which will prevent DP replicas from being colocated.
+    """
+    model_loading_config = ModelLoadingConfig(
+        model_id="deepseek",
+        model_source="deepseek-ai/DeepSeek-V2-Lite",
+    )
+    base_engine_kwargs = {
+        "tensor_parallel_size": 1,
+        "enable_expert_parallel": True,
+        "load_format": "dummy",
+        "max_model_len": 512,
+        "max_num_batched_tokens": 256,
+        "enforce_eager": True,
+    }
+
+    prefill_config = LLMConfig(
+        model_loading_config=model_loading_config,
+        engine_kwargs={
+            **base_engine_kwargs,
+            "data_parallel_size": 4,
+            "kv_transfer_config": {
+                "kv_connector": "NixlConnector",
+                "kv_role": "kv_both",
+            },
+        },
+        experimental_configs={
+            # Use ports below the Linux ephemeral range (32768-60999) to
+            # prevent conflicts with the vLLM DP coordinator's random TCP
+            # port allocations.
+            "NIXL_SIDE_CHANNEL_PORT_BASE": 15000,  # Prefill port range
+        },
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    decode_config = LLMConfig(
+        model_loading_config=model_loading_config,
+        engine_kwargs={
+            **base_engine_kwargs,
+            "data_parallel_size": 4,
+            "kv_transfer_config": {
+                "kv_connector": "NixlConnector",
+                "kv_role": "kv_both",
+            },
+        },
+        experimental_configs={
+            "NIXL_SIDE_CHANNEL_PORT_BASE": 16000,  # Decode port range (different)
+        },
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+
+    # build_pd_openai_app auto-detects DP and uses build_dp_deployment
+    app = build_pd_openai_app(
+        {
+            "prefill_config": prefill_config,
+            "decode_config": decode_config,
+        }
+    )
+    serve.run(app, blocking=False)
+
+    wait_for_condition(is_default_app_running, timeout=300)
+
+
+def test_llm_serve_data_parallelism_worker_fault():
+    """Test DP deployment recovers after replica death."""
+
+    deployment_name = "DPServer:microsoft--Phi-tiny-MoE-instruct"
+    dp_size = 2
+    num_replicas = 2
+    expected_serve_replicas = num_replicas * dp_size
+
+    llm_config = LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id="microsoft/Phi-tiny-MoE-instruct",
+            model_source="microsoft/Phi-tiny-MoE-instruct",
+        ),
+        deployment_config=dict(
+            num_replicas=2,
+            health_check_period_s=1,
+            health_check_timeout_s=5,
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            data_parallel_size=dp_size,
+            distributed_executor_backend="ray",
+            max_model_len=1024,
+            max_num_seqs=32,
+            enforce_eager=True,
+        ),
+        runtime_env={"env_vars": {"VLLM_DISABLE_COMPILE_CACHE": "1"}},
+    )
+    handle = serve.run(build_dp_deployment(llm_config), blocking=False)
+    wait_for_condition(is_default_app_running, timeout=300)
+
+    deployment_id = DeploymentID(name=deployment_name, app_name=SERVE_DEFAULT_APP_NAME)
+    controller = serve.context._global_client._controller
+
+    # Verify 4 running replicas in 2 gangs.
+    replicas = ray.get(
+        controller._dump_replica_states_for_testing.remote(deployment_id)
+    )
+    running = replicas.get([ReplicaState.RUNNING])
+    assert len(running) == expected_serve_replicas
+
+    gangs = defaultdict(list)
+    for r in running:
+        assert r.gang_context is not None
+        gangs[r.gang_context.gang_id].append(r)
+    assert len(gangs) == 2
+
+    # Pick one gang and kill one of its replicas.
+    target_gang_id = list(gangs.keys())[0]
+    victim = gangs[target_gang_id][0]
+    victim_actor = ray.get_actor(
+        victim.replica_id.to_full_id_str(), namespace=SERVE_NAMESPACE
+    )
+    original_replica_ids = {r.replica_id.unique_id for r in running}
+
+    # Background request sender
+    @ray.remote
+    class RequestSender:
+        def __init__(self, handle):
+            self._handle = handle.options(stream=True)
+            self.total = 0
+            self.errors = []
+            self._stop = False
+
+        async def run(self):
+            request = CompletionRequest(
+                model="microsoft/Phi-tiny-MoE-instruct",
+                prompt="Hello, world!",
+                max_tokens=5,
+            )
+            while not self._stop:
+                try:
+                    async for _ in self._handle.completions.remote(request):
+                        pass
+                    self.total += 1
+                except Exception as e:
+                    self.errors.append(str(e))
+                    self.total += 1
+                await asyncio.sleep(0.1)
+
+        def stop(self):
+            self._stop = True
+
+        def get_results(self):
+            return self.total, self.errors
+
+    sender = RequestSender.remote(handle)
+    sender.run.remote()
+    time.sleep(1)
+
+    # Kill the victim replica
+    ray.kill(victim_actor, no_restart=False)
+
+    # Wait for the killed gang to be replaced
+    def all_replicas_recovered():
+        replicas = ray.get(
+            controller._dump_replica_states_for_testing.remote(deployment_id)
+        )
+        running = replicas.get([ReplicaState.RUNNING])
+        if len(running) != expected_serve_replicas:
+            return False
+        current_ids = {r.replica_id.unique_id for r in running}
+        new_ids = current_ids - original_replica_ids
+        return len(new_ids) == dp_size
+
+    wait_for_condition(all_replicas_recovered, timeout=180)
+
+    # Stop sender and verify no requests were dropped
+    ray.get(sender.stop.remote())
+    total, errors = ray.get(sender.get_results.remote())
+    assert total > 0
+    assert len(errors) == 0
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", __file__])

release/llm_tests/serve/test_llm_serve_pause_resume.py

@@ -0,0 +1,197 @@
+"""Test pause/resume control plane API for Ray Serve LLM.
+
+This test verifies that the DevIngress pause/resume endpoints work correctly:
+1. Engine starts in unpaused state (is_paused=False)
+2. Pause command halts generation (is_paused=True) while keeping weights in GPU
+3. Resume command restores generation (is_paused=False)
+4. Model can still serve requests after resume
+
+Unlike sleep/wakeup which offloads weights to CPU, pause/resume keeps model
+weights in GPU memory. This is useful for quick pause/resume cycles during
+RL training where you want to pause generation for weight updates without
+the overhead of offloading/reloading weights.
+
+NOTE (Kourosh): This is part of a design in progress for integrating Ray Serve
+LLM with RL workloads. The API is not public and won't be documented until the
+end-to-end story is finalized. Class names and endpoint names may change.
+"""
+
+import time
+
+import pytest
+import requests
+from openai import OpenAI
+from ray import serve
+from ray.llm._internal.serve.core.ingress.dev_ingress import build_dev_openai_app
+from ray.serve.llm import LLMConfig, ModelLoadingConfig
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.schema import ApplicationStatus
+
+MODEL_ID = "Qwen/Qwen2-0.5B-Instruct"
+BASE_URL = "http://localhost:8000"
+
+
+def get_llm_config() -> LLMConfig:
+    """Create LLMConfig for pause/resume testing."""
+    return LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id=MODEL_ID,
+        ),
+        deployment_config=dict(
+            num_replicas=2,
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+        ),
+    )
+
+
+def is_default_app_running():
+    """Check if the default application is running successfully."""
+    try:
+        default_app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        return default_app.status == ApplicationStatus.RUNNING
+    except (KeyError, AttributeError):
+        return False
+
+
+def wait_for_server_ready(timeout: int = 240) -> None:
+    """Wait for the server to be ready to handle requests."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            test_data = {
+                "model": MODEL_ID,
+                "messages": [{"role": "user", "content": "test"}],
+                "max_tokens": 5,
+                "temperature": 0,
+            }
+            response = requests.post(
+                f"{BASE_URL}/v1/chat/completions", json=test_data, timeout=10
+            )
+            if response.status_code == 200:
+                print(f"Server at {BASE_URL} is ready to handle requests!")
+                return
+        except Exception as e:
+            print(f"Waiting for server to be ready... (error: {e})")
+
+        time.sleep(2)
+
+    raise TimeoutError(
+        f"Server at {BASE_URL} did not become ready within {timeout} seconds"
+    )
+
+
+def test_pause_resume_lifecycle():
+    """Test the complete pause/resume lifecycle."""
+    # Start Ray Serve with DevIngress
+    llm_config = get_llm_config()
+    app = build_dev_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+
+    # Wait for application to be running
+    wait_for_condition(is_default_app_running, timeout=300)
+    wait_for_server_ready(timeout=240)
+
+    try:
+        # Step 1: Verify initial state - engine should not be paused
+        print("\n=== Step 1: Checking initial state ===")
+        response = requests.get(
+            f"{BASE_URL}/is_paused?model={MODEL_ID}",
+            timeout=10,
+        )
+        assert response.status_code == 200, f"is_paused returned {response.status_code}"
+        initial_pause_state = response.json().get("is_paused", None)
+        assert (
+            initial_pause_state is False
+        ), f"Expected is_paused=False, got {initial_pause_state}"
+        print(f"✓ Initial paused state: {initial_pause_state}")
+
+        # Step 2: Verify model can serve requests before pause
+        print("\n=== Step 2: Verifying model serves requests before pause ===")
+        client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="fake-key")
+        chat_response = client.chat.completions.create(
+            model=MODEL_ID,
+            messages=[{"role": "user", "content": "Hello"}],
+            max_tokens=10,
+            temperature=0,
+        )
+        assert chat_response.choices[0].message.content is not None
+        print(
+            f"✓ Pre-pause response: {chat_response.choices[0].message.content[:50]}..."
+        )
+
+        # Step 3: Pause the engine
+        print("\n=== Step 3: Pausing engine ===")
+        pause_response = requests.post(
+            f"{BASE_URL}/pause",
+            json={
+                "model": MODEL_ID,
+                "options": {"mode": "abort", "clear_cache": True},
+            },
+            timeout=60,
+        )
+        assert (
+            pause_response.status_code == 200
+        ), f"pause returned {pause_response.status_code}"
+        print("✓ Pause command executed successfully")
+
+        # Step 4: Verify engine is paused
+        print("\n=== Step 4: Verifying engine is paused ===")
+        # Wait for pause to complete
+        wait_for_condition(
+            lambda: requests.get(f"{BASE_URL}/is_paused?model={MODEL_ID}", timeout=5)
+            .json()
+            .get("is_paused")
+            is True,
+            timeout=30,
+            retry_interval_ms=1000,
+        )
+
+        # Step 5: Resume the engine
+        print("\n=== Step 5: Resuming engine ===")
+        resume_response = requests.post(
+            f"{BASE_URL}/resume",
+            json={"model": MODEL_ID, "options": {}},
+            timeout=60,
+        )
+        assert (
+            resume_response.status_code == 200
+        ), f"resume returned {resume_response.status_code}"
+        print("✓ Resume command executed successfully")
+
+        # Step 6: Verify engine is no longer paused
+        print("\n=== Step 6: Verifying engine is resumed ===")
+        wait_for_condition(
+            lambda: requests.get(f"{BASE_URL}/is_paused?model={MODEL_ID}", timeout=5)
+            .json()
+            .get("is_paused")
+            is False,
+            timeout=30,
+            retry_interval_ms=1000,
+        )
+
+        # Step 7: Verify model can still serve requests after resume
+        print("\n=== Step 7: Verifying model can serve requests after resume ===")
+        chat_response = client.chat.completions.create(
+            model=MODEL_ID,
+            messages=[{"role": "user", "content": "What is the capital of France?"}],
+            max_tokens=10,
+            temperature=0,
+        )
+        assert chat_response.choices[0].message.content is not None
+        print(
+            f"✓ Post-resume response: {chat_response.choices[0].message.content[:50]}..."
+        )
+
+        print("\n=== All tests passed! ===")
+
+    finally:
+        # Cleanup
+        serve.shutdown()
+        time.sleep(1)
+
+
+if __name__ == "__main__":
+    pytest.main(["-xvs", __file__])

release/llm_tests/serve/test_llm_serve_rlhf.py

@@ -0,0 +1,202 @@
+"""Test collective_rpc control plane API for Ray Serve LLM.
+
+This test verifies that the DevIngress /collective_rpc endpoint works correctly
+for RLHF-style weight synchronization workflows:
+
+1. Server starts with worker_extension_cls for weight update methods
+2. Trainer initializes NCCL process group with all inference workers
+3. Trainer broadcasts weight updates to all workers via collective_rpc
+4. Workers receive and apply the weight updates
+5. Inference continues to work with updated weights
+
+This demonstrates the core RLHF workflow where:
+- Trainer and inference engine form a single NCCL communicator
+- Weights are synchronized via high-bandwidth GPU-to-GPU transfer
+- The /collective_rpc endpoint orchestrates the RPC across all replicas/workers
+
+NOTE (Kourosh): This is part of a design in progress for integrating Ray Serve
+LLM with RL workloads. The API is not public and won't be documented until the
+end-to-end story is finalized. Class names and endpoint names may change.
+"""
+
+import time
+
+import pytest
+import ray
+import requests
+from openai import OpenAI
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.llm._internal.serve.core.ingress.dev_ingress import build_dev_openai_app
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.llm import LLMConfig, ModelLoadingConfig
+from ray.serve.schema import ApplicationStatus
+from rlhf_utils import TrainerActor
+
+MODEL_ID = "facebook/opt-125m"
+BASE_URL = "http://localhost:8000"
+TENSOR_PARALLEL_SIZE = 2
+NUM_REPLICAS = 1
+
+
+def get_llm_config() -> LLMConfig:
+    """Create LLMConfig for collective_rpc testing."""
+    return LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id=MODEL_ID,
+        ),
+        deployment_config=dict(
+            num_replicas=NUM_REPLICAS,
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=TENSOR_PARALLEL_SIZE,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+            # Worker extension for RLHF weight updates
+            worker_extension_cls="rlhf_utils.WorkerExtension",
+        ),
+    )
+
+
+def is_default_app_running():
+    """Check if the default application is running successfully."""
+    try:
+        default_app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        return default_app.status == ApplicationStatus.RUNNING
+    except (KeyError, AttributeError):
+        return False
+
+
+def wait_for_server_ready(timeout: int = 240) -> None:
+    """Wait for the server to be ready to handle requests."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            test_data = {
+                "model": MODEL_ID,
+                "prompt": "Hello",
+                "max_tokens": 5,
+            }
+            response = requests.post(
+                f"{BASE_URL}/v1/completions", json=test_data, timeout=10
+            )
+            if response.status_code == 200:
+                print(f"Server at {BASE_URL} is ready!")
+                return
+        except Exception as e:
+            print(f"Waiting for server... ({e})")
+
+        time.sleep(2)
+
+    raise TimeoutError(f"Server not ready within {timeout} seconds")
+
+
+def call_collective_rpc_sync(method: str, args: list = None) -> dict:
+    """Synchronously call the /collective_rpc endpoint."""
+    response = requests.post(
+        f"{BASE_URL}/collective_rpc",
+        json={
+            "model": MODEL_ID,
+            "method": method,
+            "args": args or [],
+        },
+        timeout=60,
+    )
+    return response.json()
+
+
+# =============================================================================
+# Test
+# =============================================================================
+
+
+def test_collective_rpc_weight_sync():
+    """Test the complete RLHF weight synchronization workflow."""
+
+    # Start Ray Serve with DevIngress
+    llm_config = get_llm_config()
+    app = build_dev_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+
+    # Wait for application to be running
+    wait_for_condition(is_default_app_running, timeout=300)
+    wait_for_server_ready(timeout=240)
+
+    trainer = None  # Initialize before try block to avoid NameError in finally
+    try:
+        # Step 1: Verify model serves requests before weight update
+        print("\n=== Step 1: Verifying model serves requests before update ===")
+        client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="fake-key")
+        response = client.completions.create(
+            model=MODEL_ID,
+            prompt="Hello, my name is",
+            max_tokens=10,
+            temperature=0,
+        )
+        assert response.choices[0].text is not None
+        original_output = response.choices[0].text
+        print(f"✓ Original output: {original_output!r}")
+
+        # Step 2: Create trainer and set up weight sync group
+        print("\n=== Step 2: Setting up trainer and NCCL process group ===")
+        trainer = TrainerActor.remote(MODEL_ID, BASE_URL)
+        ray.get(
+            trainer.setup_weight_sync_group.remote(
+                tp_size=TENSOR_PARALLEL_SIZE,
+                num_replicas=NUM_REPLICAS,
+            )
+        )
+        print("✓ Weight sync group established")
+
+        # Step 3: Broadcast weight updates (zero out weights)
+        print("\n=== Step 3: Broadcasting weight updates ===")
+        start_time = time.time()
+        ray.get(trainer.update_weights.remote())
+        elapsed = time.time() - start_time
+        print(f"✓ Weight update completed in {elapsed:.2f}s")
+
+        # Step 4: Verify weights changed on inference workers
+        print("\n=== Step 4: Verifying weights changed on workers ===")
+        result = call_collective_rpc_sync("check_weights_changed")
+        print(f"check_weights_changed response: {result}")
+
+        # Verify all workers report weights changed
+        assert "results" in result, f"Expected 'results' in response: {result}"
+        for replica_result in result["results"]:
+            worker_results = replica_result.get("worker_results", [])
+            for worker_result in worker_results:
+                assert (
+                    worker_result
+                ), f"Worker reported weights not changed: {replica_result}"
+        print("✓ All workers confirmed weights updated")
+
+        # Step 5: Verify model still serves requests (with zeroed weights)
+        print("\n=== Step 5: Verifying inference works with updated weights ===")
+        response = client.completions.create(
+            model=MODEL_ID,
+            prompt="Hello, my name is",
+            max_tokens=10,
+            temperature=0,
+        )
+        assert response.choices[0].text is not None
+        updated_output = response.choices[0].text
+        print(f"✓ Output with zeroed weights: {updated_output!r}")
+
+        # Output should be different since weights are now zero
+        # (model produces garbage/different output)
+        print(f"\nOriginal: {original_output!r}")
+        print(f"Updated:  {updated_output!r}")
+
+        print("\n=== All tests passed! ===")
+
+    finally:
+        # Cleanup
+        if trainer is not None:
+            ray.kill(trainer)
+        serve.shutdown()
+        ray.shutdown()
+        time.sleep(1)
+
+
+if __name__ == "__main__":
+    pytest.main(["-xvs", __file__])

release/llm_tests/serve/test_llm_serve_sglang.py

@@ -0,0 +1,598 @@
+import sys
+
+import httpx
+import pytest
+from openai import OpenAI
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.llm._internal.serve.engines.sglang import SGLangServer
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.llm import LLMConfig, build_openai_app
+from ray.serve.schema import ApplicationStatus
+
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+RAY_MODEL_ID = "qwen-0.5b-sglang"
+
+
+def _app_is_running():
+    try:
+        return (
+            serve.status().applications[SERVE_DEFAULT_APP_NAME].status
+            == ApplicationStatus.RUNNING
+        )
+    except (KeyError, AttributeError):
+        return False
+
+
+@pytest.fixture(scope="module")
+def sglang_client():
+    """Start an SGLang server once for all tests in this module."""
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": RAY_MODEL_ID,
+            "model_source": MODEL_ID,
+        },
+        deployment_config={
+            "autoscaling_config": {
+                "min_replicas": 1,
+                "max_replicas": 1,
+            }
+        },
+        server_cls=SGLangServer,
+        engine_kwargs={
+            "model_path": MODEL_ID,
+            "tp_size": 1,
+            "mem_fraction_static": 0.8,
+        },
+    )
+
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(_app_is_running, timeout=300)
+
+    client = OpenAI(base_url="http://localhost:8000/v1", api_key="fake-key")
+    yield client
+
+    serve.shutdown()
+
+
+def test_sglang_serve_e2e(sglang_client):
+    """Verify chat and completions endpoints work end-to-end."""
+    chat_resp = sglang_client.chat.completions.create(
+        model=RAY_MODEL_ID,
+        messages=[{"role": "user", "content": "What is the capital of France?"}],
+        max_tokens=64,
+        temperature=0.0,
+    )
+    assert chat_resp.choices[0].message.content.strip()
+
+    comp_resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID,
+        prompt="The capital of France is",
+        max_tokens=64,
+        temperature=0.0,
+    )
+    assert comp_resp.choices[0].text.strip()
+
+
+def test_sglang_streaming_chat(sglang_client):
+    """Verify streaming chat completions produce incremental chunks."""
+    stream = sglang_client.chat.completions.create(
+        model=RAY_MODEL_ID,
+        messages=[{"role": "user", "content": "Count to 5"}],
+        max_tokens=64,
+        temperature=0.0,
+        stream=True,
+    )
+
+    chunks = list(stream)
+    assert len(chunks) > 1, "Expected multiple streaming chunks"
+
+    # First chunk must include the assistant role.
+    first_delta = chunks[0].choices[0].delta
+    assert first_delta.role == "assistant"
+
+    # Collect all content fragments.
+    collected_text = ""
+    finish_reason = None
+    for chunk in chunks:
+        delta = chunk.choices[0].delta
+        if delta.content is not None:
+            collected_text += delta.content
+        if chunk.choices[0].finish_reason is not None:
+            finish_reason = chunk.choices[0].finish_reason
+
+    assert collected_text.strip(), "Streaming produced no text"
+    assert finish_reason is not None, "Final chunk must have a finish_reason"
+
+
+def test_sglang_streaming_completions(sglang_client):
+    """Verify streaming completions produce incremental chunks."""
+    stream = sglang_client.completions.create(
+        model=RAY_MODEL_ID,
+        prompt="The capital of France is",
+        max_tokens=32,
+        temperature=0.0,
+        stream=True,
+    )
+
+    chunks = list(stream)
+    assert len(chunks) > 1, "Expected multiple streaming chunks"
+
+    collected_text = ""
+    finish_reason = None
+    for chunk in chunks:
+        if chunk.choices[0].text is not None:
+            collected_text += chunk.choices[0].text
+        if chunk.choices[0].finish_reason is not None:
+            finish_reason = chunk.choices[0].finish_reason
+
+    assert collected_text.strip(), "Streaming produced no text"
+    assert finish_reason is not None, "Final chunk must have a finish_reason"
+
+
+def test_sglang_tokenize(sglang_client):
+    """Verify tokenize endpoint works."""
+    resp = httpx.post(
+        "http://localhost:8000/tokenize",
+        json={"model": RAY_MODEL_ID, "prompt": "Hello world"},
+    )
+    assert resp.status_code == 200
+    data = resp.json()
+    assert "tokens" in data
+    assert "count" in data
+    assert "max_model_len" in data
+    assert isinstance(data["tokens"], list)
+    assert len(data["tokens"]) > 0
+    assert data["count"] == len(data["tokens"])
+    assert data["max_model_len"] > 0
+
+
+def test_sglang_detokenize(sglang_client):
+    """Verify detokenize endpoint works and round-trips with tokenize."""
+    # First tokenize
+    tok_resp = httpx.post(
+        "http://localhost:8000/tokenize",
+        json={"model": RAY_MODEL_ID, "prompt": "Hello world"},
+    )
+    assert tok_resp.status_code == 200
+    tokens = tok_resp.json()["tokens"]
+
+    # Then detokenize
+    detok_resp = httpx.post(
+        "http://localhost:8000/detokenize",
+        json={"model": RAY_MODEL_ID, "tokens": tokens},
+    )
+    assert detok_resp.status_code == 200
+    data = detok_resp.json()
+    assert "text" in data
+    assert "Hello world" in data["text"]
+
+
+def test_sglang_batched_completions(sglang_client):
+    """Verify that batched completions (multiple prompts) return one choice per prompt."""
+    prompts = [
+        "The capital of France is",
+        "The capital of Germany is",
+        "The capital of Japan is",
+    ]
+    batch_resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID,
+        prompt=prompts,
+        max_tokens=16,
+        temperature=0.0,
+    )
+
+    assert len(batch_resp.choices) == len(prompts)
+
+    for i, choice in enumerate(batch_resp.choices):
+        assert choice.index == i
+        assert choice.text.strip()
+
+    assert batch_resp.usage.total_tokens > 0
+
+
+@pytest.fixture(scope="module")
+def sglang_embedding_client():
+    """Start an SGLang server with is_embedding enabled for embedding tests."""
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": RAY_MODEL_ID,
+            "model_source": MODEL_ID,
+        },
+        deployment_config={
+            "autoscaling_config": {
+                "min_replicas": 1,
+                "max_replicas": 1,
+            }
+        },
+        server_cls=SGLangServer,
+        engine_kwargs={
+            "model_path": MODEL_ID,
+            "tp_size": 1,
+            "mem_fraction_static": 0.8,
+            "is_embedding": True,
+        },
+    )
+
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+    wait_for_condition(_app_is_running, timeout=300)
+
+    client = OpenAI(base_url="http://localhost:8000/v1", api_key="fake-key")
+    yield client
+
+    serve.shutdown()
+
+
+def test_sglang_embeddings(sglang_embedding_client):
+    """Verify embeddings endpoint works with single and batch inputs."""
+    # Single input
+    emb_resp = sglang_embedding_client.embeddings.create(
+        model=RAY_MODEL_ID,
+        input="Hello world",
+    )
+    assert emb_resp.data
+    assert len(emb_resp.data) == 1
+    assert emb_resp.data[0].embedding
+    assert len(emb_resp.data[0].embedding) > 0
+    assert emb_resp.usage.prompt_tokens > 0
+
+    # Batch input
+    emb_batch_resp = sglang_embedding_client.embeddings.create(
+        model=RAY_MODEL_ID,
+        input=["Hello world", "How are you"],
+    )
+    assert len(emb_batch_resp.data) == 2
+    assert emb_batch_resp.data[0].embedding
+    assert emb_batch_resp.data[1].embedding
+
+
+def test_sglang_serve_e2e_multi_gpu():
+    """Verify SGLang multi-GPU deployment works with tp_size=2.
+
+    Requires a node with at least 2 GPUs. Confirms that:
+    - Placement group bundles are correctly constructed as [{"GPU": 1, "CPU": 1}, {"GPU": 1}]
+    - The model loads and serves inference correctly across both GPUs.
+    """
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": RAY_MODEL_ID,
+            "model_source": MODEL_ID,
+        },
+        deployment_config={
+            "autoscaling_config": {
+                "min_replicas": 1,
+                "max_replicas": 1,
+            }
+        },
+        server_cls=SGLangServer,
+        engine_kwargs={
+            "model_path": MODEL_ID,
+            "tp_size": 2,
+            "mem_fraction_static": 0.8,
+        },
+    )
+
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+
+    try:
+        wait_for_condition(_app_is_running, timeout=300)
+
+        deployment_options = SGLangServer.get_deployment_options(llm_config)
+        expected_bundles = [{"GPU": 1, "CPU": 1}, {"GPU": 1}]
+        assert deployment_options["placement_group_bundles"] == expected_bundles, (
+            f"Expected placement group bundles {expected_bundles}, "
+            f"got {deployment_options['placement_group_bundles']}"
+        )
+
+        client = OpenAI(base_url="http://localhost:8000/v1", api_key="fake-key")
+
+        chat_resp = client.chat.completions.create(
+            model=RAY_MODEL_ID,
+            messages=[{"role": "user", "content": "What is the capital of France?"}],
+            max_tokens=64,
+            temperature=0.0,
+        )
+        assert chat_resp.choices[0].message.content.strip()
+
+        comp_resp = client.completions.create(
+            model=RAY_MODEL_ID,
+            prompt="The capital of France is",
+            max_tokens=64,
+            temperature=0.0,
+        )
+        assert comp_resp.choices[0].text.strip()
+    finally:
+        serve.shutdown()
+
+
+def test_sglang_serve_e2e_pipeline_parallel():
+    """Verify SGLang multi-GPU deployment works with tp_size=2, pp_size=2.
+
+    Requires a node with at least 4 GPUs. Confirms that:
+    - Placement group bundles are correctly constructed as
+      [{"GPU": 1, "CPU": 1}, {"GPU": 1}, {"GPU": 1}, {"GPU": 1}]
+    - The model loads and serves inference correctly across all 4 GPUs.
+    """
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": RAY_MODEL_ID,
+            "model_source": MODEL_ID,
+        },
+        deployment_config={
+            "autoscaling_config": {
+                "min_replicas": 1,
+                "max_replicas": 1,
+            }
+        },
+        server_cls=SGLangServer,
+        engine_kwargs={
+            "model_path": MODEL_ID,
+            "tp_size": 2,
+            "pp_size": 2,
+            "mem_fraction_static": 0.8,
+        },
+    )
+
+    app = build_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+
+    try:
+        wait_for_condition(_app_is_running, timeout=300)
+
+        # tp_size=2, pp_size=2 → num_devices=4 → 4 GPU bundles
+        # first bundle merges replica actor CPU with first GPU worker
+        deployment_options = SGLangServer.get_deployment_options(llm_config)
+        expected_bundles = [{"GPU": 1, "CPU": 1}, {"GPU": 1}, {"GPU": 1}, {"GPU": 1}]
+        assert deployment_options["placement_group_bundles"] == expected_bundles, (
+            f"Expected placement group bundles {expected_bundles}, "
+            f"got {deployment_options['placement_group_bundles']}"
+        )
+
+        client = OpenAI(base_url="http://localhost:8000/v1", api_key="fake-key")
+
+        chat_resp = client.chat.completions.create(
+            model=RAY_MODEL_ID,
+            messages=[{"role": "user", "content": "What is the capital of France?"}],
+            max_tokens=64,
+            temperature=0.0,
+        )
+        assert chat_resp.choices[0].message.content.strip()
+
+        comp_resp = client.completions.create(
+            model=RAY_MODEL_ID,
+            prompt="The capital of France is",
+            max_tokens=64,
+            temperature=0.0,
+        )
+        assert comp_resp.choices[0].text.strip()
+    finally:
+        serve.shutdown()
+
+
+def test_sglang_custom_placement_group_config():
+    """Verify explicit placement_group_config is respected by get_deployment_options.
+
+    Covers the configuration pattern used in serve_sglang_multinode_example.py
+    where users provide custom bundles and strategy for multi-node TP/PP.
+    Does not require GPUs — only tests configuration logic.
+    """
+    custom_bundles = [{"GPU": 1}] * 8
+    custom_strategy = "PACK"
+
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": RAY_MODEL_ID,
+            "model_source": MODEL_ID,
+        },
+        deployment_config={
+            "autoscaling_config": {
+                "min_replicas": 1,
+                "max_replicas": 2,
+                "target_ongoing_requests": 4,
+            }
+        },
+        placement_group_config={
+            "placement_group_bundles": custom_bundles,
+            "placement_group_strategy": custom_strategy,
+        },
+        server_cls=SGLangServer,
+        engine_kwargs={
+            "model_path": MODEL_ID,
+            "tp_size": 4,
+            "pp_size": 2,
+            "mem_fraction_static": 0.8,
+        },
+    )
+
+    deployment_options = SGLangServer.get_deployment_options(llm_config)
+    assert deployment_options["placement_group_bundles"] == custom_bundles, (
+        f"Expected custom bundles {custom_bundles}, "
+        f"got {deployment_options['placement_group_bundles']}"
+    )
+    assert deployment_options["placement_group_strategy"] == custom_strategy, (
+        f"Expected strategy '{custom_strategy}', "
+        f"got '{deployment_options['placement_group_strategy']}'"
+    )
+
+
+def test_sglang_custom_placement_group_default_strategy():
+    """Verify that custom bundles without an explicit strategy default to PACK."""
+    custom_bundles = [{"GPU": 1}] * 4
+
+    llm_config = LLMConfig(
+        model_loading_config={
+            "model_id": RAY_MODEL_ID,
+            "model_source": MODEL_ID,
+        },
+        server_cls=SGLangServer,
+        engine_kwargs={
+            "model_path": MODEL_ID,
+            "tp_size": 2,
+            "pp_size": 2,
+        },
+        placement_group_config={
+            "placement_group_bundles": custom_bundles,
+        },
+    )
+
+    deployment_options = SGLangServer.get_deployment_options(llm_config)
+    assert deployment_options["placement_group_bundles"] == custom_bundles
+    assert deployment_options["placement_group_strategy"] == "PACK"
+
+
+def _get_llm_handle(model_id: str = RAY_MODEL_ID):
+    """Return a Ray Serve handle to the LLMServer deployment for model_id."""
+    cleaned_id = model_id.replace("/", "--").replace(".", "_")
+    deployment_name = f"{SGLangServer.__name__}:{cleaned_id}"
+    return serve.get_deployment_handle(deployment_name, SERVE_DEFAULT_APP_NAME)
+
+
+@pytest.mark.asyncio
+async def test_sglang_pause_resume(sglang_client):
+    """Verify pause/resume lifecycle: server accepts requests before and after."""
+    handle = _get_llm_handle()
+
+    # Baseline: inference works before pause.
+    resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID, prompt="Hello", max_tokens=4, temperature=0.0
+    )
+    assert resp.choices[0].text is not None
+
+    # Pause with default mode ("abort").
+    await handle.pause.remote()
+    assert await handle.is_paused.remote() is True
+
+    # Resume and confirm state clears.
+    await handle.resume.remote()
+    assert await handle.is_paused.remote() is False
+
+    # Inference must work again after resume.
+    resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID, prompt="Hello", max_tokens=4, temperature=0.0
+    )
+    assert resp.choices[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_sglang_pause_resume_modes(sglang_client):
+    """Verify all three pause modes are accepted without error."""
+    handle = _get_llm_handle()
+
+    for mode in ("abort", "in_place", "retract"):
+        await handle.pause.remote(mode=mode)
+        assert await handle.is_paused.remote() is True
+        await handle.resume.remote()
+        assert await handle.is_paused.remote() is False
+
+
+@pytest.mark.asyncio
+async def test_sglang_sleep_wakeup(sglang_client):
+    """Verify sleep/wakeup lifecycle: GPU memory released then restored."""
+    handle = _get_llm_handle()
+
+    # Baseline inference.
+    resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID, prompt="Hello", max_tokens=4, temperature=0.0
+    )
+    assert resp.choices[0].text is not None
+
+    # Sleep (release all GPU memory).
+    await handle.sleep.remote()
+    assert await handle.is_sleeping.remote() is True
+
+    # Wakeup and confirm state clears.
+    await handle.wakeup.remote()
+    assert await handle.is_sleeping.remote() is False
+
+    # Inference must work again after wakeup.
+    resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID, prompt="Hello", max_tokens=4, temperature=0.0
+    )
+    assert resp.choices[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_sglang_sleep_wakeup_with_tags(sglang_client):
+    """Verify selective sleep/wakeup using component tags."""
+    handle = _get_llm_handle()
+
+    await handle.sleep.remote(tags=["kv_cache"])
+    assert await handle.is_sleeping.remote() is True
+
+    await handle.wakeup.remote(tags=["kv_cache"])
+    assert await handle.is_sleeping.remote() is False
+
+    resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID, prompt="Hello", max_tokens=4, temperature=0.0
+    )
+    assert resp.choices[0].text is not None
+
+
+@pytest.mark.asyncio
+async def test_sglang_reset_prefix_cache(sglang_client):
+    """Verify reset_prefix_cache completes and inference continues to work."""
+    handle = _get_llm_handle()
+
+    # Warm the cache with a request.
+    sglang_client.completions.create(
+        model=RAY_MODEL_ID,
+        prompt="The capital of France is",
+        max_tokens=8,
+        temperature=0.0,
+    )
+
+    # Flush the cache.
+    await handle.reset_prefix_cache.remote()
+
+    # Inference must still work after cache flush.
+    resp = sglang_client.completions.create(
+        model=RAY_MODEL_ID,
+        prompt="The capital of France is",
+        max_tokens=8,
+        temperature=0.0,
+    )
+    assert resp.choices[0].text.strip()
+
+
+# ---------------------------------------------------------------------------
+# Protocol decoupling tests — verify modules are importable without vLLM
+# and that SGLang protocol models are wired correctly.
+# ---------------------------------------------------------------------------
+
+
+class TestSGLangProtocolDecoupling:
+    """Verify modules are importable without vLLM and SGLang models are wired."""
+
+    def test_modules_importable_without_vllm(self):
+        """openai_api_models, ingress, llm_server, and ray.serve.llm should
+        all import without vLLM installed."""
+        from ray.llm._internal.serve.core.configs import openai_api_models  # noqa: F401
+        from ray.llm._internal.serve.core.ingress import ingress  # noqa: F401
+        from ray.llm._internal.serve.core.server.llm_server import LLMServer
+        import ray.serve.llm  # noqa: F401
+
+        assert LLMServer._default_engine_cls is None
+
+    def test_error_response_round_trip(self):
+        from ray.llm._internal.serve.core.configs.openai_api_models import (
+            ErrorInfo,
+            ErrorResponse,
+        )
+
+        resp = ErrorResponse(error=ErrorInfo(message="bad", code=400, type="Invalid"))
+        assert resp.error.message == "bad"
+        assert resp.error.code == 400
+        assert resp.model_dump()["error"]["message"] == "bad"
+
+    def test_score_request_is_sglang_scoring_request(self):
+        from sglang.srt.entrypoints.openai.protocol import ScoringRequest
+        from ray.llm._internal.serve.core.configs.openai_api_models import ScoreRequest
+
+        assert issubclass(ScoreRequest, ScoringRequest)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-xvs", __file__]))

release/llm_tests/serve/test_llm_serve_sleep_wakeup.py

@@ -0,0 +1,244 @@
+"""Test sleep/wakeup control plane API for Ray Serve LLM.
+
+This test verifies that the DevIngress sleep/wakeup endpoints work correctly:
+1. Engine starts in awake state (is_sleeping=False)
+2. Sleep command puts engine to sleep (is_sleeping=True) and frees GPU memory
+3. Wakeup command restores engine (is_sleeping=False) and restores GPU memory
+4. Model can still serve requests after wakeup
+
+NOTE (Kourosh): This is part of a design in progress for integrating Ray Serve
+LLM with RL workloads. The API is not public and won't be documented until the
+end-to-end story is finalized. Class names and endpoint names may change.
+"""
+
+import subprocess
+import time
+from typing import List
+
+import pytest
+import requests
+from openai import OpenAI
+from ray import serve
+from ray.llm._internal.serve.core.ingress.dev_ingress import build_dev_openai_app
+from ray.serve.llm import LLMConfig, ModelLoadingConfig
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.schema import ApplicationStatus
+
+MODEL_ID = "Qwen/Qwen2-0.5B-Instruct"
+BASE_URL = "http://localhost:8000"
+
+
+def get_llm_config() -> LLMConfig:
+    """Create LLMConfig with sleep mode enabled."""
+    return LLMConfig(
+        model_loading_config=ModelLoadingConfig(
+            model_id=MODEL_ID,
+        ),
+        deployment_config=dict(
+            num_replicas=2,
+        ),
+        engine_kwargs=dict(
+            tensor_parallel_size=2,
+            enable_sleep_mode=True,
+        ),
+    )
+
+
+def is_default_app_running():
+    """Check if the default application is running successfully."""
+    try:
+        default_app = serve.status().applications[SERVE_DEFAULT_APP_NAME]
+        return default_app.status == ApplicationStatus.RUNNING
+    except (KeyError, AttributeError):
+        return False
+
+
+def get_gpu_memory_used_mb() -> List[float]:
+    """Get GPU memory used per device via nvidia-smi.
+
+    Returns:
+        List of memory used in MB for each GPU device.
+    """
+    result = subprocess.run(
+        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return [float(x.strip()) for x in result.stdout.strip().split("\n") if x.strip()]
+
+
+def get_total_gpu_memory_mb() -> float:
+    """Get total GPU memory used across all devices."""
+    return sum(get_gpu_memory_used_mb())
+
+
+def wait_for_server_ready(timeout: int = 240) -> None:
+    """Wait for the server to be ready to handle requests."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            test_data = {
+                "model": MODEL_ID,
+                "messages": [{"role": "user", "content": "test"}],
+                "max_tokens": 5,
+                "temperature": 0,
+            }
+            response = requests.post(
+                f"{BASE_URL}/v1/chat/completions", json=test_data, timeout=10
+            )
+            if response.status_code == 200:
+                print(f"Server at {BASE_URL} is ready to handle requests!")
+                return
+        except Exception as e:
+            print(f"Waiting for server to be ready... (error: {e})")
+
+        time.sleep(2)
+
+    raise TimeoutError(
+        f"Server at {BASE_URL} did not become ready within {timeout} seconds"
+    )
+
+
+def test_sleep_wakeup_lifecycle():
+    """Test the complete sleep/wakeup lifecycle with GPU memory verification."""
+    # Start Ray Serve with DevIngress
+    llm_config = get_llm_config()
+    app = build_dev_openai_app({"llm_configs": [llm_config]})
+    serve.run(app, blocking=False)
+
+    # Wait for application to be running
+    wait_for_condition(is_default_app_running, timeout=300)
+    wait_for_server_ready(timeout=240)
+    time.sleep(5)  # Buffer time for server to be fully ready
+
+    try:
+        # Step 1: Verify initial state - engine should be awake
+        print("\n=== Step 1: Checking initial state ===")
+        response = requests.get(
+            f"{BASE_URL}/is_sleeping?model={MODEL_ID}",
+            timeout=10,
+        )
+        assert (
+            response.status_code == 200
+        ), f"is_sleeping returned {response.status_code}"
+        initial_sleep_state = response.json().get("is_sleeping", None)
+        assert (
+            initial_sleep_state is False
+        ), f"Expected is_sleeping=False, got {initial_sleep_state}"
+        print(f"✓ Initial sleeping state: {initial_sleep_state}")
+
+        # Step 2: Record baseline GPU memory
+        print("\n=== Step 2: Recording baseline GPU memory ===")
+        baseline_memory_mb = get_total_gpu_memory_mb()
+        print(f"Baseline GPU memory: {baseline_memory_mb:.2f} MB")
+        assert baseline_memory_mb > 0, "Baseline GPU memory should be > 0"
+
+        # Step 3: Put engine to sleep
+        print("\n=== Step 3: Putting engine to sleep ===")
+        sleep_response = requests.post(
+            f"{BASE_URL}/sleep",
+            json={"model": MODEL_ID, "options": {"level": 1}},
+            timeout=60,
+        )
+        assert (
+            sleep_response.status_code == 200
+        ), f"sleep returned {sleep_response.status_code}"
+        print("✓ Sleep command executed successfully")
+
+        # Wait a bit for sleep to complete
+        time.sleep(5)
+
+        # Step 4: Verify engine is sleeping
+        print("\n=== Step 4: Verifying engine is sleeping ===")
+        response = requests.get(
+            f"{BASE_URL}/is_sleeping?model={MODEL_ID}",
+            timeout=10,
+        )
+        assert response.status_code == 200
+        sleep_state = response.json().get("is_sleeping", None)
+        assert (
+            sleep_state is True
+        ), f"Expected is_sleeping=True after sleep, got {sleep_state}"
+        print(f"✓ Sleeping state: {sleep_state}")
+
+        # Step 5: Verify GPU memory reduction
+        print("\n=== Step 5: Verifying GPU memory reduction ===")
+        sleep_memory_mb = get_total_gpu_memory_mb()
+        memory_reduction_mb = baseline_memory_mb - sleep_memory_mb
+        memory_reduction_pct = (memory_reduction_mb / baseline_memory_mb) * 100
+        print(f"GPU memory after sleep: {sleep_memory_mb:.2f} MB")
+        print(
+            f"Memory reduction: {memory_reduction_mb:.2f} MB ({memory_reduction_pct:.1f}%)"
+        )
+        assert (
+            memory_reduction_pct > 50
+        ), f"Expected >50% memory reduction, got {memory_reduction_pct:.1f}%"
+        print("✓ GPU memory reduced significantly after sleep")
+
+        # Step 6: Wake up the engine
+        print("\n=== Step 6: Waking up engine ===")
+        wakeup_response = requests.post(
+            f"{BASE_URL}/wakeup",
+            json={"model": MODEL_ID, "options": {"tags": ["weights", "kv_cache"]}},
+            timeout=60,
+        )
+        assert (
+            wakeup_response.status_code == 200
+        ), f"wakeup returned {wakeup_response.status_code}"
+        print("✓ Wakeup command executed successfully")
+
+        # Wait a bit for wakeup to complete
+        time.sleep(5)
+
+        # Step 7: Verify engine is awake
+        print("\n=== Step 7: Verifying engine is awake ===")
+        response = requests.get(
+            f"{BASE_URL}/is_sleeping?model={MODEL_ID}",
+            timeout=10,
+        )
+        assert response.status_code == 200
+        wake_state = response.json().get("is_sleeping", None)
+        assert (
+            wake_state is False
+        ), f"Expected is_sleeping=False after wakeup, got {wake_state}"
+        print(f"✓ Sleeping state: {wake_state}")
+
+        # Step 8: Verify GPU memory restoration
+        print("\n=== Step 8: Verifying GPU memory restoration ===")
+        wake_memory_mb = get_total_gpu_memory_mb()
+        memory_diff_mb = abs(wake_memory_mb - baseline_memory_mb)
+        memory_diff_pct = (memory_diff_mb / baseline_memory_mb) * 100
+        print(f"GPU memory after wakeup: {wake_memory_mb:.2f} MB")
+        print(f"Baseline memory: {baseline_memory_mb:.2f} MB")
+        print(f"Memory difference: {memory_diff_mb:.2f} MB ({memory_diff_pct:.1f}%)")
+        assert (
+            memory_diff_pct < 20
+        ), f"Expected <20% memory difference from baseline, got {memory_diff_pct:.1f}%"
+        print("✓ GPU memory restored to near baseline after wakeup")
+
+        # Step 9: Verify model can still serve requests
+        print("\n=== Step 9: Verifying model can serve requests ===")
+        client = OpenAI(base_url=f"{BASE_URL}/v1", api_key="fake-key")
+        chat_response = client.chat.completions.create(
+            model=MODEL_ID,
+            messages=[{"role": "user", "content": "What is the capital of France?"}],
+            max_tokens=10,
+            temperature=0,
+        )
+        assert chat_response.choices[0].message.content is not None
+        print(
+            f"✓ Model successfully generated response: {chat_response.choices[0].message.content[:50]}..."
+        )
+
+        print("\n=== All tests passed! ===")
+
+    finally:
+        # Cleanup
+        serve.shutdown()
+        time.sleep(1)
+
+
+if __name__ == "__main__":
+    pytest.main(["-xvs", __file__])

release/llm_tests/serve/test_utils.py

@@ -0,0 +1,281 @@
+import json
+import logging
+import os
+import re
+import time
+import uuid
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Union
+
+import requests
+from openai import OpenAI
+
+import boto3
+import ray
+import yaml
+from anyscale import service
+from anyscale.compute_config.models import ComputeConfig
+from anyscale.service.models import ServiceState
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.utils import get_random_string
+
+
+logger = logging.getLogger(__file__)
+logging.basicConfig(level=logging.INFO)
+REGION_NAME = "us-west-2"
+SECRET_NAME = "llm_release_test_hf_token"
+
+# This bucket is on anyscale-dev-product account and the
+# anyscale-staging cloud is already configured to have write
+# access to this bucket
+# Buildkite is also configured to have read access to this bucket
+S3_BUCKET = "rayllm-ci-results"
+S3_PREFIX = "vllm_perf_results"
+ANYSCALE_JOB_CLUSTER_COMPUTE_NAME_ENV_VAR = "ANYSCALE_JOB_CLUSTER_COMPUTE_NAME"
+
+
+def check_service_state(
+    service_name: str, expected_state: ServiceState, cloud: Optional[str] = None
+):
+    """Check if the service is in the expected state."""
+    state = service.status(name=service_name, cloud=cloud).state
+    logger.info(
+        f"Waiting for service {service_name} to be {expected_state}, currently {state}"
+    )
+    assert (
+        state == expected_state
+    ), f"Service {service_name} is {state}, expected {expected_state}."
+    return True
+
+
+def terminate_service_if_running(service_name: str, cloud: Optional[str] = None):
+    try:
+        status = service.status(name=service_name, cloud=cloud)
+    except RuntimeError:
+        return
+
+    if status.state != ServiceState.TERMINATED:
+        logger.info(
+            f"Service {service_name} is in state {status.state}. Terminating it before running the benchmark."
+        )
+        service.terminate(name=service_name, cloud=cloud)
+        service.wait(name=service_name, cloud=cloud, state=ServiceState.TERMINATED)
+        logger.info(f"Service {service_name} is now terminated.")
+
+
+@contextmanager
+def timeit(stage: str, time_metrics: Dict[str, float]):
+    start = time.perf_counter()
+    yield
+    end = time.perf_counter()
+
+    duration = end - start
+    logger.info(f"Stage '{stage}' took {duration:.2f} seconds.")
+    time_metrics[f"time_{stage}"] = duration
+
+
+@contextmanager
+def start_service(
+    service_name: str,
+    compute_config: Union[ComputeConfig, str],
+    applications: List[Dict],
+    image_uri: Optional[str] = None,
+    working_dir: Optional[str] = None,
+    add_unique_suffix: bool = True,
+    cloud: Optional[str] = None,
+    env_vars: Optional[Dict[str, str]] = None,
+    timeout_s: int = 900,  # seconds
+):
+    """Starts an Anyscale Service with the specified configs.
+
+    Args:
+        service_name: Name of the Anyscale Service. The actual service
+            name may be modified if `add_unique_suffix` is True.
+        compute_config: The configuration for the hardware resources
+            that the cluster will utilize.
+        applications: The list of Ray Serve applications to run in the
+            service.
+        image_uri: The URI of the Docker image to use for the service.
+            If None, the image URI is fetched and constructed from the env var.
+        working_dir: The working directory for the service.
+        add_unique_suffix: Whether to append a unique suffix to the
+            service name.
+        cloud: The cloud to deploy the service to.
+        env_vars: The environment variables to set in the service.
+        timeout_s: The maximum time to wait for the service to start
+            and terminate, in seconds.
+    """
+
+    if add_unique_suffix:
+        ray_commit = (
+            ray.__commit__[:8] if ray.__commit__ != "{{RAY_COMMIT_SHA}}" else "nocommit"
+        )
+        service_name = f"{service_name}-{ray_commit}-{get_random_string()}"
+
+    if image_uri is None:
+        # We expect this environment variable to be set for all release tests
+        cluster_env = os.environ["ANYSCALE_JOB_CLUSTER_ENV_NAME"]
+        image_uri = f"anyscale/image/{cluster_env}:1"
+
+    time_metrics = {}
+    service_config = service.ServiceConfig(
+        name=service_name,
+        image_uri=image_uri,
+        compute_config=compute_config,
+        working_dir=working_dir,
+        applications=applications,
+        env_vars=env_vars,
+        query_auth_token_enabled=False,
+    )
+
+    # If the service already exists, terminate and the start a new service
+    # so the new service starts immediately. Otherwise, start a new service
+    # without a canary_percent.
+    terminate_service_if_running(service_name=service_name, cloud=cloud)
+
+    try:
+        logger.info(f"Service config: {service_config}")
+        with timeit("service_startup", time_metrics):
+            service.deploy(service_config)
+
+            wait_for_condition(
+                check_service_state,
+                service_name=service_name,
+                expected_state=ServiceState.RUNNING,
+                retry_interval_ms=10000,  # 10s
+                timeout=timeout_s,
+                cloud=cloud,
+            )
+
+        service_status = service.status(name=service_name, cloud=cloud)
+
+        yield {
+            "api_url": service_status.query_url,
+            "api_token": service_status.query_auth_token,
+            **time_metrics,
+        }
+
+    finally:
+        logger.info(f"Terminating service {service_name}.")
+        service.terminate(name=service_name, cloud=cloud)
+        wait_for_condition(
+            check_service_state,
+            service_name=service_name,
+            expected_state="TERMINATED",
+            retry_interval_ms=10000,  # 10s
+            timeout=timeout_s,
+            cloud=cloud,
+        )
+        logger.info(f"Service '{service_name}' terminated successfully.")
+
+
+def get_service_compute_config(
+    compute_config: Optional[str] = None,
+) -> str:
+    """Get the compute config to use when starting the Anyscale Service."""
+    service_compute_config = compute_config or os.environ.get(
+        ANYSCALE_JOB_CLUSTER_COMPUTE_NAME_ENV_VAR
+    )
+    if service_compute_config is None:
+        raise RuntimeError(
+            "No compute config was provided for the Anyscale Service. Set "
+            f"--compute-config or {ANYSCALE_JOB_CLUSTER_COMPUTE_NAME_ENV_VAR}; "
+            "release jobs should provide this automatically."
+        )
+    return service_compute_config
+
+
+def get_applications(serve_config_file: str) -> List[Any]:
+    """Get the applications from the serve config file."""
+    with open(serve_config_file, "r") as f:
+        loaded_llm_config = yaml.safe_load(f)
+    return loaded_llm_config["applications"]
+
+
+def setup_client_env_vars(api_url: str, api_token: Optional[str] = None):
+    """Set up the environment variables for the tests."""
+    os.environ["OPENAI_API_BASE"] = f"{api_url.rstrip('/')}/v1"
+    os.environ["OPENAI_API_KEY"] = api_token or "fake-key"
+
+
+def get_hf_token_env_var() -> Dict[str, str]:
+    """Get the environment variables for the service."""
+    session = boto3.session.Session()
+    client = session.client(service_name="secretsmanager", region_name=REGION_NAME)
+    secret_string = client.get_secret_value(SecretId=SECRET_NAME)["SecretString"]
+    return json.loads(secret_string)
+
+
+def get_python_version_from_image(image_name: str) -> str:
+    """Regex to capture the python version from the image name.
+
+    If the image name does not contain a python version, an empty string is returned.
+    """
+    if image_name is None:
+        return ""
+
+    image_python_version_regex_match = re.search(r"py[0-9]+", image_name)
+    if image_python_version_regex_match and image_python_version_regex_match.group(0):
+        return image_python_version_regex_match.group(0)
+
+    return ""
+
+
+def append_python_version_from_image(name: str, image_name: str) -> str:
+    """Regex to capture the python version from the image name and append it to the
+    given name.
+
+    If the image name does not contain a python version, the name is returned as is.
+    """
+    python_version = get_python_version_from_image(image_name)
+    if python_version:
+        return f"{name}_{python_version}"
+
+    return name
+
+
+def get_vllm_s3_storage_path() -> str:
+    build_number = os.environ.get(
+        "BUILDKITE_BUILD_NUMBER", uuid.uuid4().hex[:5].upper()
+    )
+    retry_count = os.environ.get("BUILDKITE_RETRY_COUNT", "0")
+    unique_id = f"build-{build_number}-{retry_count}"
+
+    storage_path = f"s3://{S3_BUCKET}/{S3_PREFIX}/vllm-perf-results-{unique_id}.jsonl"
+
+    return storage_path
+
+
+def create_openai_client(server_url: str) -> OpenAI:
+    return OpenAI(base_url=f"{server_url}/v1", api_key="fake-key")
+
+
+def wait_for_server_ready(
+    url: str,
+    model_id: str,
+    timeout: int = 300,
+    retry_interval: int = 2,
+) -> None:
+    """Poll the server until it's ready or timeout is reached."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            resp = requests.post(
+                f"{url}/v1/completions",
+                json={
+                    "model": model_id,
+                    "prompt": "test",
+                    "max_tokens": 5,
+                    "temperature": 0,
+                },
+                timeout=10,
+            )
+            if resp.status_code == 200:
+                print(f"Server at {url} is ready to handle requests!")
+                return
+        except Exception:
+            pass
+
+        print(f"Waiting for server at {url} to be ready...")
+        time.sleep(retry_interval)
+    raise TimeoutError(f"Server at {url} not ready within {timeout}s")