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

xinference/client/restful/restful_client.py

@@ -126,6 +126,43 @@ class RESTfulEmbeddingModelHandle(RESTfulModelHandle):
         response_data = response.json()
         return response_data
 
+    def convert_ids_to_tokens(
+        self, input: Union[List, List[List]], **kwargs
+    ) -> List[str]:
+        """
+        Convert token IDs to human readable tokens via RESTful APIs.
+
+        Parameters
+        ----------
+        input: Union[List, List[List]]
+            Input token IDs to convert, can be a single list of token IDs or a list of token ID lists.
+            To convert multiple sequences in a single request, pass a list of token ID lists.
+
+        Returns
+        -------
+        list
+            A list of decoded tokens in human readable format.
+
+        Raises
+        ------
+        RuntimeError
+            Report the failure of token conversion and provide the error message.
+
+        """
+        url = f"{self._base_url}/v1/convert_ids_to_tokens"
+        request_body = {
+            "model": self._model_uid,
+            "input": input,
+        }
+        request_body.update(kwargs)
+        response = requests.post(url, json=request_body, headers=self.auth_headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to decode token ids, detail: {_get_error_string(response)}"
+            )
+        response_data = response.json()
+        return response_data
+
 
 class RESTfulRerankModelHandle(RESTfulModelHandle):
     def rerank(
@@ -174,6 +211,7 @@ class RESTfulRerankModelHandle(RESTfulModelHandle):
             "max_chunks_per_doc": max_chunks_per_doc,
             "return_documents": return_documents,
             "return_len": return_len,
+            "kwargs": json.dumps(kwargs),
         }
         request_body.update(kwargs)
         response = requests.post(url, json=request_body, headers=self.auth_headers)
@@ -703,6 +741,8 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
             The speed of the generated audio.
         stream: bool
             Use stream or not.
+        prompt_speech: bytes
+            The audio bytes to be provided to the model.
 
         Returns
         -------
@@ -1357,7 +1397,7 @@ class Client:
         response_data = response.json()
         return response_data
 
-    def abort_request(self, model_uid: str, request_id: str):
+    def abort_request(self, model_uid: str, request_id: str, block_duration: int = 30):
         """
         Abort a request.
         Abort a submitted request. If the request is finished or not found, this method will be a no-op.
@@ -1369,13 +1409,18 @@ class Client:
             Model uid.
         request_id: str
             Request id.
+        block_duration: int
+            The duration to make the request id abort. If set to 0, the abort_request will be immediate, which may
+            prevent it from taking effect if it arrives before the request operation.
         Returns
         -------
         Dict
             Return empty dict.
         """
         url = f"{self.base_url}/v1/models/{model_uid}/requests/{request_id}/abort"
-        response = requests.post(url, headers=self._headers)
+        response = requests.post(
+            url, headers=self._headers, json={"block_duration": block_duration}
+        )
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to abort request, detail: {_get_error_string(response)}"

xinference/constants.py

@@ -88,3 +88,4 @@ XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE = os.environ.get(
     XINFERENCE_ENV_TEXT_TO_IMAGE_BATCHING_SIZE, None
 )
 XINFERENCE_LAUNCH_MODEL_RETRY = 3
+XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION = 30

xinference/core/chat_interface.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import base64
+import html
 import logging
 import os
 from io import BytesIO
@@ -137,7 +138,11 @@ class GradioInterface:
                 if "content" not in delta:
                     continue
                 else:
-                    response_content += delta["content"]
+                    # some model like deepseek-r1-distill-qwen
+                    # will generate <think>...</think> ...
+                    # in gradio, no output will be rendered,
+                    # thus escape html tags in advance
+                    response_content += html.escape(delta["content"])
                     yield response_content
 
             yield response_content

xinference/core/model.py

@@ -35,12 +35,14 @@ from typing import (
     List,
     Optional,
     Union,
+    no_type_check,
 )
 
 import sse_starlette.sse
 import xoscar as xo
 
 from ..constants import (
+    XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
     XINFERENCE_LAUNCH_MODEL_RETRY,
     XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE,
 )
@@ -57,7 +59,7 @@ import logging
 logger = logging.getLogger(__name__)
 
 from ..device_utils import empty_cache
