xinference 1.0.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

xinference/model/llm/vllm/xavier/block.py

@@ -0,0 +1,111 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# 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 asyncio
+import logging
+from typing import Any, Dict, Optional
+
+import xoscar as xo
+from vllm.core.block.interfaces import BlockId
+from vllm.core.block.prefix_caching_block import (
+    BlockTracker,
+    PrefixCachingBlockAllocator,
+)
+
+from .....isolation import Isolation
+
+logger = logging.getLogger(__name__)
+
+
+class XavierInnerBlockTracker(BlockTracker):
+    """Used to track the status of a block inside the prefix caching allocator"""
+
+    """
+    Here, two fixed attributes, `transferred` and `executed`,
+    have been added to the `BlockTracker` class to mark the status of the corresponding `block_id`.
+    We cannot directly set attributes on the `Block` object
+    because the `Block` objects are dynamically allocated with each scheduling.
+    The `Block` objects executed in two different scheduling steps may have the same `id`, `hash`, etc.,
+    but the instance objects may differ.
+    The BlockTracker object inside vllm is one-to-one with the block_id.
+    """
+    __slots__ = ("active", "last_accessed", "computed", "transferred", "executed")
+
+    def __init__(self):
+        super().__init__()
+        self.transferred = False
+        self.executed = False
+
+
+class XavierPrefixCachingBlockAllocator(PrefixCachingBlockAllocator):
+    def __init__(self, *args, run_isolation: bool = False, **kwargs):
+        super().__init__(*args, **kwargs)
+        for _id in self._block_tracker.keys():
+            self._block_tracker[_id] = XavierInnerBlockTracker()
+
+        self._xavier_config: Optional[Dict[str, Any]] = None
+        self._block_tracker_ref = None
+        if run_isolation:
+            self._isolation = Isolation(
+                asyncio.new_event_loop(), threaded=True, daemon=True
+            )
+            self._isolation.start()
+        else:
+            self._isolation = None  # type: ignore
+
+    def __del__(self):
+        if self._isolation is not None:
+            self._isolation.stop()
+
+    @property
+    def xavier_config(self):
+        return self._xavier_config
+
+    @xavier_config.setter
+    def xavier_config(self, v: Dict[str, Any]):
+        self._xavier_config = v
+
+    async def _get_block_tracker_ref(self):
+        if self._block_tracker_ref is None:
+            block_tracker_address = self.xavier_config.get("block_tracker_address")
+            block_tracker_uid = self.xavier_config.get("block_tracker_uid")
+            self._block_tracker_ref = await xo.actor_ref(
+                address=block_tracker_address, uid=block_tracker_uid
+            )
+        return self._block_tracker_ref
+
+    async def unregister_block(self, block_id: int):
+        assert self._xavier_config is not None
+        tracker_ref = await self._get_block_tracker_ref()
+        await tracker_ref.unregister_block(
+            self.xavier_config.get("virtual_engine"),
+            self.xavier_config.get("rank"),
+            block_id,
+        )
+
+    def _maybe_allocate_evicted_block_id(self) -> Optional[BlockId]:
+        """
+        This is the only entry point where the `block_id` is evicted from the cache.
+        Therefore, when the `block_id` is evicted, the tracker actor needs to unregister the block information.
+        At the same time, make sure to reset the attributes corresponding to that `block_id`.
+        """
+        evicted_block_id = super()._maybe_allocate_evicted_block_id()
+        logger.debug(f"block_id: {evicted_block_id} will be evicted from the cache.")
+        if evicted_block_id is not None and self._isolation is not None:
+            tracker = self._block_tracker[evicted_block_id]
+            assert isinstance(tracker, XavierInnerBlockTracker)
+            tracker.transferred = False
+            tracker.executed = False
+            self._isolation.call(self.unregister_block(evicted_block_id))
+            logger.debug(f"block_id: {evicted_block_id} will be used again.")
+        return evicted_block_id

xinference/model/llm/vllm/xavier/block_manager.py

