xinference 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

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__)

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

@@ -64,14 +64,13 @@ class XavierExecutor(GPUExecutorAsync):
         )
 
     async def _get_block_tracker_ref(self):
-        from .block_tracker import VLLMBlockTracker
-
         if self._block_tracker_ref is None:
             block_tracker_address = self.vllm_config.xavier_config.get(
                 "block_tracker_address"
             )
+            block_tracker_uid = self.vllm_config.xavier_config.get("block_tracker_uid")
             self._block_tracker_ref = await xo.actor_ref(
-                address=block_tracker_address, uid=VLLMBlockTracker.default_uid()
+                address=block_tracker_address, uid=block_tracker_uid
             )
         return self._block_tracker_ref
 
@@ -86,8 +85,8 @@ class XavierExecutor(GPUExecutorAsync):
             )
         return self._transfer_ref
 
-    def get_rank_address(self) -> str:
-        return self.vllm_config.xavier_config.get("rank_address")
+    def get_rank(self) -> int:
+        return self.vllm_config.xavier_config.get("rank")
 
     async def execute_model_async(
         self,
@@ -100,7 +99,7 @@ class XavierExecutor(GPUExecutorAsync):
         virtual_engine = execute_model_req.virtual_engine
         block_tracker_ref = await self._get_block_tracker_ref()
         scheduler = self.scheduler[virtual_engine]  # type: ignore
-        rank_address = self.get_rank_address()
+        rank = self.get_rank()
         executed_blocks_details: Set[Tuple[int, int]] = set()
         for meta in execute_model_req.seq_group_metadata_list:
             block_tables = meta.block_tables
@@ -117,16 +116,19 @@ class XavierExecutor(GPUExecutorAsync):
 
         res = await super().execute_model_async(execute_model_req)
 
-        """
-        Why not collect and register the information after execution?
-        Because after execution, the model's execution callback hook will release the block_id,
-        causing the block manager to lose access to the correct information.
-        """
-        await block_tracker_ref.register_blocks(
-            virtual_engine, list(executed_blocks_details), rank_address
-        )
+        if executed_blocks_details:
+            """
+            Why not collect and register the information after execution?
+            Because after execution, the model's execution callback hook will release the block_id,
+            causing the block manager to lose access to the correct information.
+            """
+            await block_tracker_ref.register_blocks(
+                virtual_engine, list(executed_blocks_details), rank
+            )
 
-        for _, _id in executed_blocks_details:
-            scheduler.block_manager.set_block_status_by_block_id("executed", _id, True)
+            for _, _id in executed_blocks_details:
+                scheduler.block_manager.set_block_status_by_block_id(
+                    "executed", _id, True
+                )
 
         return res

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

@@ -72,12 +72,11 @@ class XavierScheduler(Scheduler):
         self._transfer_status: Dict[SequenceGroup, Set[int]] = {}
 
     async def _get_block_tracker_ref(self):
-        from .block_tracker import VLLMBlockTracker
-
         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=VLLMBlockTracker.default_uid()
+                address=block_tracker_address, uid=block_tracker_uid
             )
         return self._block_tracker_ref
 
@@ -97,7 +96,12 @@ class XavierScheduler(Scheduler):
         virtual_engine: int,
         block_tables: Dict[int, List[int]],
         seq_group: SequenceGroup,