-from .utils import json_dumps, log_async
+from .utils import CancelMixin, json_dumps, log_async
 
 try:
     from torch.cuda import OutOfMemoryError
@@ -77,6 +79,9 @@ XINFERENCE_BATCHING_ALLOWED_VISION_MODELS = [
 ]
 
 XINFERENCE_TEXT_TO_IMAGE_BATCHING_ALLOWED_MODELS = ["FLUX.1-dev", "FLUX.1-schnell"]
+XINFERENCE_TEST_OUT_OF_MEMORY_ERROR = bool(
+    os.getenv("XINFERENCE_TEST_OUT_OF_MEMORY_ERROR", False)
+)
 
 
 def request_limit(fn):
@@ -90,21 +95,26 @@ def request_limit(fn):
         logger.debug(
             f"Request {fn.__name__}, current serve request count: {self._serve_count}, request limit: {self._request_limits} for the model {self.model_uid()}"
         )
-        if self._request_limits is not None:
-            if 1 + self._serve_count <= self._request_limits:
-                self._serve_count += 1
-            else:
-                raise RuntimeError(
-                    f"Rate limit reached for the model. Request limit {self._request_limits} for the model: {self.model_uid()}"
-                )
+        if 1 + self._serve_count <= self._request_limits:
+            self._serve_count += 1
+        else:
+            raise RuntimeError(
+                f"Rate limit reached for the model. Request limit {self._request_limits} for the model: {self.model_uid()}"
+            )
+        ret = None
         try:
             ret = await fn(self, *args, **kwargs)
         finally:
-            if self._request_limits is not None:
+            if ret is not None and (
+                inspect.isasyncgen(ret) or inspect.isgenerator(ret)
+            ):
+                # stream case, let client call model_ref to decrease self._serve_count
+                pass
+            else:
                 self._serve_count -= 1
-            logger.debug(
-                f"After request {fn.__name__}, current serve request count: {self._serve_count} for the model {self.model_uid()}"
-            )
+                logger.debug(
+                    f"After request {fn.__name__}, current serve request count: {self._serve_count} for the model {self.model_uid()}"
+                )
         return ret
 
     return wrapped_func
@@ -112,20 +122,25 @@ def request_limit(fn):
 
 def oom_check(fn):
     @functools.wraps(fn)
-    def _wrapper(*args, **kwargs):
+    def _wrapper(self, *args, **kwargs):
         try:
-            return fn(*args, **kwargs)
-        except OutOfMemoryError:
-            logger.exception("Model actor is out of memory.")
-            os._exit(1)
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
+            return fn(self, *args, **kwargs)
+        except OutOfMemoryError as ex:
+            assert self._loop is not None
+            asyncio.run_coroutine_threadsafe(
+                self._handle_oom_error(ex), loop=self._loop
+            )
 
     @functools.wraps(fn)
-    async def _async_wrapper(*args, **kwargs):
+    async def _async_wrapper(self, *args, **kwargs):
         try:
-            return await fn(*args, **kwargs)
-        except OutOfMemoryError:
-            logger.exception("Model actor is out of memory.")
-            os._exit(1)
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
+            return await fn(self, *args, **kwargs)
+        except OutOfMemoryError as ex:
+            await self._handle_oom_error(ex)
 
     assert not inspect.isasyncgen(fn)
     assert not inspect.isgenerator(fn)
@@ -136,7 +151,7 @@ def oom_check(fn):
         return _wrapper
 
 
-class ModelActor(xo.StatelessActor):
+class ModelActor(xo.StatelessActor, CancelMixin):
     _replica_model_uid: Optional[str]
 
     @classmethod
@@ -172,6 +187,16 @@ class ModelActor(xo.StatelessActor):
         if hasattr(self._model, "stop") and callable(self._model.stop):
             self._model.stop()
 
+        if isinstance(self._model, LLMVLLMModel):
+            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, address: {self.address}, error: {e}"
+                    )
+
         if (
             isinstance(self._model, (LLMPytorchModel, LLMVLLMModel, SGLANGModel))
             and self._model.model_spec.model_format == "pytorch"
@@ -200,6 +225,7 @@ class ModelActor(xo.StatelessActor):
         replica_model_uid: str,
         model_description: Optional["ModelDescription"] = None,
         request_limits: Optional[int] = None,
+        xavier_config: Optional[Dict] = None,
     ):
         super().__init__()
         from ..model.llm.lmdeploy.core import LMDeployModel
