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

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

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

@@ -0,0 +1,247 @@
+# 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 Dict, List, Optional, Type, Union
+
+from vllm import AsyncEngineArgs, EmbeddingRequestOutput, RequestOutput
+from vllm.config import VllmConfig
+from vllm.engine.async_llm_engine import AsyncLLMEngine, _AsyncLLMEngine
+from vllm.engine.llm_engine import SchedulerOutputState
+from vllm.engine.metrics_types import StatLoggerBase
+from vllm.executor.executor_base import ExecutorBase
+from vllm.sequence import ExecuteModelRequest
+from vllm.usage.usage_lib import UsageContext
+
+from .executor import XavierExecutor
+from .scheduler import XavierScheduler
+
+logger = logging.getLogger(__name__)
+
+
+class XavierInternalEngine(_AsyncLLMEngine):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._xavier_config = kwargs["vllm_config"].xavier_config
+        self.scheduler = [
+            XavierScheduler(
+                self.scheduler_config,
+                self.cache_config,
+                self.lora_config,
+                self.parallel_config.pipeline_parallel_size,
+                self.async_callbacks[v_id]
+                if self.model_config.use_async_output_proc
+                else None,
+                xavier_config=self._xavier_config,
+                virtual_engine=v_id,
+            )
+            for v_id in range(self.parallel_config.pipeline_parallel_size)
+        ]
+        self.output_processor.scheduler = self.scheduler
+        self.model_executor.scheduler = self.scheduler
+
+    async def step_async(
+        self, virtual_engine: int
+    ) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
+        """Performs one decoding iteration and returns newly generated results.
+        The workers are ran asynchronously if possible.
+
+        This function performs one decoding iteration of the engine. It first
+        schedules the sequences to be executed in the next iteration and the
+        token blocks to be swapped in/out/copy. Then, it executes the model
+        and updates the scheduler with the model outputs. Finally, it decodes
+        the sequences and returns the newly generated results.
+        """
+        # these are cached outputs from previous iterations. None if on first
+        # iteration
+        cached_outputs = self.cached_scheduler_outputs[virtual_engine]
+        seq_group_metadata_list = cached_outputs.seq_group_metadata_list
+        scheduler_outputs = cached_outputs.scheduler_outputs
+        allow_async_output_proc = cached_outputs.allow_async_output_proc
+
+        ctx = self.scheduler_contexts[virtual_engine]
+
+        # Clear outputs for each new scheduler iteration
+        ctx.request_outputs.clear()
+
+        # skip the scheduler if there are any remaining steps in the seq groups.
+        # This ensures that the scheduler is only called again when the current
+        # batch has completed.
+        if not self._has_remaining_steps(seq_group_metadata_list):
+            # Schedule iteration
+            """Xinference Change!!!
+            Why copy the entire function code of vllm:
+            The purpose here is to modify the way the `schedule` function is invoked to asynchronous calling.
+            No other modifications were made elsewhere.
+            """
+            (
+                seq_group_metadata_list,
+                scheduler_outputs,
+                allow_async_output_proc,
+            ) = await self.scheduler[virtual_engine].schedule()
+
+            ctx.seq_group_metadata_list = seq_group_metadata_list
+            ctx.scheduler_outputs = scheduler_outputs
+
+            # Maybe switch from async mode to sync mode
+            if not allow_async_output_proc and len(ctx.output_queue) > 0:
+                self._process_model_outputs(ctx=ctx)
+
+            if (
+                self.scheduler_config.is_multi_step
+                and scheduler_outputs.num_lookahead_slots > 0
+            ):
+                # cache the scheduler outputs for the next iteration if we have
+                # lookahead slots
+                self._cache_scheduler_outputs_for_multi_step(
+                    virtual_engine,
+                    seq_group_metadata_list,
+                    scheduler_outputs,
+                    allow_async_output_proc,
+                )
+
+        assert seq_group_metadata_list is not None
+        assert scheduler_outputs is not None
+
+        if not scheduler_outputs.is_empty():
+            finished_requests_ids = self.scheduler[
+                virtual_engine
+            ].get_and_reset_finished_requests_ids()
+
+            # Check if we have a cached last_output from the previous iteration.
+            # For supporting PP this is probably the best way to pass the
+            # sampled_token_ids, as a separate broadcast over all the PP stages
+            # will cause one virtual engine's microbatch to block the pipeline.
+            last_sampled_token_ids = self._get_last_sampled_token_ids(virtual_engine)
+
+            execute_model_req = ExecuteModelRequest(
+                seq_group_metadata_list=seq_group_metadata_list,
+                blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
+                blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
+                blocks_to_copy=scheduler_outputs.blocks_to_copy,
+                virtual_engine=virtual_engine,
+                num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
+                running_queue_size=scheduler_outputs.running_queue_size,
+                finished_requests_ids=finished_requests_ids,
+                # We use ExecuteModelRequest to pass the last sampled_token_ids
+                # to each of the non-last PP stages for in-place prepare_input.
+                last_sampled_token_ids=last_sampled_token_ids,
+            )
+
+            if allow_async_output_proc:
+                execute_model_req.async_callback = self.async_callbacks[virtual_engine]
+
+            # Execute the model.
+            outputs = await self.model_executor.execute_model_async(execute_model_req)
+
+            # we need to do this here so that last step's sampled_token_ids can
+            # be passed to the next iteration for PP.
+            if self.scheduler_config.is_multi_step:
+                self._update_cached_scheduler_output(virtual_engine, outputs)
+        else:
+            if len(ctx.output_queue) > 0:
+                self._process_model_outputs(ctx=ctx)
+            outputs = []
+
+        # Finish the current step for all the sequence groups.
+        if self.scheduler_config.is_multi_step:
+            for seq_group in seq_group_metadata_list:
+                seq_group.finish_step()
+
+        if not self._has_remaining_steps(seq_group_metadata_list):
+            # Clear the cache if we have finished all the steps
+            if self.scheduler_config.is_multi_step:
+                self.cached_scheduler_outputs[virtual_engine] = SchedulerOutputState()
+
+            # is_first_step_output is True only when the num_steps of all
+            # the sequences are 1. When the num_steps > 1,
+            # multi_step_model_runner does the first-step output append.
+            is_first_step_output: bool = (
+                False
+                if not seq_group_metadata_list
+                else seq_group_metadata_list[0].state.num_steps == 1
+            )
+
+            ctx.append_output(
+                outputs=outputs,
+                seq_group_metadata_list=seq_group_metadata_list,
+                scheduler_outputs=scheduler_outputs,
+                is_async=allow_async_output_proc,
+                is_last_step=True,
+                is_first_step_output=is_first_step_output,
+            )
+
+            if outputs and allow_async_output_proc:
+                assert (
+                    len(outputs) == 1
+                ), "Async postprocessor expects only a single output set"
+                self._advance_to_next_step(
+                    outputs[0],
+                    seq_group_metadata_list,
+                    scheduler_outputs.scheduled_seq_groups,
+                )
+
+            if not allow_async_output_proc:
+                self._process_model_outputs(ctx=ctx)
+
+                # Log stats.
+                self.do_log_stats(scheduler_outputs, outputs)
+
+                # Tracing
+                self.do_tracing(scheduler_outputs)
+
+        else:
+            # Multi-step case
+            return ctx.request_outputs
+
+        if not self.has_unfinished_requests():
+            # Drain async postprocessor (if exists)
+            if len(ctx.output_queue) > 0:
+                self._process_model_outputs(ctx=ctx)
+            assert len(ctx.output_queue) == 0
+
+        return ctx.request_outputs
+
+
+class XavierEngine(AsyncLLMEngine):
+    _engine_class: Type[_AsyncLLMEngine] = XavierInternalEngine
+    _xavier_config: Optional[Dict] = None
+
+    @classmethod
+    def _get_executor_cls(cls, engine_config: VllmConfig) -> Type[ExecutorBase]:
+        logger.debug(f"Initializing Xavier executor.")
+        return XavierExecutor
+
+    @classmethod
+    def from_engine_args(
+        cls,
+        engine_args: AsyncEngineArgs,
+        engine_config: Optional[VllmConfig] = None,
+        start_engine_loop: bool = True,
+        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
+        stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,
+        xavier_config: Optional[Dict] = None,
+    ) -> "AsyncLLMEngine":
+        cls._xavier_config = xavier_config
+        return super().from_engine_args(
+            engine_args, engine_config, start_engine_loop, usage_context, stat_loggers
+        )
+
+    def __init__(self, *args, **kwargs):
+        # set xavier_config to `vllm_config`,
+        # because it may be needed everywhere in the vllm internal components
+        kwargs["vllm_config"].xavier_config = self._xavier_config
+        super().__init__(*args, **kwargs)
+
+    async def init_xavier(self):
+        await self.engine.model_executor.init_transfer()