tpu-inference 0.0.1rc1__py3-none-any.whl → 0.11.1.dev202511130813__py3-none-any.whl

tpu_inference/distributed/tpu_connector.py

@@ -60,6 +60,7 @@ D workflow:
 
 import copy
 import functools
+import os
 import threading
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -85,7 +86,6 @@ if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
     from vllm.v1.request import Request
 
-from tpu_inference import envs
 from tpu_inference.distributed.utils import (get_host_ip, get_kv_ips,
                                              get_kv_ports,
                                              get_kv_transfer_port, get_node_id,
@@ -441,7 +441,8 @@ class TPUConnectorWorker:
 
         self.runner: TPUModelRunner = None
         self.mesh: Mesh = None
-        self.multi_host = envs.TPU_MULTIHOST_BACKEND == "ray"
+        self.multi_host = os.getenv("TPU_MULTIHOST_BACKEND",
+                                    "").lower() == "ray"
         # NOTE(xiang): This can not be the worker rank set in RayDistributedExecutor.
         # The worker rank is assigned with vLLM's sorting logic, which does not work
         # for TPU host topology.
@@ -457,6 +458,7 @@ class TPUConnectorWorker:
         self.side_channel_port = get_side_channel_port()
 
         self.kv_transfer_server = None
+        self._maybe_start_p2p_server()
         self.zmq_cxt = zmq.Context()
         if self.is_producer:
             ready_event = threading.Event()
@@ -498,7 +500,6 @@ class TPUConnectorWorker:
         self.shape = list(kv_layer.shape)
         self.dtype = kv_layer.dtype
         self.sharding = kv_layer.sharding
-        self._maybe_start_p2p_server()
 
     def _maybe_start_p2p_server(self):
         if self.kv_transfer_server is not None:

tpu_inference/distributed/utils.py

@@ -2,7 +2,6 @@ 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__)
@@ -18,7 +17,7 @@ def set_node_kv_ip_port(ip_port: tuple[int, str, int]):
 
 
 def get_kv_ips() -> str:
-    if envs.TPU_MULTIHOST_BACKEND == "ray":
+    if os.getenv("TPU_MULTIHOST_BACKEND", "").lower() == "ray":
         num_nodes = len(_NODES_KV_IP_PORT)
         ips = []
         for node_id in range(num_nodes):
@@ -29,7 +28,7 @@ def get_kv_ips() -> str:
 
 
 def get_kv_ports() -> str:
-    if envs.TPU_MULTIHOST_BACKEND == "ray":
+    if os.getenv("TPU_MULTIHOST_BACKEND", "").lower() == "ray":
         num_nodes = len(_NODES_KV_IP_PORT)
         ports = []
         for node_id in range(num_nodes):

tpu_inference/envs.py

@@ -15,64 +15,18 @@ if TYPE_CHECKING:
     PREFILL_SLICES: str = ""
     DECODE_SLICES: str = ""
     SKIP_JAX_PRECOMPILE: bool = False
-    VLLM_XLA_CHECK_RECOMPILATION: bool = False
     MODEL_IMPL_TYPE: str = "flax_nnx"
     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"
 
