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

tpu_inference/core/disagg_executor.py

@@ -0,0 +1,118 @@
+# 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.network_utils import (get_distributed_init_method, get_ip,
+                                      get_open_port)
+from vllm.v1.executor.abstract import Executor
+from vllm.v1.outputs import AsyncModelRunnerOutput
+from vllm.v1.serial_utils import run_method
+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,49 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Tuple
+
+from tpu_inference import envs
+
+
+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 bool(envs.PREFILL_SLICES)
+
+
+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 not envs.PREFILL_SLICES:
+        return ()
+    return _parse_slices(envs.PREFILL_SLICES)
+
+
+def get_decode_slices() -> Tuple[int, ...]:
+    if not envs.DECODE_SLICES:
+        return ()
+    return _parse_slices(envs.DECODE_SLICES)

tpu_inference/core/sched/__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.