@@ -0,0 +1,71 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# 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 logging
+from typing import Any, Dict, Optional
+
+from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
+from vllm.core.block.interfaces import Block
+from vllm.core.block_manager import SelfAttnBlockSpaceManager
+from vllm.sequence import SequenceGroup, SequenceStatus
+from vllm.utils import Device
+
+from .allocator import XavierCpuGpuBlockAllocator
+
+logger = logging.getLogger(__name__)
+
+
+class XavierBlockManager(SelfAttnBlockSpaceManager):
+    def __init__(self, *args, **kwargs):
+        # Monkey patch
+        CpuGpuBlockAllocator.create = XavierCpuGpuBlockAllocator.create
+        super().__init__(*args, **kwargs)
+        self._xavier_config: Optional[Dict[str, Any]] = None
+        logger.debug("Init xavier block manager done.")
+
+    @property
+    def xavier_config(self):
+        return self._xavier_config
+
+    @xavier_config.setter
+    def xavier_config(self, value: Dict[str, Any]):
+        self._xavier_config = value
+        self.block_allocator.xavier_config = value
+
+    def get_block_by_block_id(self, seq_id: int, block_id: int) -> Block:
+        table = self.block_tables[seq_id]
+        for b in table.blocks:
+            if b.block_id == block_id:
+                return b
+
+    def get_block_status_by_block_id(self, status_name: str, block_id: int) -> bool:
+        tracker = self.block_allocator._allocators[Device.GPU]._block_tracker[block_id]
+        return getattr(tracker, status_name)
+
+    def set_block_status_by_block_id(
+        self, status_name: str, block_id: int, status: bool
+    ) -> None:
+        tracker = self.block_allocator._allocators[Device.GPU]._block_tracker[block_id]
+        assert getattr(tracker, status_name, None) is not None
+        setattr(tracker, status_name, status)
+
+    def allocate(self, seq_group: SequenceGroup) -> None:
+        """
+        If the `seq_group` has the `transferred` attribute,
+        it indicates that the `seq_group` has gone through the transfer process,
+        so the block allocation logic should not be executed again.
+        """
+        waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)
+        if all([getattr(s, "transferred", False) for s in waiting_seqs]):
+            return
+        super().allocate(seq_group)

xinference/model/llm/vllm/xavier/block_tracker.py

@@ -0,0 +1,129 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# 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 random
+from typing import Dict, List, Optional, Set, Tuple
+
+import xoscar as xo
+
+
+class VLLMBlockTracker(xo.StatelessActor):
+    @classmethod
+    def default_uid(cls):
+        return f"vllm-block-tracker-actor"
+
+    def __init__(self):
+        super().__init__()
+        # engine -> hash -> (rank, block_id)
+        self._hash_to_rank_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {}
+        # engine -> rank -> (hash, block_id)
+        self._rank_to_hash_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {}
+        self._unavailable_ranks: Set[int] = set()
+
+    def register_blocks(
+        self, virtual_engine: int, block_infos: List[Tuple[int, int]], rank: int
+    ):
+        # Update query meta
+        if virtual_engine not in self._hash_to_rank_and_block_id:
+            self._hash_to_rank_and_block_id[virtual_engine] = {}
+        hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine]
+        for hash_content, block_id in block_infos:
+            if hash_content not in hash_to_rank_and_block_id:
+                hash_to_rank_and_block_id[hash_content] = {
+                    (rank, block_id),
+                }
+            else:
+                hash_to_rank_and_block_id[hash_content].add((rank, block_id))
+
+        # Update remove meta
+        if virtual_engine not in self._rank_to_hash_and_block_id:
+            self._rank_to_hash_and_block_id[virtual_engine] = {}
+        rank_to_hash_and_block_id = self._rank_to_hash_and_block_id[virtual_engine]
+        if rank not in rank_to_hash_and_block_id:
+            rank_to_hash_and_block_id[rank] = set()
+        rank_to_hash_and_block_id[rank].update(block_infos)
+
+    def query_blocks(
+        self, virtual_engine: int, hash_contents: List[Tuple[int, int]]
+    ) -> Dict[int, Set[Tuple[int, int, int]]]:
+        if virtual_engine not in self._hash_to_rank_and_block_id:
+            return {}
+        hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine]
+        remote: Dict[int, Set[Tuple[int, int, int]]] = {}
+        for hash_content, _id in hash_contents:
+            if (
+                hash_content in hash_to_rank_and_block_id
+            ) and hash_to_rank_and_block_id[hash_content]:
+                # exclude ranks that are in the recovery process
+                rank_and_block_id = [
+                    (r, b)
+                    for r, b in hash_to_rank_and_block_id[hash_content]
+                    if r not in self._unavailable_ranks
+                ]
+                if rank_and_block_id:
+                    # TODO: Randomly select here, and try to distribute requests as evenly as possible.
+                    # There may be better methods in the future.
+                    rank, block_id = random.choice(rank_and_block_id)
+                    if rank not in remote:
+                        remote[rank] = {
+                            (hash_content, block_id, _id),
+                        }
+                    else:
+                        remote[rank].add((hash_content, block_id, _id))
+        return remote
+
+    def unregister_block(self, virtual_engine: int, rank: int, block_id: int):
+        if (virtual_engine not in self._rank_to_hash_and_block_id) or (
+            virtual_engine not in self._hash_to_rank_and_block_id
+        ):
+            return
+
+        # Update remove meta
+        rank_to_hash_and_block_id = self._rank_to_hash_and_block_id[virtual_engine]
+        if rank not in rank_to_hash_and_block_id:
+            return
+        hash_and_block_id = rank_to_hash_and_block_id[rank]
+        detail: Optional[Tuple[int, int]] = None
+        for hash_content, _id in hash_and_block_id.copy():
+            if _id == block_id:
+                detail = (hash_content, block_id)
+                hash_and_block_id.discard(detail)
+                break
+
+        # Update query meta
+        if detail is not None:
+            hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine]
+            _hash = detail[0]
+            if _hash in hash_to_rank_and_block_id:
+                hash_to_rank_and_block_id[_hash].discard((rank, detail[1]))
+
+    def unregister_rank(self, rank: int):
+        """
+        This rank is in the recovery process, and its query results will be excluded.
+        """
+        self._unavailable_ranks.add(rank)
+
+    def register_rank(self, rank: int):
+        """
+        After recovery is successful, clear all stale data of the rank and mark the rank as available.
+        """
+        for _, rank_to_hash_and_block_id in self._rank_to_hash_and_block_id.items():
+            rank_to_hash_and_block_id.pop(rank, None)
+
+        for _, hash_to_rank_and_block_id in self._hash_to_rank_and_block_id.items():
+            for _, rank_and_block_id in hash_to_rank_and_block_id.items():
+                to_delete = [(r, b) for r, b in rank_and_block_id if r == rank]
+                if to_delete:
+                    rank_and_block_id.difference_update(to_delete)
+
+        self._unavailable_ranks.discard(rank)

