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

tpu_inference/distributed/utils.py

@@ -0,0 +1,112 @@
+# 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.
+
+import os
+
+from vllm.utils.network_utils import get_ip
+
+from tpu_inference import envs
+from tpu_inference.logger import init_logger
+
+logger = init_logger(__name__)
+
+# For multi-host usage only, to collect IP and port for all nodes.
+_NODES_KV_IP_PORT = dict()
+
+
+def set_node_kv_ip_port(ip_port: tuple[int, str, int]):
+    global _NODES_KV_IP_PORT
+    node_id, ip, port = ip_port
+    _NODES_KV_IP_PORT[node_id] = (ip, port)
+
+
+def get_kv_ips() -> str:
+    if envs.TPU_MULTIHOST_BACKEND == "ray":
+        num_nodes = len(_NODES_KV_IP_PORT)
+        ips = []
+        for node_id in range(num_nodes):
+            ips.append(_NODES_KV_IP_PORT[node_id][0])
+        return ips
+    else:
+        return get_host_ip()
+
+
+def get_kv_ports() -> str:
+    if envs.TPU_MULTIHOST_BACKEND == "ray":
+        num_nodes = len(_NODES_KV_IP_PORT)
+        ports = []
+        for node_id in range(num_nodes):
+            ports.append(_NODES_KV_IP_PORT[node_id][1])
+        return ports
+    else:
+        return get_kv_transfer_port()
+
+
+def get_host_ip() -> str:
+    """Use `VLLM_HOST_IP` if set, otherwise use default network interface IP."""
+    return get_ip()
+
+
+def get_kv_transfer_port() -> str:
+    port = os.getenv("TPU_KV_TRANSFER_PORT", "9100")
+    return port
+
+
+def get_side_channel_port() -> str:
+    port = os.getenv("TPU_SIDE_CHANNEL_PORT", "9600")
+    return port
+
+
+def get_device_topology_order_id(local_devices, global_devices) -> int:
+    """
+    Calculates the topology order ID for the local device set within the global topology.
+
+    This function determines the rank of the current host/process based on the
+    coordinate of its TPU devices relative to all devices in the topology.
+
+    Args:
+        local_devices: A list of TpuDevice objects available to the current process.
+        global_devices: A list of all TpuDevice objects in the global topology.
+
+    Returns:
+        The topology order ID (rank) of the local devices.
+    """
+    if not local_devices:
+        raise ValueError("local_devices cannot be empty")
+    if not global_devices:
+        raise ValueError("global_devices cannot be empty")
+
+    # 1. Find the 'anchor' (minimum coordinate) for the local devices.
+    #    This represents the physical top-left corner of the local machine.
+    local_anchor = min(d.coords for d in local_devices)
+
+    # 2. Group global devices by process to find the anchor for EVERY process.
+    process_anchors = {}
+    for d in global_devices:
+        pid = d.process_index
+        # Update the minimum coordinate found for this process so far
+        if pid not in process_anchors or d.coords < process_anchors[pid]:
+            process_anchors[pid] = d.coords
+
+    # 3. Sort the unique anchors to establish the canonical topology order.
+    #    Tuples (x, y, z) sort lexicographically (x first, then y, then z).
+    sorted_anchors = sorted(process_anchors.values())
+
+    # 4. Return the index (rank) of the local anchor in the sorted list.
+    try:
+        return sorted_anchors.index(local_anchor)
+    except ValueError:
+        raise ValueError(
+            f"Local devices: {local_devices} do not exist in the global device: {global_devices} list."
+        )

tpu_inference/env_override.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the tpu-inference project
+
+import os
+
+# Disable CUDA-specific shared experts stream for TPU
+# This prevents errors when trying to create CUDA streams on TPU hardware
+# The issue was introduced by vllm-project/vllm#26440
+os.environ["VLLM_DISABLE_SHARED_EXPERTS_STREAM"] = "1"