@@ -214,7 +240,9 @@ class ModelActor(xo.StatelessActor):
         self._model_description = (
             model_description.to_dict() if model_description else {}
         )
-        self._request_limits = request_limits
+        self._request_limits = (
+            float("inf") if request_limits is None else request_limits
+        )
         self._pending_requests: asyncio.Queue = asyncio.Queue()
         self._handle_pending_requests_task = None
         self._lock = (
@@ -239,6 +267,11 @@ class ModelActor(xo.StatelessActor):
         self._scheduler_ref = None
         self._text_to_image_scheduler_ref = None
 
+        if isinstance(self._model, VLLMModel):
+            self._xavier_config = xavier_config
+            self._model.set_xavier_config(xavier_config)
+            self._transfer_ref = None
+
     async def __post_create__(self):
         self._loop = asyncio.get_running_loop()
 
@@ -267,6 +300,32 @@ class ModelActor(xo.StatelessActor):
     def __repr__(self) -> str:
         return f"ModelActor({self._replica_model_uid})"
 
+    def decrease_serve_count(self):
+        self._serve_count -= 1
+
+    @no_type_check
+    async def start_transfer_for_vllm(self, rank_addresses: List[str]):
+        from ..model.llm.vllm.core import VLLMModel
+        from ..model.llm.vllm.xavier.transfer import TransferActor
+
+        assert isinstance(self._model, VLLMModel)
+        rank = self._xavier_config.get("rank")  # type: ignore
+        self._transfer_ref = await xo.create_actor(
+            TransferActor,
+            address=self.address,
+            uid=f"{TransferActor.default_uid()}-{rank}",
+            rank=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,
+        )
+        await self._model.init_xavier()
+        logger.debug(
+            f"Init transfer actor: {self._transfer_ref.address}, rank: {rank} done for vllm."  # type: ignore
+        )
+
     async def _record_completion_metrics(
         self, duration, completion_tokens, prompt_tokens
     ):
@@ -429,11 +488,24 @@ class ModelActor(xo.StatelessActor):
             )
         )
 
+    async def _handle_oom_error(self, ex):
+        error_message = (
+            f"Model actor is out of memory, model id: {self.model_uid()}, error: {ex}"
+        )
+        logger.exception(error_message)
+        worker_ref = await self._get_worker_ref()
+        await worker_ref.update_model_status(
+            self._replica_model_uid, last_error=error_message
+        )
+        os._exit(1)
+
     def _to_generator(self, output_type: str, gen: types.GeneratorType):
         start_time = time.time()
         time_to_first_token = None
         final_usage = None
         try:
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
             for v in gen:
                 if time_to_first_token is None:
                     time_to_first_token = (time.time() - start_time) * 1000
@@ -445,11 +517,11 @@ class ModelActor(xo.StatelessActor):
                         output_type == "binary"
                     ), f"Unknown output type '{output_type}'"
                 yield sse_starlette.sse.ensure_bytes(v, None)
-        except OutOfMemoryError:
-            logger.exception(
-                "Model actor is out of memory, model id: %s", self.model_uid()
+        except OutOfMemoryError as ex:
+            assert self._loop is not None
+            asyncio.run_coroutine_threadsafe(
+                self._handle_oom_error(ex), loop=self._loop
             )
-            os._exit(1)
         finally:
             if self._loop is not None and time_to_first_token is not None:
                 coro = self.record_metrics(
@@ -471,6 +543,8 @@ class ModelActor(xo.StatelessActor):
         time_to_first_token = None
         final_usage = None
         try:
+            if XINFERENCE_TEST_OUT_OF_MEMORY_ERROR:
+                raise OutOfMemoryError("Test Out of Memory Error")
             async for v in gen:
                 if time_to_first_token is None:
                     time_to_first_token = (time.time() - start_time) * 1000
@@ -483,11 +557,8 @@ class ModelActor(xo.StatelessActor):
                         output_type == "binary"
                     ), f"Unknown output type '{output_type}'"
                 yield await asyncio.to_thread(sse_starlette.sse.ensure_bytes, v, None)
-        except OutOfMemoryError:
-            logger.exception(
-                "Model actor is out of memory, model id: %s", self.model_uid()
-            )
-            os._exit(1)
+        except OutOfMemoryError as ex:
+            await self._handle_oom_error(ex)
         finally:
             coros = []
             if time_to_first_token is not None:
@@ -553,6 +624,7 @@ class ModelActor(xo.StatelessActor):
 
     @oom_check
     async def _call_wrapper(self, output_type: str, fn: Callable, *args, **kwargs):
+        self._add_running_task(kwargs.get("request_id"))
         if self._lock is None:
             if inspect.iscoroutinefunction(fn):
                 ret = await fn(*args, **kwargs)
@@ -761,9 +833,14 @@ class ModelActor(xo.StatelessActor):
                     prompt_tokens,
                 )
 
