tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tests/lora/test_lora.py

@@ -0,0 +1,133 @@
+# https://github.com/vllm-project/vllm/blob/ed10f3cea199a7a1f3532fbe367f5c5479a6cae9/tests/tpu/lora/test_lora.py
+import os
+import time
+
+import pytest
+import vllm
+from vllm.lora.request import LoRARequest
+
+# This file contains tests to ensure that LoRA works correctly on the TPU
+# backend. We use a series of custom trained adapters for Qwen2.5-3B-Instruct
+# for this. The adapters are:
+# Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter, where x ranges
+# from 1 to 4.
+
+# These adapters are trained using a standard huggingface peft training script,
+# where all the inputs are "What is 1+1? \n" and all the outputs are "x". We run
+# 100 training iterations with a training batch size of 100.
+
+
+def setup_vllm(num_loras: int, tp: int = 1) -> vllm.LLM:
+    return vllm.LLM(model="Qwen/Qwen2.5-3B-Instruct",
+                    max_model_len=256,
+                    max_num_batched_tokens=64,
+                    max_num_seqs=8,
+                    tensor_parallel_size=tp,
+                    enable_lora=True,
+                    max_loras=num_loras,
+                    max_lora_rank=8)
+
+
+# For multi-chip test, we only use TP=2 because the base model Qwen/Qwen2.5-3B-Instruct has 2 kv heads and the current attention kernel requires it to be divisible by tp_size.
+TP = [2] if os.environ.get("USE_V6E8_QUEUE", False) else [1]
+
+
+@pytest.mark.parametrize("tp", TP)
+def test_single_lora(tp):
+    """
+    This test ensures we can run a single LoRA adapter on the TPU backend.
+    We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_2_adapter" which
+    will force Qwen2.5-3B-Instruct to claim 1+1=2.
+    """
+
+    llm = setup_vllm(1, tp)
+
+    prompt = "What is 1+1? \n"
+
+    lora_request = LoRARequest(
+        "lora_adapter_2", 2,
+        "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_2_adapter")
+    output = llm.generate(prompt,
+                          sampling_params=vllm.SamplingParams(max_tokens=16,
+                                                              temperature=0),
+                          lora_request=lora_request)[0].outputs[0].text
+
+    answer = output.strip()[0]
+
+    assert answer.isdigit()
+    assert int(answer) == 2
+
+    del llm
+    time.sleep(10)
+
+
+@pytest.mark.parametrize("tp", TP)
+def test_lora_hotswapping(tp):
+    """
+    This test ensures we can run multiple LoRA adapters on the TPU backend, even
+    if we only have space to store 1.
+
+    We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter" which
+    will force Qwen2.5-3B-Instruct to claim 1+1=x, for a range of x.
+    """
+
+    lora_name_template = \
+        "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_{}_adapter"
+    lora_requests = [
+        LoRARequest(f"lora_adapter_{i}", i, lora_name_template.format(i))
+        for i in range(1, 5)
+    ]
+
+    llm = setup_vllm(1, tp)
+
+    prompt = "What is 1+1? \n"
+
+    for i, req in enumerate(lora_requests):
+        output = llm.generate(prompt,
+                              sampling_params=vllm.SamplingParams(
+                                  max_tokens=16, temperature=0),
+                              lora_request=req)[0].outputs[0].text
+        answer = output.strip()[0]
+
+        assert answer.isdigit()
+        assert int(answer) == i + 1, f"Expected {i + 1}, got {answer}"
+
+    del llm
+    time.sleep(10)
+
+
+@pytest.mark.parametrize("tp", TP)
+def test_multi_lora(tp):
+    """
+    This test ensures we can run multiple LoRA adapters on the TPU backend, when
+    we have enough space to store all of them.
+
+    We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter" which
+    will force Qwen2.5-3B-Instruct to claim 1+1=x, for a range of x.
+    """
+    lora_name_template = \
+        "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_{}_adapter"
+    lora_requests = [
+        LoRARequest(f"lora_adapter_{i}", i, lora_name_template.format(i))
+        for i in range(1, 5)
+    ]
+
+    llm = setup_vllm(4, tp)
+
+    prompt = "What is 1+1? \n"
+
+    for i, req in enumerate(lora_requests):
+        output = llm.generate(prompt,
+                              sampling_params=vllm.SamplingParams(
+                                  max_tokens=16, temperature=0),
+                              lora_request=req)[0].outputs[0].text
+
+        answer = output.strip()[0]
+
+        assert answer.isdigit()
+        assert int(
+            output.strip()
+            [0]) == i + 1, f"Expected {i + 1}, got {int(output.strip()[0])}"
+
+    del llm
+    time.sleep(10)

