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

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-12-27T18:14:37+0800",
+ "date": "2025-01-10T17:24:10+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "d3428697115cc4666b38b32925ba28bdc1a21957",
- "version": "1.1.1"
+ "full-revisionid": "df45f11115051929d6296a0c138b99472abf497f",
+ "version": "1.2.0"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -1214,6 +1214,19 @@ class RESTfulAPI(CancelMixin):
     async def get_address(self) -> JSONResponse:
         return JSONResponse(content=self._supervisor_address)
 
+    async def _get_model_last_error(self, replica_model_uid: bytes, e: Exception):
+        if not isinstance(e, xo.ServerClosed):
+            return e
+        try:
+            model_status = await (await self._get_supervisor_ref()).get_model_status(
+                replica_model_uid.decode("utf-8")
+            )
+            if model_status is not None and model_status.last_error:
+                return Exception(model_status.last_error)
+        except Exception as ex:
+            return ex
+        return e
+
     async def create_completion(self, request: Request) -> Response:
         raw_body = await request.json()
         body = CreateCompletionRequest.parse_obj(raw_body)
@@ -1272,6 +1285,7 @@ class RESTfulAPI(CancelMixin):
                     )
                     return
                 except Exception as ex:
+                    ex = await self._get_model_last_error(model.uid, ex)
                     logger.exception("Completion stream got an error: %s", ex)
                     await self._report_error_event(model_uid, str(ex))
                     # https://github.com/openai/openai-python/blob/e0aafc6c1a45334ac889fe3e54957d309c3af93f/src/openai/_streaming.py#L107
@@ -1286,6 +1300,7 @@ class RESTfulAPI(CancelMixin):
                 data = await model.generate(body.prompt, kwargs, raw_params=raw_kwargs)
                 return Response(data, media_type="application/json")
             except Exception as e:
+                e = await self._get_model_last_error(model.uid, e)
                 logger.error(e, exc_info=True)
                 await self._report_error_event(model_uid, str(e))
                 self.handle_request_limit_error(e)
@@ -1317,14 +1332,11 @@ class RESTfulAPI(CancelMixin):
         try:
             embedding = await model.create_embedding(body.input, **kwargs)
             return Response(embedding, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def convert_ids_to_tokens(self, request: Request) -> Response:
@@ -1352,14 +1364,11 @@ class RESTfulAPI(CancelMixin):
         try:
             decoded_texts = await model.convert_ids_to_tokens(body.input, **kwargs)
             return Response(decoded_texts, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def rerank(self, request: Request) -> Response:
@@ -1393,14 +1402,11 @@ class RESTfulAPI(CancelMixin):
                 **parsed_kwargs,
             )
             return Response(scores, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_transcriptions(
@@ -1445,13 +1451,11 @@ class RESTfulAPI(CancelMixin):
                 **parsed_kwargs,
             )
             return Response(content=transcription, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_translations(
@@ -1496,13 +1500,11 @@ class RESTfulAPI(CancelMixin):
                 **parsed_kwargs,
             )
             return Response(content=translation, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_speech(
@@ -1558,14 +1560,11 @@ class RESTfulAPI(CancelMixin):
                 )
             else:
                 return Response(media_type="application/octet-stream", content=out)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def get_progress(self, request_id: str) -> JSONResponse:
@@ -1611,14 +1610,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def sdapi_options(self, request: Request) -> Response:
@@ -1689,14 +1685,11 @@ class RESTfulAPI(CancelMixin):
                 **kwargs,
             )
             return Response(content=image_list, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def sdapi_img2img(self, request: Request) -> Response:
@@ -1723,14 +1716,11 @@ class RESTfulAPI(CancelMixin):
                 **kwargs,
             )
             return Response(content=image_list, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_variations(
@@ -1779,13 +1769,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_inpainting(
@@ -1841,13 +1829,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_ocr(
@@ -1887,13 +1873,11 @@ class RESTfulAPI(CancelMixin):
             logger.error(err_str)
             await self._report_error_event(model_uid, err_str)
             raise HTTPException(status_code=409, detail=err_str)
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model_ref.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_flexible_infer(self, request: Request) -> Response:
@@ -1920,14 +1904,11 @@ class RESTfulAPI(CancelMixin):
         try:
             result = await model.infer(**kwargs)
             return Response(result, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_videos(self, request: Request) -> Response:
@@ -1952,14 +1933,11 @@ class RESTfulAPI(CancelMixin):
                 **kwargs,
             )
             return Response(content=video_list, media_type="application/json")
-        except RuntimeError as re:
-            logger.error(re, exc_info=True)
-            await self._report_error_event(model_uid, str(re))
-            self.handle_request_limit_error(re)
-            raise HTTPException(status_code=400, detail=str(re))
         except Exception as e:
+            e = await self._get_model_last_error(model.uid, e)
             logger.error(e, exc_info=True)
             await self._report_error_event(model_uid, str(e))
+            self.handle_request_limit_error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     async def create_chat_completion(self, request: Request) -> Response:
@@ -2084,6 +2062,7 @@ class RESTfulAPI(CancelMixin):
                     # TODO: Cannot yield here. Yield here would leads to error for the next streaming request.
                     return
                 except Exception as ex:
+                    ex = await self._get_model_last_error(model.uid, ex)
                     logger.exception("Chat completion stream got an error: %s", ex)
                     await self._report_error_event(model_uid, str(ex))
                     # https://github.com/openai/openai-python/blob/e0aafc6c1a45334ac889fe3e54957d309c3af93f/src/openai/_streaming.py#L107
@@ -2102,6 +2081,7 @@ class RESTfulAPI(CancelMixin):
                 )
                 return Response(content=data, media_type="application/json")
             except Exception as e:
+                e = await self._get_model_last_error(model.uid, e)
                 logger.error(e, exc_info=True)
                 await self._report_error_event(model_uid, str(e))
                 self.handle_request_limit_error(e)
@@ -2383,6 +2363,10 @@ class RESTfulAPI(CancelMixin):
                 kwargs["guided_whitespace_pattern"] = raw_extra_body.get(
                     "guided_whitespace_pattern"
                 )
+            if raw_extra_body.get("platform") is not None:
+                kwargs["platform"] = raw_extra_body.get("platform")
+            if raw_extra_body.get("format") is not None:
+                kwargs["format"] = raw_extra_body.get("format")
 
         return kwargs
 

xinference/core/model.py

@@ -78,6 +78,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):
@@ -118,20 +121,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)
@@ -178,6 +186,16 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         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"
@@ -206,6 +224,7 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         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
@@ -247,6 +266,11 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         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()
 
@@ -278,6 +302,28 @@ class ModelActor(xo.StatelessActor, CancelMixin):
     def decrease_serve_count(self):
         self._serve_count -= 1
 
+    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
     ):
@@ -440,11 +486,24 @@ class ModelActor(xo.StatelessActor, CancelMixin):
             )
         )
 
+    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
@@ -456,11 +515,11 @@ class ModelActor(xo.StatelessActor, CancelMixin):
                         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(
@@ -482,6 +541,8 @@ class ModelActor(xo.StatelessActor, CancelMixin):
         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
@@ -494,11 +555,8 @@ class ModelActor(xo.StatelessActor, CancelMixin):
                         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:

xinference/core/supervisor.py

@@ -267,6 +267,10 @@ class SupervisorActor(xo.StatelessActor):
                 signal.SIGTERM, lambda: asyncio.create_task(signal_handler())
             )
 
+        from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
+
+        self._block_tracker: Optional[xo.ActorRefType[VLLMBlockTracker]] = None
+
     @typing.no_type_check
     async def get_cluster_device_info(self, detailed: bool = False) -> List:
         import psutil
@@ -956,17 +960,37 @@ class SupervisorActor(xo.StatelessActor):
         ]:
             raise ValueError("Tensorizer is not supported for %s." % model_name)
 