-    async def abort_request(self, request_id: str) -> str:
+    async def abort_request(
+        self,
+        request_id: str,
+        block_duration: int = XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
+    ) -> str:
         from .utils import AbortRequestMessage
 
+        self._cancel_running_task(request_id, block_duration)
         if self.allow_batching():
             if self._scheduler_ref is None:
                 return AbortRequestMessage.NOT_FOUND.name
@@ -787,6 +864,19 @@ class ModelActor(xo.StatelessActor):
             f"Model {self._model.model_spec} is not for creating embedding."
         )
 
+    @request_limit
+    @log_async(logger=logger)
+    async def convert_ids_to_tokens(
+        self, input: Union[List, List[List]], *args, **kwargs
+    ):
+        kwargs.pop("request_id", None)
+        if hasattr(self._model, "convert_ids_to_tokens"):
+            return await self._call_wrapper_json(
+                self._model.convert_ids_to_tokens, input, *args, **kwargs
+            )
+
+        raise AttributeError(f"Model {self._model.model_spec} can convert token id.")
+
     @request_limit
     @log_async(logger=logger)
     async def rerank(

xinference/core/supervisor.py

@@ -35,6 +35,7 @@ from typing import (
 import xoscar as xo
 
 from ..constants import (
+    XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
     XINFERENCE_DISABLE_HEALTH_CHECK,
     XINFERENCE_HEALTH_CHECK_FAILURE_THRESHOLD,
     XINFERENCE_HEALTH_CHECK_INTERVAL,
@@ -266,6 +267,14 @@ class SupervisorActor(xo.StatelessActor):
                 signal.SIGTERM, lambda: asyncio.create_task(signal_handler())
             )
 
+        from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
+        from ..model.llm.vllm.xavier.collective_manager import CollectiveManager
+
+        self._block_tracker_mapping: Dict[str, xo.ActorRefType[VLLMBlockTracker]] = {}
+        self._collective_manager_mapping: Dict[
+            str, xo.ActorRefType[CollectiveManager]
+        ] = {}
+
     @typing.no_type_check
     async def get_cluster_device_info(self, detailed: bool = False) -> List:
         import psutil
@@ -958,29 +967,83 @@ class SupervisorActor(xo.StatelessActor):
         if model_uid is None:
             model_uid = self._gen_model_uid(model_name)
 
+        # Xavier-related
+        enable_xavier: bool = (
+            bool(kwargs.pop("enable_xavier", False))
+            and model_engine is not None
+            and model_engine.lower() == "vllm"
+        )
+        store_address = None
+        store_port = None
+        world_size = None
+        if enable_xavier:
+            if replica <= 1:
+                logger.warning(f"Enabling xavier when `replica<=1` is meaningless.")
+                enable_xavier = False
+            else:
+                from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
+                from ..model.llm.vllm.xavier.collective_manager import CollectiveManager
+
+                self._block_tracker_mapping[model_uid] = await xo.create_actor(
+                    VLLMBlockTracker,
+                    address=self.address,
+                    uid=f"{VLLMBlockTracker.default_uid()}-{model_uid}",
+                )
+                world_size = replica + 1
+                logger.info(f"Going to start xavier with world size: {world_size}")
+                self._collective_manager_mapping[model_uid] = await xo.create_actor(
+                    CollectiveManager,
+                    address=self.address,
+                    uid=f"{CollectiveManager.default_uid()}-{model_uid}",
+                    model_uid=model_uid,
+                )
+                logger.info(f"Start collective manager for {model_uid} done.")
+
         model_size = str(model_size_in_billions) if model_size_in_billions else ""
         logger.debug(
             f"Enter launch_builtin_model, model_uid: {model_uid}, model_name: {model_name}, model_size: {model_size}, "
-            f"model_format: {model_format}, quantization: {quantization}, replica: {replica}, "
+            f"model_format: {model_format}, quantization: {quantization}, replica: {replica}, enable_xavier: {enable_xavier}, "
             f"kwargs: {kwargs}"
         )
 
-        async def _launch_one_model(_replica_model_uid):
+        async def _launch_one_model(worker_ref, _replica_model_uid, rank: int):
             if _replica_model_uid in self._replica_model_uid_to_worker:
                 raise ValueError(
                     f"Model is already in the model list, uid: {_replica_model_uid}"
                 )
+
+            nonlocal store_address
+            nonlocal store_port
+            xavier_config = (
+                {
+                    "block_tracker_uid": self._block_tracker_mapping[model_uid].uid,
+                    "block_tracker_address": self._block_tracker_mapping[
+                        model_uid
+                    ].address,
+                    "rank": rank,
+                    "world_size": world_size,
+                    "store_address": store_address,
+                    "store_port": store_port,
+                }
+                if enable_xavier
+                else None
+            )
+
+            if enable_xavier and rank == 0:
+                rank0_address, _port = await worker_ref.launch_rank0_model(
+                    _replica_model_uid, xavier_config
+                )
+                self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref
+                store_address = rank0_address.split(":")[0]
+                store_port = _port
+                return rank0_address
+
             replica_gpu_idx = assign_replica_gpu(_replica_model_uid, replica, gpu_idx)
             nonlocal model_type
 
-            worker_ref = (
-                target_ip_worker_ref
-                if target_ip_worker_ref is not None
-                else await self._choose_worker()
-            )
             # LLM as default for compatibility
             model_type = model_type or "LLM"
-            await worker_ref.launch_builtin_model(
+            subpool_address = await worker_ref.launch_builtin_model(
                 model_uid=_replica_model_uid,
                 model_name=model_name,
                 model_size_in_billions=model_size_in_billions,
@@ -994,14 +1057,64 @@ class SupervisorActor(xo.StatelessActor):
                 gpu_idx=replica_gpu_idx,
                 download_hub=download_hub,
                 model_path=model_path,
+                xavier_config=xavier_config,
                 **kwargs,
             )
             self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref
+            return subpool_address
 
         async def _launch_model():
             try:
-                for rep_model_uid in iter_replica_model_uid(model_uid, replica):
-                    await _launch_one_model(rep_model_uid)
+                worker_refs = []
+                rank_addresses = []
+                for _idx, rep_model_uid in enumerate(
+                    iter_replica_model_uid(model_uid, replica)
+                ):
+                    worker_ref = (
+                        target_ip_worker_ref
+                        if target_ip_worker_ref is not None
+                        else await self._choose_worker()
+                    )
+                    if enable_xavier and _idx == 0:
+                        """
+                        Start the rank 0 model actor on the worker that holds the rank 1 replica,
+                        solely for constructing the collective communication world.
+                        """
+                        _uid = model_uid + "-rank0"
+                        rank0_address = await _launch_one_model(worker_ref, _uid, 0)
+                        worker_refs.append((worker_ref, _uid))
+                        rank_addresses.append(rank0_address)
+
+                    subpool_address = await _launch_one_model(
+                        worker_ref, rep_model_uid, _idx + 1
+                    )
+                    worker_refs.append((worker_ref, rep_model_uid))
+                    rank_addresses.append(subpool_address)
+
+                # For xavier, start all the vllm instances first,
+                # and then start the transfer component,
+                # because the transfer actor needs all the rank addresses used for collective communication
+                if enable_xavier:
+                    logger.debug(f"Init transfer component for xavier...")
+                    collective_manager_ref = self._collective_manager_mapping[model_uid]
+                    tasks = []
+                    for worker_ref, rep_model_uid in worker_refs:
+                        tasks.append(
+                            worker_ref.start_transfer_for_vllm(
+                                rep_model_uid, rank_addresses
+                            )
+                        )
+                    # Here you must use asyncio.gather, not a for loop,
+                    # or you will get stuck.
+                    await asyncio.gather(*tasks)
+
+                    # init collective_manager
+                    for idx, addr in enumerate(rank_addresses):
+                        await collective_manager_ref.register_rank(
+                            idx, addr, update=False
+                        )
+
+                    logger.debug(f"Init transfer component for xavier done.")
             except Exception:
                 # terminate_model will remove the replica info.
                 await self.terminate_model(model_uid, suppress_exception=True)
@@ -1130,6 +1243,38 @@ class SupervisorActor(xo.StatelessActor):
                     raise
         self._model_uid_to_replica_info.pop(model_uid, None)
 
+        # clear for xavier
+        rank0_uid = model_uid + "-rank0"
+        if rank0_uid in self._replica_model_uid_to_worker:
+            await _terminate_one_model(rank0_uid)
+
+        collective_manager_ref = self._collective_manager_mapping.pop(model_uid, None)
+        if collective_manager_ref is not None:
+            try:
+                await xo.destroy_actor(collective_manager_ref)
+            except Exception as e:
+                logger.debug(
+                    "Destroy collective_manager_ref failed, model uid: %s, error: %s",
+                    model_uid,
+                    e,
+                )
+            finally:
+                logger.debug(
+                    f"Destroy collective_manager_ref done. model uid: {model_uid}"
+                )
+        block_tracker_ref = self._block_tracker_mapping.pop(model_uid, None)
+        if block_tracker_ref is not None:
+            try:
+                await xo.destroy_actor(block_tracker_ref)
+            except Exception as e:
+                logger.debug(
+                    "Destroy block_tracker_ref failed, model uid: %s, error: %s",
+                    model_uid,
+                    e,
+                )
+            finally:
+                logger.debug(f"Destroy block_tracker_ref done. model uid: {model_uid}")
+
     @log_async(logger=logger)
     async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
         replica_info = self._model_uid_to_replica_info.get(model_uid, None)
@@ -1147,6 +1292,15 @@ class SupervisorActor(xo.StatelessActor):
             )
         return await worker_ref.get_model(model_uid=replica_model_uid)
 
+    @log_async(logger=logger)
+    async def get_model_status(self, replica_model_uid: str):
+        worker_ref = self._replica_model_uid_to_worker.get(replica_model_uid, None)
+        if worker_ref is None:
+            raise ValueError(
+                f"Model not found in the model list, uid: {replica_model_uid}"
+            )
+        return await worker_ref.get_model_status(replica_model_uid)
+
     @log_async(logger=logger)
     async def describe_model(self, model_uid: str) -> Dict[str, Any]:
         replica_info = self._model_uid_to_replica_info.get(model_uid, None)
@@ -1213,7 +1367,12 @@ class SupervisorActor(xo.StatelessActor):
         return cached_models
 
     @log_async(logger=logger)
-    async def abort_request(self, model_uid: str, request_id: str) -> Dict:
+    async def abort_request(
+        self,
+        model_uid: str,
+        request_id: str,
+        block_duration: int = XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
+    ) -> Dict:
         from .scheduler import AbortRequestMessage
 
         res = {"msg": AbortRequestMessage.NO_OP.name}
@@ -1228,7 +1387,7 @@ class SupervisorActor(xo.StatelessActor):
             if worker_ref is None:
                 continue
             model_ref = await worker_ref.get_model(model_uid=rep_mid)
-            result_info = await model_ref.abort_request(request_id)
+            result_info = await model_ref.abort_request(request_id, block_duration)
             res["msg"] = result_info
             if result_info == AbortRequestMessage.DONE.name:
                 break
@@ -1371,3 +1530,12 @@ class SupervisorActor(xo.StatelessActor):
 
     async def get_progress(self, request_id: str) -> float:
         return await self._progress_tracker.get_progress(request_id)
+
+    async def call_collective_manager(
+        self, model_uid: str, func_name: str, *args, **kwargs
+    ):
+        """
+        Used by worker.
+        """
+        collective_manager_ref = self._collective_manager_mapping[model_uid]
+        await getattr(collective_manager_ref, func_name)(*args, **kwargs)