tests/lora/utils.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
+
+
+# https://github.com/vllm-project/vllm/blob/279a5f31b3faa6f40759516efa5c742f637ab8b7/tests/lora/utils.py
+class DummyLoRAManager:
+
+    def __init__(self, device: torch.device = "cuda:0"):
+        super().__init__()
+        self._loras: dict[str, LoRALayerWeights] = {}
+        self._device = device
+
+    def set_module_lora(self, module_name: str, lora: LoRALayerWeights):
+        self._loras[module_name] = lora
+
+    def get_module_lora(self, module_name: str) -> LoRALayerWeights:
+        return self._loras[module_name]
+
+    def init_random_lora(
+        self,
+        module_name: str,
+        weight: torch.Tensor,
+        rank: int = 8,
+        generate_embeddings_tensor: int = 0,
+    ):
+        lora = LoRALayerWeights(
+            module_name,
+            rank=rank,
+            lora_alpha=1,
+            lora_a=torch.rand([rank, weight.shape[1]],
+                              dtype=weight.dtype,
+                              device=self._device),
+            lora_b=torch.rand([weight.shape[0], rank],
+                              dtype=weight.dtype,
+                              device=self._device),
+        )
+        if generate_embeddings_tensor:
+            lora.embeddings_tensor = torch.rand(
+                5,
+                generate_embeddings_tensor,
+                dtype=weight.dtype,
+                device=self._device,
+            )
+        self.set_module_lora(module_name, lora)
+
+        return lora
+
+    def init_lora(
+        self,
+        module_name: str,
+        input_dim: int,
+        output_dim: int,
+        rank=8,
+        noop=False,
+        embeddings_tensor=None,
+    ):
+        lora = LoRALayerWeights(
+            module_name,
+            rank=rank,
+            lora_alpha=1,
+            lora_a=torch.rand([rank, input_dim], device="cuda"),
+            lora_b=torch.rand([output_dim, input_dim], device="cuda"),
+            embeddings_tensor=embeddings_tensor,
+        )
+        self.set_module_lora(module_name, lora)
+        return lora
+
+    def reset_lora(self):
+        self._loras = {}
+
+    def init_packed_lora(
+        self,
+        module_name: str,
+        input_dim: int,
+        output_dims: list[int],
+        noop_lora_index: list[int] | None = None,
+        rank: int = 8,
+    ):
+        base_loras: list[LoRALayerWeights] = []
+        noop_lora_index_set = set(noop_lora_index or [])
+
+        for i, out_dim in enumerate(output_dims):
+            base_lora = self.init_lora(
+                module_name + "_000_" + str(i),
+                input_dim,
+                out_dim,
+                rank=rank,
+                noop=i in noop_lora_index_set,
+            )
+            base_loras.append(base_lora)
+        packed_lora = PackedLoRALayerWeights.pack(base_loras)
+        self.set_module_lora(module_name, packed_lora)
+        return packed_lora

tests/test_base.py

