tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tests/e2e/test_local_disagg.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import time
+from dataclasses import asdict
+from unittest.mock import patch
+
+import pytest
+import vllm.envs as vllm_envs
+from vllm import LLM, EngineArgs, SamplingParams
+
+from tpu_inference.core.core_tpu import DisaggEngineCore, DisaggEngineCoreProc
+
+
+@pytest.fixture
+def test_prompts():
+    """Simple test prompts for disaggregated serving testing."""
+    return [
+        "Hello, my name is",
+        "The capital of France is",
+        "The colors of the rainbow are",
+        "The future of AI is",
+        "The president of the United States is",
+        "How many players are on a standard soccer team on the field at one time?",
+        "In Greek mythology, who is the god of the sea?",
+        "In what year did the Titanic sink?",
+        "In which museum is the Mona Lisa displayed?",
+        "Mount Everest is located in which mountain range?",
+        "What ancient empire was ruled by Julius Caesar?",
+        "What are the four fundamental forces of nature?",
+        'What does "CPU" stand for?',
+        'What does "HTML" stand for?',
+        "What is the capital of Australia?",
+        "What is the chemical symbol for gold?",
+        "What is the currency of Switzerland?",
+        "What is the distance from the Earth to the Sun called?",
+        "What is the freezing point of water in Celsius?",
+        "What is the hardest known natural substance on Earth?",
+        "What is the largest planet in our solar system?",
+        "What is the longest river in the world?",
+        "What is the main function of the kidneys in the human body?",
+        "What is the main ingredient in guacamole?",
+        "What is the most spoken language in the world by number of native speakers?",
+        "What is the process by which plants use sunlight to create food?",
+        "Which country is known as the Land of the Rising Sun?",
+        "Who developed the theory of general relativity?",
+        'Who directed the original "Star Wars" trilogy?',
+        "Who is credited with inventing the telephone?",
+        "Who painted the ceiling of the Sistine Chapel?",
+        "Who was the first female Prime Minister of the United Kingdom?",
+        "Who was the first person to walk on the moon?",
+        "Who wrote the American Declaration of Independence?",
+        'Who wrote the novel "Pride and Prejudice"?',
+    ]
+
+
+@pytest.fixture
+def sampling_params():
+    """Standard sampling parameters for testing."""
+    return SamplingParams(
+        temperature=0.0,
+        max_tokens=32,
+        ignore_eos=True,
+        logprobs=1,
+    )
+
+
+def test_disaggregated_serving(test_prompts, sampling_params):
+    """
+    Test disaggregated serving end-to-end.
+
+    Equivalent to:
+    PREFILL_SLICES=4 DECODE_SLICES=4 python examples/offline_inference.py \
+        --model=meta-llama/Meta-Llama-3.1-8B-Instruct --task=generate \
+        --max_model_len=2048 --tensor_parallel_size 4
+    """
+    # Set environment variables for disaggregated serving
+    # Using 4 slices for prefill and 4 for decode as requested
+    # Note: The user example used PREFILL_SLICES=4 DECODE_SLICES=4
+    # But usually slices are specified as "2x2" or similar if they are TPU topology.
+    # However, disagg_utils.py _parse_slices handles "4" as well (1D).
+    # We will stick to the user's example values.
+
+    # We need to mock the environment variables for this test
+    with patch.dict(
+            os.environ, {
+                "PREFILL_SLICES": "4",
+                "DECODE_SLICES": "4",
+                "SKIP_JAX_PRECOMPILE": "1",
+                "VLLM_XLA_CHECK_RECOMPILATION": "0"
+            }):
+        # Patch the EngineCore classes to use Disagg versions
+        with patch("vllm.v1.engine.core.EngineCore", DisaggEngineCore), \
+             patch("vllm.v1.engine.core.EngineCoreProc", DisaggEngineCoreProc):
+
+            model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+            os.system(f"rm -rf {vllm_envs.VLLM_XLA_CACHE_PATH}/*")
+            engine_args = EngineArgs(
+                model=model_name,
+                max_model_len=2048,
+                tensor_parallel_size=4,
+                gpu_memory_utilization=0.90,
+                enforce_eager=False,
+            )
+
+            llm = LLM(**asdict(engine_args))
+
+            try:
+                outputs = llm.generate(test_prompts, sampling_params)
+
+                # Verify outputs
+                assert len(outputs) == len(test_prompts)
+                for output in outputs:
+                    assert len(output.outputs) > 0
+                    assert len(output.outputs[0].text.strip()) > 0
+                    print(f"Prompt: {output.prompt!r}")
+                    print(f"Generated: {output.outputs[0].text!r}")
+
+            finally:
+                del llm
+                time.sleep(10)
+                pass
+
+
+def _run_inference(model_name: str,
+                   test_prompts: list,
+                   sampling_params: SamplingParams,
+                   tensor_parallel_size: int = 1,
+                   is_disagg: bool = False,
+                   prefill_slices: str = "4",
+                   decode_slices: str = "4") -> list:
+    """Helper function to run inference with specified configuration."""
+
+    # Define the inner execution logic
+    def run_inner():
+        engine_args = EngineArgs(
+            model=model_name,
+            max_model_len=2048,
+            tensor_parallel_size=tensor_parallel_size,
+            gpu_memory_utilization=0.90,
+            enforce_eager=False,
+        )
+
+        llm = LLM(**asdict(engine_args))
+        try:
+            return llm.generate(test_prompts, sampling_params)
+        finally:
+            del llm
+            time.sleep(10)
+            pass
+
+    if is_disagg:
+        # Mock environment variables and patch classes for disagg
+        with patch.dict(
+                os.environ, {
+                    "PREFILL_SLICES": prefill_slices,
+                    "DECODE_SLICES": decode_slices,
+                    "SKIP_JAX_PRECOMPILE": "1",
+                    "VLLM_XLA_CHECK_RECOMPILATION": "0"
+                }):
+            with patch("vllm.v1.engine.core.EngineCore", DisaggEngineCore), \
+                 patch("vllm.v1.engine.core.EngineCoreProc", DisaggEngineCoreProc):
+                return run_inner()
+    else:
+        # Run standard inference
+        # We still set some env vars to ensure consistent behavior if needed
+        # but for baseline we want it as standard as possible.
+        # However, to match the disagg run's potential jax settings:
+        with patch.dict(os.environ, {
+                "SKIP_JAX_PRECOMPILE": "1",
+                "VLLM_XLA_CHECK_RECOMPILATION": "0"
+        }):
+            return run_inner()
+
+
+def test_disaggregated_serving_correctness(test_prompts, sampling_params):
+    """
+    Test that disaggregated serving produces consistent results compared to a baseline.
+    """
+    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    # Use a smaller subset of prompts for correctness testing
+    small_prompts = test_prompts[:20]
+    sampling_params.max_tokens = 16
+
+    # Run baseline (standard execution)
+    # We use tensor_parallel_size=4 to match the disagg resources if we assume
+    # the user has enough chips, or if we are just mocking.
+    # Since the original test used tp=4, we stick to it.
+    print("Running Baseline Inference...")
+    baseline_outputs = _run_inference(model_name=model_name,
+                                      test_prompts=small_prompts,
+                                      sampling_params=sampling_params,
+                                      tensor_parallel_size=4,
+                                      is_disagg=False)
+
+    # Run disaggregated inference
+    os.system(f"rm -rf {vllm_envs.VLLM_XLA_CACHE_PATH}/*")
+    print("Running Disaggregated Inference...")
+
+    disagg_outputs = _run_inference(model_name=model_name,
+                                    test_prompts=small_prompts,
+                                    sampling_params=sampling_params,
+                                    tensor_parallel_size=4,
+                                    is_disagg=True,
+                                    prefill_slices="4",
+                                    decode_slices="4")
+
+    # Compare outputs
+    assert len(baseline_outputs) == len(disagg_outputs)
+
+    text_matches = 0
+    text_mismatches = 0
+    token_mismatches = 0
+
+    for i, (baseline,
+            disagg) in enumerate(zip(baseline_outputs, disagg_outputs)):
+        baseline_text = baseline.outputs[0].text.strip()
+        disagg_text = disagg.outputs[0].text.strip()
+
+        # Check text output
+        if baseline_text == disagg_text:
+            text_matches += 1
+        else:
+            text_mismatches += 1
+            print(f"Text mismatch found in prompt {i}:")
+            print(f"  Baseline: {baseline_text}")
+            print(f"  Disagg:   {disagg_text}")
+
+        # Check log probabilities (tokens) if available
+        baseline_logprobs = baseline.outputs[0].logprobs
+        disagg_logprobs = disagg.outputs[0].logprobs
+
+        if baseline_logprobs is not None and disagg_logprobs is not None:
+            assert len(baseline_logprobs) == len(disagg_logprobs), \
+                f"Logprobs length mismatch: {len(baseline_logprobs)} vs {len(disagg_logprobs)}"
+
+            for token_idx, (base_lp, disagg_lp) in enumerate(
+                    zip(baseline_logprobs, disagg_logprobs)):
+                if base_lp and disagg_lp:
+                    # Compare the top token IDs
+                    base_top_token = list(base_lp.keys())[0]
+                    disagg_top_token = list(disagg_lp.keys())[0]
+
+                    if base_top_token != disagg_top_token:
+                        token_mismatches += 1
+                        print(
+                            f"Token mismatch in prompt {i}, token {token_idx}:"
+                        )
+                        print(f"  Baseline: {base_top_token}")
+                        print(f"  Disagg:   {disagg_top_token}")
+
+    print("✓ Correctness test results:")
+    print(f"  Text: {text_matches} matches, {text_mismatches} mismatches")
+    print(f"  Token mismatches in logprobs: {token_mismatches}")
+    assert text_mismatches <= 5, f"Found {text_mismatches} text mismatches"
+    assert token_mismatches <= 40, f"Found {token_mismatches} token mismatches"