xinference/model/llm/vllm/xavier/collective.py

@@ -0,0 +1,74 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# 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 logging
+from typing import List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class CollectiveRank:
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        rank_address: str,
+        store_address: str,
+        store_port: int,
+        world_addresses: List[str],
+    ):
+        self._rank = rank
+        self._world_size = world_size
+        self._rank_address = rank_address
+        self._world_addresses = world_addresses
+        self._store_address = store_address
+        self._store_port = store_port
+        self._device = None
+        self._tcp_store = None
+        self._context = None
+
+    def init_rank(self):
+        from xoscar.collective import xoscar_pygloo as xp
+
+        self._context = xp.rendezvous.Context(self._rank, self._world_size)
+
+        attr = xp.transport.tcp.attr(self._rank_address.split(":")[0])
+        self._device = xp.transport.tcp.CreateDevice(attr)
+
+        opt = xp.rendezvous.TCPStoreOptions()
+        opt.port = self._store_port
+        opt.numWorkers = self._world_size
+        opt.isServer = self._rank == 0
+        opt.waitWorkers = False
+
+        self._tcp_store = xp.rendezvous.TCPStore(self._store_address, opt)
+        if self._world_addresses:
+            self.connect_full_mesh()
+
+    def connect_full_mesh(
+        self, prefix: Optional[str] = None, world_addresses: Optional[List[str]] = None
+    ):
+        from xoscar.collective import xoscar_pygloo as xp
+
+        assert self._device is not None
+        assert self._tcp_store is not None
+        assert self._context is not None
+        if world_addresses is not None:
+            self._world_addresses = world_addresses
+        prefix_store = xp.rendezvous.PrefixStore(
+            prefix or str(self._world_size), self._tcp_store
+        )
+        self._context.connectFullMesh(prefix_store, self._device)
+        logger.debug(
+            f"Rank {self._rank} arrives successfully, world addresses: {self._world_addresses}"
+        )

xinference/model/llm/vllm/xavier/collective_manager.py