@@ -0,0 +1,201 @@
+import logging
+import unittest
+import warnings
+from dataclasses import dataclass, field, fields
+from typing import Any, List, Mapping
+
+from tpu_inference.layers.jax.base import Config
+
+# Use the 'warnings' module to globally ignore warnings within this block
+vllm_logger = logging.getLogger("vllm")
+original_level = vllm_logger.level
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+
+    # Set the vLLM logger to ERROR to suppress its messages
+    vllm_logger.setLevel(logging.ERROR)
+
+    # Import the class; all warnings will be suppressed
+    from vllm.config import ModelConfig
+
+vllm_logger.setLevel(logging.WARNING)
+
+
+def setup_vllm_config(subconfig_types: List[str],
+                      overrides: List[Mapping[str, Any]]):
+    vllm_config = SimpleVllmConfig()
+    for (subconfig_type, override) in zip(subconfig_types, overrides):
+        if subconfig_type == "model":
+            for key in override:
+                setattr(vllm_config.model_config, key, override[key])
+        else:
+            for key in override:
+                setattr(vllm_config, key, override[key])
+    return vllm_config
+
+
+@dataclass
+class SimpleVllmConfig():
+    additional_config: Mapping[str, Any] = field(default_factory=dict)
+    # Set default max_model_len to turn off warnings.
+    model_config: ModelConfig = field(
+        default_factory=lambda: ModelConfig(max_model_len=1024))
+
+
+@dataclass
+class SimpleConfig(Config):
+    vllm_config: SimpleVllmConfig
+    arg1: str
+    arg2: str
+    arg3: int
+
+    def is_equal(self, other: Config):
+        for f in fields(self):
+            if f.name != "vllm_config":
+                if getattr(self, f.name) != getattr(other, f.name):
+                    return False
+        return True
+
+
+class ConfigOverrideTests(unittest.TestCase):
+
+    def test_additional_config_overrides(self):
+        subconfig_types = ['']
+        overrides = [{"additional_config": {"arg1": "val1", "arg2": "val2"}}]
+        override_vllm_config = setup_vllm_config(subconfig_types, overrides)
+        default_vllm_config = SimpleVllmConfig()
+        config = SimpleConfig(vllm_config=override_vllm_config,
+                              arg1="foo",
+                              arg2="bar",
+                              arg3=123)
+        expected_config = SimpleConfig(vllm_config=default_vllm_config,
+                                       arg1="val1",
+                                       arg2="val2",
+                                       arg3=123)
+        self.assertTrue(config.is_equal(expected_config))
+
+    def test_hf_overrides(self):
+        subconfig_types = ['model']
+        overrides = [{"hf_overrides": {"arg2": "val2", "arg3": 456}}]
+        default_vllm_config = SimpleVllmConfig()
+        override_vllm_config = setup_vllm_config(subconfig_types, overrides)
+        config = SimpleConfig(vllm_config=override_vllm_config,
+                              arg1="foo",
+                              arg2="bar",
+                              arg3=123)
+        expected_config = SimpleConfig(vllm_config=default_vllm_config,
+                                       arg1="foo",
+                                       arg2="val2",
+                                       arg3=456)
+        self.assertTrue(config.is_equal(expected_config))
+
+    def test_additional_and_hf_overrides(self):
+        subconfig_types = ['', 'model']
+        overrides = [{
+            "additional_config": {
+                "arg1": "val1",
+                "arg2": "val2"
+            }
+        }, {
+            "hf_overrides": {
+                "arg2": "val3",
+                "arg3": 456
+            }
+        }]
+        default_vllm_config = SimpleVllmConfig()
+        override_vllm_config = setup_vllm_config(subconfig_types, overrides)
+        config = SimpleConfig(vllm_config=override_vllm_config,
+                              arg1="foo",
+                              arg2="bar",
+                              arg3=123)
+        expected_config = SimpleConfig(vllm_config=default_vllm_config,
+                                       arg1="val1",
+                                       arg2="val3",
+                                       arg3=456)
+        self.assertTrue(config.is_equal(expected_config))
+
+    def test_additional_and_generate_overrides(self):
+        subconfig_types = ['', 'model']
+        overrides = [{
+            "additional_config": {
+                "arg1": "val1",
+                "arg2": "val2"
+            }
+        }, {
+            "override_generation_config": {
+                "arg2": "val3",
+                "arg3": 456
+            }
+        }]
+        default_vllm_config = SimpleVllmConfig()
+        override_vllm_config = setup_vllm_config(subconfig_types, overrides)
+        config = SimpleConfig(vllm_config=override_vllm_config,
+                              arg1="foo",
+                              arg2="bar",
+                              arg3=123)
+        expected_config = SimpleConfig(vllm_config=default_vllm_config,
+                                       arg1="val1",
+                                       arg2="val3",
+                                       arg3=456)
+        self.assertTrue(config.is_equal(expected_config))
+
+    def test_hf_and_generate_overrides(self):
+        subconfig_types = ['model', 'model']
+        overrides = [{
+            "hf_overrides": {
+                "arg2": "val2",
+                "arg3": 456
+            }
+        }, {
+            "override_generation_config": {
+                "arg2": "val4",
+                "arg3": 789
+            }
+        }]
+        default_vllm_config = SimpleVllmConfig()
+        override_vllm_config = setup_vllm_config(subconfig_types, overrides)
+        config = SimpleConfig(vllm_config=override_vllm_config,
+                              arg1="foo",
+                              arg2="bar",
+                              arg3=123)
+        expected_config = SimpleConfig(vllm_config=default_vllm_config,
+                                       arg1="foo",
+                                       arg2="val4",
+                                       arg3=789)
+        self.assertTrue(config.is_equal(expected_config))
+
+    def test_additional_and_hf_and_generate_overrides(self):
+        subconfig_types = ['', 'model', 'model']
+        overrides = [{
+            "additional_config": {
+                "arg1": "val1",
+                "arg2": "val2"
+            }
+        }, {
+            "hf_overrides": {
+                "arg2": "val2",
+                "arg3": 456
+            }
+        }, {
+            "override_generation_config": {
+                "arg1": "val3",
+                "arg2": "val4",
+                "arg3": 789
+            }
+        }]
+        default_vllm_config = SimpleVllmConfig()
+        override_vllm_config = setup_vllm_config(subconfig_types, overrides)
+        config = SimpleConfig(vllm_config=override_vllm_config,
+                              arg1="foo",
+                              arg2="bar",
+                              arg3=123)
+        expected_config = SimpleConfig(vllm_config=default_vllm_config,
+                                       arg1="val3",
+                                       arg2="val4",
+                                       arg3=789)
+        self.assertTrue(config.is_equal(expected_config))
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_envs.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the tpu-inference project
+
+import pytest
+
+import tpu_inference.envs as envs
+from tpu_inference.envs import enable_envs_cache, environment_variables
+
+
+def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch):
+    assert envs.JAX_PLATFORMS == ""
+    assert envs.PHASED_PROFILING_DIR == ""
+    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
+    monkeypatch.setenv("PHASED_PROFILING_DIR", "/tmp/profiling")
+    assert envs.JAX_PLATFORMS == "tpu"
+    assert envs.PHASED_PROFILING_DIR == "/tmp/profiling"
+
+    assert envs.TPU_NAME is None
+    assert envs.TPU_ACCELERATOR_TYPE is None
+    monkeypatch.setenv("TPU_NAME", "my-tpu")
+    monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v5litepod-16")
+    assert envs.TPU_NAME == "my-tpu"
+    assert envs.TPU_ACCELERATOR_TYPE == "v5litepod-16"
+
+    # __getattr__ is not decorated with functools.cache
+    assert not hasattr(envs.__getattr__, "cache_info")
+
+
+def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
+    monkeypatch.setenv("TPU_NAME", "my-tpu")
+
+    # __getattr__ is not decorated with functools.cache
+    assert not hasattr(envs.__getattr__, "cache_info")
+
+    enable_envs_cache()
+
+    # __getattr__ is decorated with functools.cache
+    assert hasattr(envs.__getattr__, "cache_info")
+    start_hits = envs.__getattr__.cache_info().hits
+
+    # 2 more hits due to JAX_PLATFORMS and TPU_NAME accesses
+    assert envs.JAX_PLATFORMS == "tpu"
+    assert envs.TPU_NAME == "my-tpu"
+    assert envs.__getattr__.cache_info().hits == start_hits + 2
+
+    # All environment variables are cached
+    for environment_variable in environment_variables:
+        envs.__getattr__(environment_variable)
+    assert envs.__getattr__.cache_info(
+    ).hits == start_hits + 2 + len(environment_variables)
+
+    # Reset envs.__getattr__ back to non-cached version to
+    # avoid affecting other tests
+    envs.__getattr__ = envs.__getattr__.__wrapped__
+
+
+def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
+    # Test SKIP_JAX_PRECOMPILE (default False)
+    assert envs.SKIP_JAX_PRECOMPILE is False
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "1")
+    assert envs.SKIP_JAX_PRECOMPILE is True
+    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0")
+    assert envs.SKIP_JAX_PRECOMPILE is False
+
+    # Test NEW_MODEL_DESIGN (default False)
+    assert envs.NEW_MODEL_DESIGN is False
+    monkeypatch.setenv("NEW_MODEL_DESIGN", "1")
+    assert envs.NEW_MODEL_DESIGN is True
+
+    # Test USE_MOE_EP_KERNEL (default False)
+    assert envs.USE_MOE_EP_KERNEL is False
+    monkeypatch.setenv("USE_MOE_EP_KERNEL", "1")
+    assert envs.USE_MOE_EP_KERNEL is True
+
+
+def test_integer_env_vars(monkeypatch: pytest.MonkeyPatch):
+    assert envs.PYTHON_TRACER_LEVEL == 1
+    monkeypatch.setenv("PYTHON_TRACER_LEVEL", "3")
+    assert envs.PYTHON_TRACER_LEVEL == 3
+    monkeypatch.setenv("PYTHON_TRACER_LEVEL", "0")
+    assert envs.PYTHON_TRACER_LEVEL == 0
+
+
+def test_lowercase_conversion(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "GRPC")
+    assert envs.TPU_MULTIHOST_BACKEND == "grpc"
+
+    monkeypatch.setenv("MODEL_IMPL_TYPE", "FLAX_NNX")
+    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
+
+
+def test_string_env_vars_defaults(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("JAX_PLATFORMS", raising=False)
+    monkeypatch.delenv("PREFILL_SLICES", raising=False)
+    monkeypatch.delenv("DECODE_SLICES", raising=False)
+
+    assert envs.JAX_PLATFORMS == ""
+    assert envs.PREFILL_SLICES == ""
+    assert envs.DECODE_SLICES == ""
+    assert envs.PHASED_PROFILING_DIR == ""
+
+
+def test_none_default_env_vars(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("TPU_ACCELERATOR_TYPE", raising=False)
+    monkeypatch.delenv("TPU_NAME", raising=False)
+    monkeypatch.delenv("TPU_WORKER_ID", raising=False)
+
+    assert envs.TPU_ACCELERATOR_TYPE is None
+    assert envs.TPU_NAME is None
+    assert envs.TPU_WORKER_ID is None
+
+
+def test_ray_env_vars(monkeypatch: pytest.MonkeyPatch):
+    assert envs.RAY_USAGE_STATS_ENABLED == "0"
+    monkeypatch.setenv("RAY_USAGE_STATS_ENABLED", "1")
+    assert envs.RAY_USAGE_STATS_ENABLED == "1"
+
+    assert envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "shm"
+    monkeypatch.setenv("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "nccl")
+    assert envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE == "nccl"
+
+
+def test_invalid_attribute_raises_error():
+    with pytest.raises(AttributeError,
+                       match="has no attribute 'NONEXISTENT_VAR'"):
+        _ = envs.NONEXISTENT_VAR
+
+
+def test_dir_returns_all_env_vars():
+    env_vars = envs.__dir__()
+    assert isinstance(env_vars, list)
+    assert len(env_vars) == len(environment_variables)
+    assert "JAX_PLATFORMS" in env_vars
+    assert "TPU_NAME" in env_vars
+    assert "SKIP_JAX_PRECOMPILE" in env_vars
+    assert "MODEL_IMPL_TYPE" in env_vars
+
+
+def test_tpu_multihost_env_vars(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("TPU_WORKER_ID", "0")
+    assert envs.TPU_WORKER_ID == "0"
+
+    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "grpc")
+    assert envs.TPU_MULTIHOST_BACKEND == "grpc"
+
+    monkeypatch.setenv("TPU_MULTIHOST_BACKEND", "xla")
+    assert envs.TPU_MULTIHOST_BACKEND == "xla"
+
+
+def test_disaggregated_serving_env_vars(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PREFILL_SLICES", "0,1,2,3")
+    assert envs.PREFILL_SLICES == "0,1,2,3"
+
+    monkeypatch.setenv("DECODE_SLICES", "4,5,6,7")
+    assert envs.DECODE_SLICES == "4,5,6,7"
+
+
+def test_model_impl_type_default(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.delenv("MODEL_IMPL_TYPE", raising=False)
+    assert envs.MODEL_IMPL_TYPE == "flax_nnx"
+
+
+def test_cache_preserves_values_across_env_changes(
+        monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("JAX_PLATFORMS", "tpu")
+
+    enable_envs_cache()
+
+    assert envs.JAX_PLATFORMS == "tpu"
+
+    # Change environment variable
+    monkeypatch.setenv("JAX_PLATFORMS", "cpu")
+
+    # Cached value should still be "tpu"
+    assert envs.JAX_PLATFORMS == "tpu"
+
+    # Reset envs.__getattr__ back to non-cached version
+    envs.__getattr__ = envs.__getattr__.__wrapped__
+
+    # Now it should reflect the new value
+    assert envs.JAX_PLATFORMS == "cpu"