-    ) -> Tuple[Set[int], Dict[str, Set[Tuple[int, int, int]]]]:
+    ) -> Tuple[Set[int], Dict[int, Set[Tuple[int, int, int]]]]:
+        # If the `seq_group` has the `force_calculation` attribute set to `True`,
+        # it indicates that there were issues during the transmission process.
+        # In this case, force the computation and exclude it from the Xavier process.
+        if getattr(seq_group, "force_calculation", False):
+            return set(), dict()
         """
         Retrieve information from other replicas to check if any blocks have already been computed,
         for the purpose of data transfer.
@@ -132,48 +136,63 @@ class XavierScheduler(Scheduler):
                     )
                 ):
                     details.add(detail)
-        tracker_ref = await self._get_block_tracker_ref()
-        remote = await tracker_ref.query_blocks(virtual_engine, list(details))
-        # Not all queried blocks have corresponding results in other replicas.
-        # Therefore, it is necessary to record which local block data was actually transferred.
-        local: Set[int] = set()
-        for _, remote_details in remote.items():
-            for _, _, local_block_id in remote_details:
-                local.add(local_block_id)
-        if local:
-            logger.debug(
-                f"Data in local blocks: {local} will be transmitted from the remote."
-            )
-        return local, remote
+
+        if details:
+            tracker_ref = await self._get_block_tracker_ref()
+            remote = await tracker_ref.query_blocks(virtual_engine, list(details))
+            # Not all queried blocks have corresponding results in other replicas.
+            # Therefore, it is necessary to record which local block data was actually transferred.
+            local: Set[int] = set()
+            for _, remote_details in remote.items():
+                for _, _, local_block_id in remote_details:
+                    local.add(local_block_id)
+            if local:
+                logger.debug(
+                    f"Data in local blocks: {local} will be transmitted from the remote."
+                )
+            return local, remote
+        else:
+            return set(), dict()
 
     async def _do_transfer_inner(
-        self, virtual_engine: int, remote: Dict[str, Set[Tuple[int, int, int]]]
+        self, virtual_engine: int, remote: Dict[int, Set[Tuple[int, int, int]]]
     ):
         transfer_ref = await self._get_transfer_ref()
-        for addr, hash_and_block_id in remote.items():
+        for from_rank, hash_and_block_id in remote.items():
             src_to_dst: Dict[int, int] = {x[1]: x[2] for x in hash_and_block_id}
-            await transfer_ref.recv(virtual_engine, addr, src_to_dst)
+            await transfer_ref.recv(virtual_engine, from_rank, src_to_dst)
 
     async def _do_transfer(
         self,
         virtual_engine: int,
         local: Set[int],
-        remote: Dict[str, Set[Tuple[int, int, int]]],
+        remote: Dict[int, Set[Tuple[int, int, int]]],
         seq_group: SequenceGroup,
-        is_prefill: bool,
     ):
-        await self._do_transfer_inner(virtual_engine, remote)
-        # After the transfer is completed, update the corresponding metadata.
-        self._transfer_status[seq_group] = local
-        for _id in local:
-            self.block_manager.set_block_status_by_block_id("transferred", _id, True)
-        # After the transfer, place the `seq_group` back into the appropriate queue to
-        # wait for the next scheduling execution.
-        if is_prefill:
+        try:
+            await self._do_transfer_inner(virtual_engine, remote)
+        except Exception as e:
+            """
+            The exception here is most likely due to the sender triggering recovery during the transmission process.
+            In this case, fallback to performing computation during the prefill stage.
+            """
+            logger.error(f"Transfer failed: {e}")
+            # Force this `seq_group` to perform computation.
+            seq_group.force_calculation = True
+            self._transfer_status.pop(seq_group, None)
             self.waiting.appendleft(seq_group)
+            self._transferring.remove(seq_group)
         else:
-            self.running.appendleft(seq_group)
-        self._transferring.remove(seq_group)
+            # After the transfer is completed, update the corresponding metadata.
+            self._transfer_status[seq_group] = local
+            for _id in local:
+                self.block_manager.set_block_status_by_block_id(
+                    "transferred", _id, True
+                )
+            # After the transfer, place the `seq_group` back into the `waiting` queue to
+            # wait for the next scheduling execution.
+            self.waiting.appendleft(seq_group)
+            self._transferring.remove(seq_group)
 
     @no_type_check
     async def schedule(
@@ -240,39 +259,36 @@ class XavierScheduler(Scheduler):
             After completing the scheduling, the blocks have been allocated.
             Therefore, it is possible to check whether some blocks have already been computed on other replicas based on this information,
             and subsequently initiate the transfer.
+            According to the internal code comments in vllm,
+            whether `token_chunk_size` is 1 can indicate whether the `seq_group` is in the decode or prefill stage.
+            It is noted that data transmission is only applied during the prefill stage.
+            In the decode stage, it only applies to the last token of the block, which can negatively impact throughput.
             """