@@ -0,0 +1,147 @@
+# Copyright 2022-2025 XProbe Inc.
+#
+# 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 asyncio
+import logging
+import traceback
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, no_type_check
+
+import xoscar as xo
+
+from .block_tracker import VLLMBlockTracker
+
+if TYPE_CHECKING:
+    from .transfer import Rank0TransferActor, TransferActor
+
+
+logger = logging.getLogger(__name__)
+
+
+class Rank0ModelActor(xo.StatelessActor):
+    @classmethod
+    def default_uid(cls):
+        return "rank0-model-actor"
+
+    def __init__(self, xavier_config: Dict[str, Any]):
+        super().__init__()
+        self._rank = 0
+        self._xavier_config = xavier_config
+        self._transfer_ref: Optional[xo.ActorRefType["Rank0TransferActor"]] = None
+
+    async def __pre_destroy__(self):
+        if self._transfer_ref is not None:
+            try:
+                await xo.destroy_actor(self._transfer_ref)
+                del self._transfer_ref
+            except Exception as e:
+                logger.debug(
+                    f"Destroy transfer actor failed, rank: {self._rank}, address: {self.address}, error: {e}"
+                )
+
+    @no_type_check
+    async def start_transfer_for_vllm(self, rank_addresses: List[str]):
+        from .transfer import Rank0TransferActor
+
+        self._transfer_ref = await xo.create_actor(
+            Rank0TransferActor,
+            address=self.address,
+            uid=f"{Rank0TransferActor.default_uid()}-{self._rank}",
+            rank=self._rank,
+            world_size=self._xavier_config.get("world_size"),  # type: ignore
+            rank_address=self._xavier_config.get("rank_address"),  # type: ignore
+            store_address=self._xavier_config.get("store_address"),  # type: ignore
+            store_port=self._xavier_config.get("store_port"),  # type: ignore
+            world_addresses=rank_addresses,
+        )
+        logger.debug(
+            f"Init transfer actor: {self._transfer_ref.address}, rank: {self._rank} done for vllm."  # type: ignore
+        )
+
+
+def with_lock(method):
+    async def wrapper(self, *args, **kwargs):
+        async with self._lock:
+            return await method(self, *args, **kwargs)
+
+    return wrapper
+
+
+class CollectiveManager(xo.StatelessActor):
+    @classmethod
+    def default_uid(cls):
+        return f"xavier-collective-manager"
+
+    def __init__(self, model_uid: str):
+        super().__init__()
+        self._model_uid = model_uid
+        self._tracker_ref: Optional[xo.ActorRefType["VLLMBlockTracker"]] = None
+        self._rank_to_ref: Dict[int, xo.ActorRefType["TransferActor"]] = {}
+        self._lock = asyncio.Lock()
+
+    async def __post_create__(self):
+        self._tracker_ref = await xo.actor_ref(
+            address=self.address,
+            uid=f"{VLLMBlockTracker.default_uid()}-{self._model_uid}",
+        )
+
+    async def unregister_rank(self, rank: int):
+        self._rank_to_ref.pop(rank, None)
+        await self._tracker_ref.unregister_rank(rank)  # type: ignore
+        logger.debug(f"Unregister rank: {rank}")
+
+    async def register_rank(self, rank: int, address: str, update: bool = False):
+        from .transfer import TransferActor
+
+        rank_ref = await xo.actor_ref(
+            address=address, uid=f"{TransferActor.default_uid()}-{rank}"
+        )
+        self._rank_to_ref[rank] = rank_ref
+        logger.debug(f"Register rank: {rank}, address: {address}")
+        if update:
+            await self._update_world()
+            await self._tracker_ref.register_rank(rank)  # type: ignore
+
+    @with_lock
+    async def _update_world(self):
+        """
+        Locking is used to prevent chaos when multiple replicas trigger recovery simultaneously.
+        """
+        from .....core.utils import gen_random_string
+
+        prefix = gen_random_string(6)
+        tasks = []
+        rank_to_ref = self._rank_to_ref.copy()
+        world_addresses = [ref.address for _, ref in sorted(rank_to_ref.items())]
+        for rank, ref in rank_to_ref.items():
+            tasks.append(ref.connect_full_mesh(prefix, world_addresses))
+        try:
+            logger.debug(
+                f"Rebuild collective communication with world_addresses: {world_addresses}, prefix: {prefix}"
+            )
+            await asyncio.gather(*tasks)
+            logger.debug(
+                f"Rebuild collective communication with world_addresses: {world_addresses}, prefix: {prefix} done."
+            )
+        except Exception as e:
+            """
+            The exception here is most likely due to another replica triggering recovery during the recovery process,
+            causing `connect_full_mesh` to time out.
+            Simply log the exception and
+            let the subsequent update process handle the reconstruction of the collective communication world.
+            """
+            logger.error(
+                f"Rebuild collective communication with world_addresses: {world_addresses} failed. "
+                f"Exception: {e}"
+            )
+            # Print the complete error stack
+            traceback.print_exception(type(e), e, e.__traceback__)