tpu_inference/envs.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the tpu-inference project
+
+import functools
+import os
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    JAX_PLATFORMS: str = ""
+    TPU_ACCELERATOR_TYPE: str | None = None
+    TPU_NAME: str | None = None
+    TPU_WORKER_ID: str | None = None
+    TPU_MULTIHOST_BACKEND: str = ""
+    PREFILL_SLICES: str = ""
+    DECODE_SLICES: str = ""
+    SKIP_JAX_PRECOMPILE: bool = False
+    VLLM_XLA_CHECK_RECOMPILATION: bool = False
+    MODEL_IMPL_TYPE: str = "auto"
+    NEW_MODEL_DESIGN: bool = False
+    PHASED_PROFILING_DIR: str = ""
+    PYTHON_TRACER_LEVEL: int = 1
+    USE_MOE_EP_KERNEL: bool = False
+    NUM_SLICES: int = 1
+    RAY_USAGE_STATS_ENABLED: str = "0"
+    VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "shm"
+    ENABLE_QUANTIZED_MATMUL_KERNEL: bool = False
+
+
+def env_with_choices(
+    env_name: str,
+    default: str | None,
+    choices: list[str] | Callable[[], list[str]],
+    case_sensitive: bool = True,
+) -> Callable[[], str | None]:
+    """
+    Create a lambda that validates environment variable against allowed choices
+
+    Args:
+        env_name: Name of the environment variable
+        default: Default value if not set (can be None)
+        choices: List of valid string options or callable that returns list
+        case_sensitive: Whether validation should be case sensitive
+
+    Returns:
+        Lambda function for environment_variables dict
+    """
+
+    def _get_validated_env() -> str | None:
+        value = os.getenv(env_name)
+        if value is None:
+            return default
+
+        # Resolve choices if it's a callable (for lazy loading)
+        actual_choices = choices() if callable(choices) else choices
+
+        if not case_sensitive:
+            check_value = value.lower()
+            check_choices = [choice.lower() for choice in actual_choices]
+        else:
+            check_value = value
+            check_choices = actual_choices
+
+        if check_value not in check_choices:
+            raise ValueError(f"Invalid value '{value}' for {env_name}. "
+                             f"Valid options: {actual_choices}.")
+
+        return value
+
+    return _get_validated_env
+
+
+def env_bool(env_name: str, default: bool = False) -> Callable[[], bool]:
+    """
+    Accepts both numeric strings ("0", "1") and boolean strings
+    ("true", "false", "True", "False").
+
+    Args:
+        env_name: Name of the environment variable
+        default: Default boolean value if not set
+    """
+
+    def _get_bool_env() -> bool:
+        value = os.getenv(env_name)
+        if value is None or value == "":
+            return default
+
+        value_lower = value.lower()
+        if value_lower in ("true", "1"):
+            return True
+        elif value_lower in ("false", "0"):
+            return False
+        else:
+            raise ValueError(
+                f"Invalid boolean value '{value}' for {env_name}. "
+                f"Valid options: '0', '1', 'true', 'false', 'True', 'False'.")
+
+    return _get_bool_env
+
+
+environment_variables: dict[str, Callable[[], Any]] = {
+    # JAX platform selection (e.g., "tpu", "cpu", "proxy")
+    "JAX_PLATFORMS":
+    lambda: os.getenv("JAX_PLATFORMS", "").lower(),
+    # TPU accelerator type (e.g., "v5litepod-16", "v4-8")
+    "TPU_ACCELERATOR_TYPE":
+    lambda: os.getenv("TPU_ACCELERATOR_TYPE", None),
+    # Name of the TPU resource
+    "TPU_NAME":
+    lambda: os.getenv("TPU_NAME", None),
+    # Worker ID for multi-host TPU setups
+    "TPU_WORKER_ID":
+    lambda: os.getenv("TPU_WORKER_ID", None),
+    # Backend for multi-host communication on TPU
+    "TPU_MULTIHOST_BACKEND":
+    env_with_choices("TPU_MULTIHOST_BACKEND", "", ["ray"]),
+    # Slice configuration for disaggregated prefill workers
+    "PREFILL_SLICES":
+    lambda: os.getenv("PREFILL_SLICES", ""),
+    # Slice configuration for disaggregated decode workers
+    "DECODE_SLICES":
+    lambda: os.getenv("DECODE_SLICES", ""),
+    # Skip JAX precompilation step during initialization
+    "SKIP_JAX_PRECOMPILE":
+    env_bool("SKIP_JAX_PRECOMPILE", default=False),
+    # Check for XLA recompilation during execution
+    "VLLM_XLA_CHECK_RECOMPILATION":
+    env_bool("VLLM_XLA_CHECK_RECOMPILATION", default=False),
+    # Model implementation type (e.g., "flax_nnx")
+    "MODEL_IMPL_TYPE":
+    env_with_choices("MODEL_IMPL_TYPE", "auto",
+                     ["auto", "vllm", "flax_nnx", "jetpack"]),
+    # Enable new experimental model design
+    "NEW_MODEL_DESIGN":
+    env_bool("NEW_MODEL_DESIGN", default=False),
+    # Directory to store phased profiling output
+    "PHASED_PROFILING_DIR":
+    lambda: os.getenv("PHASED_PROFILING_DIR", ""),
+    # Python tracer level for profiling
+    "PYTHON_TRACER_LEVEL":
+    lambda: int(os.getenv("PYTHON_TRACER_LEVEL") or "1"),
+    # Use custom expert-parallel kernel for MoE (Mixture of Experts)
+    "USE_MOE_EP_KERNEL":
+    env_bool("USE_MOE_EP_KERNEL", default=False),
+    # Number of TPU slices for multi-slice mesh
+    "NUM_SLICES":
+    lambda: int(os.getenv("NUM_SLICES") or "1"),
+    # Enable/disable Ray usage statistics collection
+    "RAY_USAGE_STATS_ENABLED":
+    lambda: os.getenv("RAY_USAGE_STATS_ENABLED", "0"),
+    # Ray compiled DAG channel type for TPU
+    "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE":
+    env_with_choices("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "shm", ["shm"]),
+    "ENABLE_QUANTIZED_MATMUL_KERNEL":
+    lambda: bool(int(os.getenv("ENABLE_QUANTIZED_MATMUL_KERNEL") or "0")),
+}
+
+
+def __getattr__(name: str) -> Any:
+    """
+    Gets environment variables lazily.
+
+    NOTE: After enable_envs_cache() invocation (which triggered after service
+    initialization), all environment variables will be cached.
+    """
+    if name in environment_variables:
+        return environment_variables[name]()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def enable_envs_cache() -> None:
+    """
+    Enables caching of environment variables by wrapping the module's __getattr__
+    function with functools.cache(). This improves performance by avoiding
+    repeated re-evaluation of environment variables.
+
+    NOTE: This should be called after service initialization. Once enabled,
+    environment variable values are cached and will not reflect changes to
+    os.environ until the process is restarted.
+    """
+    # Tag __getattr__ with functools.cache
+    global __getattr__
+    __getattr__ = functools.cache(__getattr__)
+
+    # Cache all environment variables
+    for key in environment_variables:
+        __getattr__(key)
+
+
+def __dir__() -> list[str]:
+    return list(environment_variables.keys())

tpu_inference/executors/__init__.py

@@ -0,0 +1,13 @@
+# 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.