+        enable_xavier: bool = (
+            bool(kwargs.pop("enable_xavier", False))
+            and model_engine is not None
+            and model_engine.lower() == "vllm"
+        )
+        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
+
+                self._block_tracker = await xo.create_actor(
+                    VLLMBlockTracker,
+                    address=self.address,
+                    uid=VLLMBlockTracker.default_uid(),
+                )
+
         if model_uid is None:
             model_uid = self._gen_model_uid(model_name)
 
         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, store_port: 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}"
@@ -974,14 +998,9 @@ class SupervisorActor(xo.StatelessActor):
             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,
@@ -995,14 +1014,57 @@ class SupervisorActor(xo.StatelessActor):
                 gpu_idx=replica_gpu_idx,
                 download_hub=download_hub,
                 model_path=model_path,
+                xavier_config={
+                    "block_tracker_address": self._block_tracker.address
+                    if self._block_tracker is not None
+                    else None,
+                    "rank": rank,
+                    "world_size": replica,
+                    "store_address": self.address.split(":")[0],
+                    "store_port": store_port,
+                }
+                if enable_xavier
+                else None,
                 **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)
+                store_port = xo.utils.get_next_port()
+                worker_refs = []
+                rank_addresses = []
+                for rank, 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()
+                    )
+                    subpool_address = await _launch_one_model(
+                        worker_ref, rep_model_uid, rank, store_port
+                    )
+                    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...")
+                    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)
+                    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)
@@ -1148,6 +1210,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)

xinference/core/utils.py

@@ -276,8 +276,8 @@ def get_nvidia_gpu_info() -> Dict:
 
 
 def assign_replica_gpu(
-    _replica_model_uid: str, replica: int, gpu_idx: Union[int, List[int]]
-) -> List[int]:
+    _replica_model_uid: str, replica: int, gpu_idx: Optional[Union[int, List[int]]]
+) -> Optional[List[int]]:
     model_uid, rep_id = parse_replica_model_uid(_replica_model_uid)
     rep_id, replica = int(rep_id), int(replica)
     if isinstance(gpu_idx, int):