tests/e2e/test_model_loader.py

@@ -0,0 +1,268 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# tests/e2e/test_model_loader.py
+
+import os
+import re
+import signal
+import subprocess
+import sys
+import tempfile
+import time
+
+import pytest
+import requests
+import torch
+from flax import nnx
+from vllm.model_executor.models.registry import ModelRegistry
+
+from tpu_inference.models.common.model_loader import (_MODEL_REGISTRY,
+                                                      register_model)
+
+
+@pytest.fixture
+def cleanup_registries():
+    """Cleans up the model registries before and after each test."""
+    _MODEL_REGISTRY.clear()
+    # vLLM's ModelRegistry uses a class-level dictionary to store model classes.
+    # We need to clear it to ensure test isolation.
+    if hasattr(ModelRegistry, "models"):
+        ModelRegistry.models.clear()
+    yield
+    _MODEL_REGISTRY.clear()
+    if hasattr(ModelRegistry, "models"):
+        ModelRegistry.models.clear()
+
+
+class DummyGoodModel(nnx.Module):
+    """A valid model that conforms to the expected interface."""
+
+    def __init__(self, vllm_config=None, rng=None, mesh=None):
+        pass
+
+    def __call__(self,
+                 kv_caches=None,
+                 input_ids=None,
+                 attention_metadata=None):
+        pass
+
+
+def test_register_model_success(cleanup_registries):
+    """Tests that a valid model is registered successfully."""
+    arch = "DummyGoodModelForCausalLM"
+    register_model(arch, DummyGoodModel)
+
+    # Check tpu_inference registry
+    assert arch in _MODEL_REGISTRY
+
+    class MockModelConfig:
+
+        def __init__(self, architectures):
+            self.hf_config = self._MockHfConfig(architectures)
+            self.model_impl = "flax_nnx"
+
+        class _MockHfConfig:
+
+            def __init__(self, architectures):
+                self.architectures = architectures
+
+    model_config = MockModelConfig(architectures=[arch])
+    vllm_compatible_model, _ = ModelRegistry.resolve_model_cls(
+        architectures=[arch], model_config=model_config)
+    assert vllm_compatible_model is not None
+    assert issubclass(vllm_compatible_model, torch.nn.Module)
+    assert issubclass(vllm_compatible_model, DummyGoodModel)
+
+
+try:
+    # Attempt to import vLLM's interface validation function
+    from vllm.model_executor.models.interfaces_base import is_vllm_model
+    VLLM_INTERFACE_CHECK_AVAILABLE = True
+except ImportError:
+    VLLM_INTERFACE_CHECK_AVAILABLE = False
+
+
+@pytest.mark.skipif(not VLLM_INTERFACE_CHECK_AVAILABLE,
+                    reason="is_vllm_model could not be imported from vllm.")
+def test_registered_model_passes_vllm_interface_check(cleanup_registries):
+    """
+    Ensures the wrapped model passes vLLM's own interface validation.
+
+    This test is future-proof. If vLLM adds new requirements to its
+    model interface, this test will fail, signaling that the wrapper
+    in `register_model` needs to be updated.
+    """
+    arch = "DummyGoodModelForCausalLM"
+    register_model(arch, DummyGoodModel)
+
+    class MockModelConfig:
+
+        def __init__(self, architectures):
+            self.hf_config = self._MockHfConfig(architectures)
+            self.model_impl = "flax_nnx"
+
+        class _MockHfConfig:
+
+            def __init__(self, architectures):
+                self.architectures = architectures
+
+    model_config = MockModelConfig(architectures=[arch])
+    vllm_compatible_model, _ = ModelRegistry.resolve_model_cls(
+        architectures=[arch], model_config=model_config)
+
+    # This directly uses vLLM's checker, so it's always up-to-date.
+    # We assume is_vllm_model returns True for a valid model, and either
+    # returns False or raises an exception for an invalid one.
+    assert is_vllm_model(vllm_compatible_model)
+
+
+def _run_server_and_bench(model_name: str, model_impl_type: str,
+                          port: int) -> float:
+    env = os.environ.copy()
+    env["MODEL_IMPL_TYPE"] = model_impl_type
+
+    # Start server
+    server_cmd = [
+        sys.executable,
+        "-m",
+        "vllm.entrypoints.cli.main",
+        "serve",
+        model_name,
+        "--port",
+        str(port),
+        "--max-model-len",
+        "2048",
+        "--tensor-parallel-size",
+        "1",
+        "--disable-log-requests",
+        "--no-enable-prefix-caching",
+        "--gpu-memory-utilization",
+        "0.90",
+    ]
+
+    print(f"Starting server ({model_impl_type}) on port {port}...")
+    # Use a new process group so we can kill the server and its children
+    # Use temporary files for stdout/stderr to avoid pipe buffer deadlocks
+    stdout_file = tempfile.TemporaryFile(mode='w+b')
+    stderr_file = tempfile.TemporaryFile(mode='w+b')
+    server_process = subprocess.Popen(server_cmd,
+                                      env=env,
+                                      stdout=stdout_file,
+                                      stderr=stderr_file,
+                                      preexec_fn=os.setsid)
+
+    try:
+        # Wait for server to be ready
+        start_time = time.time()
+        server_ready = False
+        while time.time() - start_time < 600:  # 10 minutes timeout
+            try:
+                if requests.get(
+                        f"http://localhost:{port}/health").status_code == 200:
+                    server_ready = True
+                    break
+            except requests.exceptions.RequestException:
+                pass
+
+            if server_process.poll() is not None:
+                stdout_file.seek(0)
+                stderr_file.seek(0)
+                stdout = stdout_file.read().decode("utf-8", errors="replace")
+                stderr = stderr_file.read().decode("utf-8", errors="replace")
+                raise RuntimeError(
+                    f"Server process exited unexpectedly.\nStdout: {stdout}\nStderr: {stderr}"
+                )
+
+            time.sleep(5)
+
+        if not server_ready:
+            stdout_file.seek(0)
+            stderr_file.seek(0)
+            stdout = stdout_file.read().decode("utf-8", errors="replace")
+            stderr = stderr_file.read().decode("utf-8", errors="replace")
+            raise RuntimeError(
+                f"Server failed to start within timeout.\nStdout: {stdout}\nStderr: {stderr}"
+            )
+
+        print("Server is ready. Running benchmark...")
+
+        # Run benchmark
+        bench_cmd = [
+            "vllm", "bench", "serve", "--model", model_name, "--port",
+            str(port), "--dataset-name", "random", "--random-input-len", "50",
+            "--random-output-len", "128", "--num-prompts", "20"
+        ]
+
+        result = subprocess.run(bench_cmd,
+                                env=env,
+                                capture_output=True,
+                                text=True)
+
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"Benchmark failed.\nStdout: {result.stdout}\nStderr: {result.stderr}"
+            )
+
+        # Parse throughput
+        # Output example: "Request throughput (req/s): 12.34"
+        match = re.search(r"Request throughput \(req/s\):\s+([\d\.]+)",
+                          result.stdout)
+        if not match:
+            raise ValueError(
+                f"Could not parse throughput from output:\n{result.stdout}")
+
+        throughput = float(match.group(1))
+        return throughput
+
+    finally:
+        print("Stopping server...")
+        try:
+            os.killpg(os.getpgid(server_process.pid), signal.SIGTERM)
+        except ProcessLookupError:
+            pass
+        server_process.wait()
+        stdout_file.close()
+        stderr_file.close()
+        # Wait for TPU cleanup
+        time.sleep(5)
+
+
+def test_flax_nnx_vs_vllm_performance():
+    """
+    Compares the performance of flax_nnx and vllm model implementations.
+
+    This test ensures that the JAX-native (`flax_nnx`) implementation's
+    performance is not significantly different from the vLLM-native PyTorch
+    (`vllm`) implementation. It measures the request throughput for both
+    backends and asserts that the percentage
+    difference is within a reasonable threshold.
+    """
+    model_name = "Qwen/Qwen3-4B"
+    # This should be 2-3% but 6% reduces flakiness.
+    percentage_difference_threshold = 0.06
+
+    throughput_vllm = _run_server_and_bench(model_name, "vllm", 8001)
+    throughput_flax = _run_server_and_bench(model_name, "flax_nnx", 8002)
+
+    print(f"vLLM (PyTorch) throughput: {throughput_vllm:.2f} req/s.")
+    print(f"flax_nnx (JAX) throughput: {throughput_flax:.2f} req/s.")
+
+    percentage_diff = abs(throughput_flax - throughput_vllm) / throughput_vllm
+    print(f"Percentage difference in throughput: {percentage_diff:.2%}.")
+
+    assert percentage_diff < percentage_difference_threshold, (
+        f"The performance difference between flax_nnx and vllm is too high. "
+        f"Difference: {percentage_diff:.2%}, Threshold: {percentage_difference_threshold:.2%}"
+    )