-
-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
-
-
 environment_variables: dict[str, Callable[[], Any]] = {
     # JAX platform selection (e.g., "tpu", "cpu", "proxy")
     "JAX_PLATFORMS":
-    lambda: os.getenv("JAX_PLATFORMS", "").lower(),
+    lambda: os.getenv("JAX_PLATFORMS", ""),
     # TPU accelerator type (e.g., "v5litepod-16", "v4-8")
     "TPU_ACCELERATOR_TYPE":
     lambda: os.getenv("TPU_ACCELERATOR_TYPE", None),
@@ -84,7 +38,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     lambda: os.getenv("TPU_WORKER_ID", None),
     # Backend for multi-host communication on TPU
     "TPU_MULTIHOST_BACKEND":
-    env_with_choices("TPU_MULTIHOST_BACKEND", "", ["ray"]),
+    lambda: os.getenv("TPU_MULTIHOST_BACKEND", "").lower(),
     # Slice configuration for disaggregated prefill workers
     "PREFILL_SLICES":
     lambda: os.getenv("PREFILL_SLICES", ""),
@@ -93,35 +47,28 @@ environment_variables: dict[str, Callable[[], Any]] = {
     lambda: os.getenv("DECODE_SLICES", ""),
     # Skip JAX precompilation step during initialization
     "SKIP_JAX_PRECOMPILE":
-    lambda: bool(int(os.getenv("SKIP_JAX_PRECOMPILE") or "0")),
-    # Check for XLA recompilation during execution
-    "VLLM_XLA_CHECK_RECOMPILATION":
-    lambda: bool(int(os.getenv("VLLM_XLA_CHECK_RECOMPILATION") or "0")),
+    lambda: bool(int(os.getenv("SKIP_JAX_PRECOMPILE", "0"))),
     # Model implementation type (e.g., "flax_nnx")
     "MODEL_IMPL_TYPE":
-    env_with_choices("MODEL_IMPL_TYPE", "flax_nnx",
-                     ["vllm", "flax_nnx", "jetpack"]),
+    lambda: os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower(),
     # Enable new experimental model design
     "NEW_MODEL_DESIGN":
-    lambda: bool(int(os.getenv("NEW_MODEL_DESIGN") or "0")),
+    lambda: bool(int(os.getenv("NEW_MODEL_DESIGN", "0"))),
     # 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"),
+    lambda: int(os.getenv("PYTHON_TRACER_LEVEL", "1")),
     # Use custom expert-parallel kernel for MoE (Mixture of Experts)
     "USE_MOE_EP_KERNEL":
-    lambda: bool(int(os.getenv("USE_MOE_EP_KERNEL") or "0")),
-    # Number of TPU slices for multi-slice mesh
-    "NUM_SLICES":
-    lambda: int(os.getenv("NUM_SLICES") or "1"),
+    lambda: bool(int(os.getenv("USE_MOE_EP_KERNEL", "0"))),
     # 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"]),
+    lambda: os.getenv("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "shm"),
 }
 
 

tpu_inference/executors/ray_distributed_executor.py

@@ -108,9 +108,6 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
             ip_port = self.collective_rpc("get_node_kv_ip_port")
             for item in ip_port:
                 set_node_kv_ip_port(item)
-        self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
-            self.vllm_config.ec_transfer_config is None
-            or not self.vllm_config.ec_transfer_config.is_ec_producer)
 
     def _initialize_ray_cluster(self) -> None:
         """Initialize the distributed cluster with Ray.
@@ -134,21 +131,10 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
                 f"current platform {current_platform.device_name} does not "
                 "support ray.")
 
-        pp_size = self.parallel_config.pipeline_parallel_size
-        placement_group_specs: List[Dict[str, float]] = []
-
-        ray_nodes = ray.nodes()
-        logger.info(f"RayDistributedExecutor | ray_nodes={ray_nodes}")
-
-        if pp_size == 1:
-            placement_group_specs = [{
-                device_str: node['Resources'][device_str]
-            } for node in ray_nodes]
-        else:
-            num_devices_per_pp_rank = self.vllm_config.sharding_config.total_devices
-            placement_group_specs = [{
-                device_str: num_devices_per_pp_rank
-            } for _ in range(pp_size)]
+        placement_group_specs: List[Dict[str, float]] = [{
+            device_str:
+            node['Resources'][device_str]
+        } for node in ray.nodes()]
 
         # vLLM engine is also a worker to execute model with an accelerator,
         # so it requires to have the device in a current node. Check if
@@ -343,8 +329,6 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
         all_kwargs = []
         for rank, (node_id, _) in enumerate(worker_node_and_tpu_ids):
             local_rank = node_workers[node_id].index(rank)
-            ip = sorted_worker_metadata[rank].ip
-            prev_ip = sorted_worker_metadata[rank - 1].ip if rank > 0 else ""
             kwargs = dict(
                 vllm_config=self.vllm_config,
                 local_rank=local_rank,
@@ -352,26 +336,22 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
                 distributed_init_method=distributed_init_method,
                 is_driver_worker=(not self.parallel_config)
                 or (rank % self.parallel_config.tensor_parallel_size == 0),
-                ip=ip,
-                prev_worker_ip=prev_ip,
             )
             all_kwargs.append(kwargs)
         self.collective_rpc("init_worker", args=(all_kwargs, ))
         self.collective_rpc("init_device")
-        if self.parallel_config.pipeline_parallel_size > 1:
-            self.collective_rpc("initialize_pp_transfer_connect")
         self.collective_rpc("load_model")
 
         if self.use_ray_spmd_worker:
             for pp_rank in range(self.parallel_config.pipeline_parallel_size):
                 self.pp_tp_workers.append([])
-                num_tp_workers = int(
-                    self.parallel_config.tensor_parallel_size //
-                    num_tpu_per_worker)
-                for tp_rank in range(num_tp_workers):
-                    # PP=2, TP=4, num_tpu_per_worker=2
-                    # pp_tp_workers = [[0, 1], [2, 3]]
-                    rank = (pp_rank * num_tp_workers) + tp_rank
+                for tp_rank in range(
+                        int(self.parallel_config.tensor_parallel_size //
+                            num_tpu_per_worker)):
+                    # PP=2, TP=4
+                    # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
+                    rank = (pp_rank * self.parallel_config.tensor_parallel_size
+                            ) + tp_rank
                     assert len(self.pp_tp_workers[pp_rank]) == tp_rank
                     assert pp_rank < len(self.pp_tp_workers)
                     self.pp_tp_workers[pp_rank].append(self.workers[rank])