tpu-inference 0.11.1__py3-none-any.whl

tpu_inference/core/disagg_executor.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+from concurrent.futures import Future
+from multiprocessing import Lock
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from vllm.logger import init_logger
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.cache import worker_receiver_cache_from_config
+from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
+                        run_method)
+from vllm.v1.executor.abstract import Executor
+from vllm.v1.outputs import AsyncModelRunnerOutput
+from vllm.v1.worker.worker_base import WorkerWrapperBase
+
+logger = init_logger(__name__)
+
+
+class DisaggExecutor(Executor):
+
+    def _init_executor(self) -> None:
+        """Initialize the worker and load the model.
+        """
+        self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
+                                               rpc_rank=0)
+        slice_config = getattr(self.vllm_config.device_config, "slice")
+        idx = slice_config[0]
+        jax_devices = slice_config[-1]
+        devices = []
+        if isinstance(idx, int):
+            sizes = slice_config[1]
+            start = sum(sizes[0:idx])
+            end = start + sizes[idx]
+
+            devices = jax_devices[start:end]
+            setattr(self.vllm_config.device_config, "slice",
+                    (idx + 1, sizes, jax_devices))
+            logger.debug(
+                f"Creating DisaggExecutor with {devices}, index: {start} -> {end}"
+            )
+        elif isinstance(idx, tuple):
+            slice_idx = slice_config[1]
+            sizes = slice_config[2][slice_idx]
+            start_row, start_col = idx
+            selected_devices = []
+            max_row, max_col = 0, 0
+            for device in jax_devices:
+                coords = device.coords
+                max_row = max(max_row, coords[0])
+                max_col = max(max_col, coords[1])
+                if coords[0] >= start_row and coords[0] < start_row + sizes[0]:
+                    if coords[1] >= start_col and coords[
+                            1] < start_col + sizes[1]:
+                        selected_devices.append(device)
+            max_row, max_col = max_row + 1, max_col + 1
+
+            devices = selected_devices
+            if start_col + sizes[1] >= max_col:
+                start_row += sizes[0]
+                start_col = 0
+            else:
+                start_col += sizes[1]
+
+            setattr(self.vllm_config.device_config, "slice",
+                    ((start_row, start_col), slice_idx + 1, slice_config[2],
+                     jax_devices))
+            logger.debug(
+                f"Creating DisaggExecutor with {devices}, next start: {((start_row, start_col), slice_idx+1, slice_config[2])}"
+            )
+
+        distributed_init_method = get_distributed_init_method(
+            get_ip(), get_open_port())
+        local_rank = 0
+        rank = 0
+        is_driver_worker = True
+        kwargs = dict(
+            vllm_config=self.vllm_config,
+            local_rank=local_rank,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            is_driver_worker=is_driver_worker,
+            devices=devices,
+        )
+        self.mm_receiver_cache = worker_receiver_cache_from_config(
+            self.vllm_config, MULTIMODAL_REGISTRY, Lock())
+        self.collective_rpc("init_worker", args=([kwargs], ))
+        self.collective_rpc("init_device")
+        self.collective_rpc("load_model")
+
+    def collective_rpc(self,
+                       method: Union[str, Callable],
+                       timeout: Optional[float] = None,
+                       args: Tuple = (),
+                       kwargs: Optional[Dict] = None,
+                       non_block: bool = False) -> List[Any]:
+        if kwargs is None:
+            kwargs = {}
+
+        if not non_block:
+            return [run_method(self.driver_worker, method, args, kwargs)]
+
+        try:
+            result = run_method(self.driver_worker, method, args, kwargs)
+            if isinstance(result, AsyncModelRunnerOutput):
+                if (async_thread := self.async_output_thread) is not None:
+                    return [async_thread.submit(result.get_output)]
+                result = result.get_output()
+            future = Future[Any]()
+            future.set_result(result)
+        except Exception as e:
+            future = Future[Any]()
+            future.set_exception(e)
+        return [future]
+
+    def check_health(self) -> None:
+        # DisaggExecutor will always be healthy as long as
+        # it's running.
+        return