-            local, remote = await self._get_transfer_details(
-                virtual_engine, block_tables, seq_group
-            )
-            if remote:
-                running_seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
-                # According to the internal code comments in vllm,
-                # whether `token_chunk_size` is 1 can indicate whether the `seq_group` is in the decode or prefill stage.
-                is_prefill = token_chunk_size != 1
-                for seq in running_seqs:
-                    seq.status = (
-                        SequenceStatus.WAITING if is_prefill else SequenceStatus.RUNNING
-                    )
-                    # Additional attribute `transferred` to mark that this `seq_group` involves a transfer process.
-                    # During the next scheduling, block allocation will no longer be required
-                    # since it has already been completed.
-                    seq.transferred = True
-                    seq.data._stage = (
-                        SequenceStage.PREFILL if is_prefill else SequenceStage.DECODE
-                    )
-                self._transfer_status[seq_group] = set()
-                # Use `create_task` to avoid blocking subsequent scheduling.
-                asyncio.create_task(
-                    self._do_transfer(
-                        virtual_engine, local, remote, seq_group, is_prefill
-                    )
+            is_prefill: bool = token_chunk_size != 1
+            if is_prefill:
+                local, remote = await self._get_transfer_details(
+                    virtual_engine, block_tables, seq_group
                 )
-                # The `seq_group` that is currently being transferred enters a new queue.
-                self._transferring.append(seq_group)
-                has_transferring = True
-                continue
-            else:
-                scheduled_seq_groups.append(seq_group)
+                if remote:
+                    running_seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
+                    for seq in running_seqs:
+                        seq.status = SequenceStatus.WAITING
+                        # Additional attribute `transferred` to mark that this `seq_group` involves a transfer process.
+                        # During the next scheduling, block allocation will no longer be required
+                        # since it has already been completed.
+                        seq.transferred = True
+                        seq.data._stage = SequenceStage.PREFILL
+                    self._transfer_status[seq_group] = set()
+                    # Use `create_task` to avoid blocking subsequent scheduling.
+                    asyncio.create_task(
+                        self._do_transfer(virtual_engine, local, remote, seq_group)
+                    )
+                    # The `seq_group` that is currently being transferred enters a new queue.
+                    self._transferring.append(seq_group)
+                    has_transferring = True
+                    continue
+                else:
+                    scheduled_seq_groups.append(seq_group)
 
             if self.cache_config.enable_prefix_caching:
                 common_computed_block_nums = (

xinference/model/llm/vllm/xavier/test/test_xavier.py

@@ -21,11 +21,11 @@ from ..block_tracker import VLLMBlockTracker
 
 
 class ExtendedBlockTracker(VLLMBlockTracker):
-    def get_hash_to_address_and_block_id(self):
-        return self._hash_to_address_and_block_id
+    def get_hash_to_rank_and_block_id(self):
+        return self._hash_to_rank_and_block_id
 
-    def get_address_to_hash_and_block_id(self):
-        return self._address_to_hash_and_block_id
+    def get_rank_to_hash_and_block_id(self):
+        return self._rank_to_hash_and_block_id
 
 
 @pytest.fixture
@@ -53,53 +53,54 @@ async def test_block_tracker(actor_pool_context):
     )
 
     virtual_engine = 0
+    rank = 0
     block_infos = [(123, 0), (456, 1), (789, 2)]
 
     # register blocks
-    await tracker_ref.register_blocks(virtual_engine, block_infos, addr)
+    await tracker_ref.register_blocks(virtual_engine, block_infos, rank)
 
     # query blocks
     res = await tracker_ref.query_blocks(virtual_engine, [(123, 4), (789, 5)])
     assert len(res) == 1
-    assert addr in res
-    assert len(res[addr]) == 2
-    assert {x[0] for x in res[addr]} == {123, 789}
-    assert {x[1] for x in res[addr]} == {0, 2}
-    assert {x[2] for x in res[addr]} == {4, 5}
+    assert rank in res
+    assert len(res[rank]) == 2
+    assert {x[0] for x in res[rank]} == {123, 789}
+    assert {x[1] for x in res[rank]} == {0, 2}
+    assert {x[2] for x in res[rank]} == {4, 5}
 
     # query with extra info
     res = await tracker_ref.query_blocks(virtual_engine, [(123, 4), (789, 5), (110, 6)])
     assert len(res) == 1
-    assert addr in res
-    assert len(res[addr]) == 2
-    assert {x[0] for x in res[addr]} == {123, 789}
-    assert {x[1] for x in res[addr]} == {0, 2}
-    assert {x[2] for x in res[addr]} == {4, 5}
+    assert rank in res
+    assert len(res[rank]) == 2
+    assert {x[0] for x in res[rank]} == {123, 789}
+    assert {x[1] for x in res[rank]} == {0, 2}
+    assert {x[2] for x in res[rank]} == {4, 5}
 
     # unregister block
-    await tracker_ref.unregister_block(virtual_engine, addr, 1)
+    await tracker_ref.unregister_block(virtual_engine, rank, 1)
     res = await tracker_ref.query_blocks(virtual_engine, [(123, 4), (456, 7)])
     assert len(res) == 1
-    assert addr in res
-    assert len(res[addr]) == 1
-    assert {x[0] for x in res[addr]} == {123}
-    assert {x[1] for x in res[addr]} == {
+    assert rank in res
+    assert len(res[rank]) == 1
+    assert {x[0] for x in res[rank]} == {123}
+    assert {x[1] for x in res[rank]} == {
         0,
     }
-    assert {x[2] for x in res[addr]} == {
+    assert {x[2] for x in res[rank]} == {
         4,
     }
     # nothing happens
-    await tracker_ref.unregister_block(virtual_engine, addr, 3)
+    await tracker_ref.unregister_block(virtual_engine, rank, 3)
     res = await tracker_ref.query_blocks(virtual_engine, [(123, 4), (456, 7)])
     assert len(res) == 1
-    assert addr in res
-    assert len(res[addr]) == 1
-    assert {x[0] for x in res[addr]} == {123}
-    assert {x[1] for x in res[addr]} == {
+    assert rank in res
+    assert len(res[rank]) == 1
+    assert {x[0] for x in res[rank]} == {123}
+    assert {x[1] for x in res[rank]} == {
         0,
     }
-    assert {x[2] for x in res[addr]} == {
+    assert {x[2] for x in res[rank]} == {
         4,
     }
     # query returns empty
@@ -107,16 +108,40 @@ async def test_block_tracker(actor_pool_context):
     assert res == {}
 
     # check internal data
-    hash_to_address_and_block_id = await tracker_ref.get_hash_to_address_and_block_id()
-    assert virtual_engine in hash_to_address_and_block_id
-    assert hash_to_address_and_block_id[virtual_engine] == {
+    hash_to_rank_and_block_id = await tracker_ref.get_hash_to_rank_and_block_id()
+    assert virtual_engine in hash_to_rank_and_block_id
+    assert hash_to_rank_and_block_id[virtual_engine] == {
         123: {
-            (addr, 0),
+            (rank, 0),
         },
         456: set(),
-        789: {(addr, 2)},
+        789: {(rank, 2)},
     }
 
-    address_to_hash_and_block_id = await tracker_ref.get_address_to_hash_and_block_id()
-    assert virtual_engine in address_to_hash_and_block_id
-    assert address_to_hash_and_block_id[virtual_engine] == {addr: {(123, 0), (789, 2)}}
+    rank_to_hash_and_block_id = await tracker_ref.get_rank_to_hash_and_block_id()
+    assert virtual_engine in rank_to_hash_and_block_id
+    assert rank_to_hash_and_block_id[virtual_engine] == {rank: {(123, 0), (789, 2)}}
+
+    # register blocks
+    new_rank = 1
+    block_infos = [(111, 7), (222, 8), (333, 9), (123, 10)]
+    await tracker_ref.register_blocks(virtual_engine, block_infos, new_rank)
+
+    # test unregister rank
+    await tracker_ref.unregister_rank(0)
+    res = await tracker_ref.query_blocks(virtual_engine, [(789, 5)])
+    assert len(res) == 0
+    res = await tracker_ref.query_blocks(virtual_engine, [(123, 6)])
+    assert len(res) == 1
+    assert new_rank in res
+
+    # check internal data
+    rank_to_hash_and_block_id = await tracker_ref.get_rank_to_hash_and_block_id()
+    assert rank in rank_to_hash_and_block_id[virtual_engine]
+    assert new_rank in rank_to_hash_and_block_id[virtual_engine]
+
+    # test register rank
+    await tracker_ref.register_rank(0)
+    rank_to_hash_and_block_id = await tracker_ref.get_rank_to_hash_and_block_id()
+    assert rank not in rank_to_hash_and_block_id[virtual_engine]
+    assert new_rank in rank_to_hash_and_block_id[virtual_engine]