tests/e2e/test_multi_modal_inference.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# A simplified example to run multi-modal inference and verify the output.
+# This script is a self-contained test that runs a single prompt and
+# compares the output to a known-good output.
+
+import difflib
+import os
+from dataclasses import asdict
+
+import pytest
+from vllm import LLM, EngineArgs, SamplingParams
+from vllm.assets.image import ImageAsset
+from vllm.multimodal.image import convert_image_mode
+
+# Expected partial text output from the model. This is based on a previous
+# run and is used for verification. The test is considered passed if the
+# generated output match with this text.
+EXPECTED_TEXT = (
+    "The image depicts a tall, cylindrical tower with a lattice-like structure, surrounded by cherry blossom trees in full bloom. The cherry blossoms are in various stages of opening, with pink petals covering the branches. The sky is clear and blue, providing a vibrant backdrop to the scene. The tower appears to be a significant landmark"
+)
+
+
+# NOTE: Could be extended to more mm models/configs as needed
+@pytest.mark.parametrize("enable_dynamic_image_sizes", [False, True])
+def test_multi_modal_inference(monkeypatch, enable_dynamic_image_sizes):
+    """
+    Runs multi-modal inference and verifies the output.
+    """
+    os.environ['SKIP_JAX_PRECOMPILE'] = '1'  # Skip warmup to save time.
+    os.environ[
+        'VLLM_XLA_CHECK_RECOMPILATION'] = '0'  # Allow compilation during execution.
+
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    # --- Configuration ---
+    model = "Qwen/Qwen2.5-VL-3B-Instruct"
+    tensor_parallel_size = 1
+    temperature = 0.0
+    max_tokens = 64
+    max_model_len = 4096
+    gpu_memory_utilization = 0.5
+    modality = "image"
+
+    print("Preparing for multi-modal inference...")
+
+    # --- Prepare Inputs ---
+    image = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+    question = "What is the content of this image?"
+
+    # Using Qwen2.5-VL prompt template
+    # NOTE: other models may be different
+    prompt = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+              f"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+              f"{question}<|im_end|>\n"
+              "<|im_start|>assistant\n")
+
+    # --- Setup vLLM Engine ---
+    engine_args = EngineArgs(
+        model=model,
+        max_model_len=max_model_len,
+        tensor_parallel_size=tensor_parallel_size,
+        gpu_memory_utilization=gpu_memory_utilization,
+        max_num_seqs=1,
+        mm_processor_kwargs={
+            "min_pixels": 28 * 28,
+            "max_pixels": 1280 * 28 * 28,
+            "fps": 1,
+        },
+        limit_mm_per_prompt={modality: 1},
+    )
+    engine_args = asdict(engine_args)
+    if engine_args.get("additional_config") is None:
+        engine_args["additional_config"] = {}
+
+    engine_args["additional_config"][
+        "enable_dynamic_image_sizes"] = enable_dynamic_image_sizes
+    llm = LLM(**engine_args)
+
+    sampling_params = SamplingParams(
+        temperature=temperature,
+        max_tokens=max_tokens,
+    )
+
+    inputs = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "image": image
+        },
+    }
+
+    # --- Run Inference ---
+    print("Running inference...")
+    outputs = llm.generate(inputs, sampling_params)
+
+    # --- Verification ---
+    generated_text = outputs[0].outputs[0].text.strip()
+
+    print("-" * 50)
+    print("Generated Text:")
+    print(generated_text)
+    print("-" * 50)
+
+    # Check output
+    similarity_score = difflib.SequenceMatcher(None, generated_text,
+                                               EXPECTED_TEXT).ratio()
+    print(f"Similarity Score: {similarity_score:.4f}")
+    assert similarity_score >= 0.85, (
+        f"Text similarity too low ({similarity_score:.2f}).\n"
+        f"Expected: {EXPECTED_TEXT}\n"
+        f"Actual:   {generated_text}")