tpu_inference/core/disagg_utils.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Tuple
+
+PREFILL_SLICES = 'PREFILL_SLICES'
+DECODE_SLICES = 'DECODE_SLICES'
+
+
+def is_disagg_enabled() -> bool:
+    # We triggrer our code path as long as prefill slices are set. This
+    # allows us to test interleave mode effectively with the code path
+    # for comparison purposes.
+    return PREFILL_SLICES in os.environ
+
+
+def _parse_slices(slices_str: str) -> Tuple[int, ...]:
+    """Parse slices environment variable and return the a list of integers, each the size of a slice.
+
+    For example, if slices_str is set to `2x2,2x1,2x4`, we should return `(4, 2, 8)`.
+
+    Throws exception if the slice str is malformed.
+    """
+    if not slices_str:
+        return ()
+
+    try:
+        slice_sizes = []
+        for s in slices_str.split(','):
+            dims = s.split('x')
+            if len(dims) == 1:
+                slice_sizes.append(int(dims[0]))
+            elif len(dims) == 2:
+                slice_sizes.append((int(dims[0]), int(dims[1])))
+            else:
+                raise ValueError("Each slice must be in 'N' or 'NxM' format.")
+        return tuple(slice_sizes)
+    except ValueError as e:
+        raise ValueError(f"Malformed slice string: '{slices_str}'") from e
+
+
+def get_prefill_slices() -> Tuple[int, ...]:
+    if PREFILL_SLICES not in os.environ:
+        return ()
+    return _parse_slices(os.environ[PREFILL_SLICES])
+
+
+def get_decode_slices() -> Tuple[int, ...]:
+    if DECODE_SLICES not in os.environ:
+        return ()
+    return _parse_slices(os.environ[DECODE_SLICES])

tpu_inference/di/abstracts.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC
+
+
+class AbstractModelRunnerOutput(ABC):
+    """Abstract base class for model runner output."""
+    pass
+
+
+class AbstractSchedulerOutput(ABC):
+    """Abstract base class for scheduler output."""
+    pass
+
+
+class AbstractLoRARequest(ABC):
+    """Abstract base class for LoRA request."""
+    pass
+
+
+class AbstractKVCacheConfig(ABC):
+    """Abstract base class for KV cache config."""
+    pass
+
+
+class AbstractKVCacheSpec(ABC):
+    """Abstract base class for KV cache spec."""
+    pass

tpu_inference/di/host.py

@@ -0,0 +1,76 @@
+"""
+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.
+"""
+from typing import Any, Callable, Dict, Type
+
+
+class DIHost:
+    """
+    A simple dependency injection host.
+
+    This host manages a graph of functions, where each function is a provider
+    for a specific data type and declares its own dependencies.
+    """
+
+    def __init__(self):
+        self._providers: Dict[Type, Callable[..., Any]] = {}
+        self._dependencies: Dict[Callable[..., Any], Dict[str, Type]] = {}
+
+    def register(self,
+                 provider: Callable[..., Any],
+                 output_type: Type,
+                 dependencies: Dict[str, Type] = None):
+        """
+        Registers a provider function with the host.
+
+        Args:
+            provider: The function that produces the output.
+            output_type: The data type that the function produces.
+            dependencies: A dictionary mapping argument names of the provider
+                          to the data types they require.
+        """
+        self._providers[output_type] = provider
+        if dependencies:
+            self._dependencies[provider] = dependencies
+
+    def resolve(self, target_type: Type) -> Any:
+        """
+        Resolves a dependency by creating an instance of the target type.
+
+        This method will recursively resolve all dependencies required to call
+        the provider for the target type.
+
+        Args:
+            target_type: The data type to be resolved.
+
+        Returns:
+            An instance of the target type.
+        """
+        if target_type not in self._providers:
+            raise ValueError(
+                f"No provider registered for type {target_type.__name__}")
+
+        provider = self._providers[target_type]
+
+        if provider not in self._dependencies:
+            # Provider has no dependencies, so just call it.
+            return provider()
+
+        # Resolve dependencies for the provider.
+        kwargs = {}
+        for arg_name, dep_type in self._dependencies[provider].items():
+            kwargs[arg_name] = self.resolve(dep_type)
+
+        return provider(**kwargs)

tpu_inference/di/interfaces.py

@@ -0,0 +1,51 @@
+"""
+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 abc
+
+
+class HostInterface(abc.ABC):
+    """
+    An interface that the host system (e.g., SGLang, vLLM) must implement.
+    This defines the contract for how the backend can call back into the host.
+    """
+
+    @abc.abstractmethod
+    def get_next_batch_to_run(self):
+        """
+        The backend calls this to get the next batch of requests to process.
+        """
+        pass
+
+    @abc.abstractmethod
+    def process_batch_result(self, batch_result):
+        """
+        The backend calls this to return the results of a processed batch.
+        """
+        pass
+
+
+class BackendInterface(abc.ABC):
+    """
+    An interface that the backend system (e.g., tpu_inference) must implement.
+    This defines the contract for how the host can call into the backend.
+    """
+
+    @abc.abstractmethod
+    def launch_tpu_batch(self, batch_to_launch):
+        """
+        The host calls this to launch a batch of requests on the backend.
+        """
+        pass