xinference 0.10.2.post1__py3-none-any.whl → 0.11.0__py3-none-any.whl

xinference/core/supervisor.py

@@ -80,12 +80,12 @@ class ReplicaInfo:
 class SupervisorActor(xo.StatelessActor):
     def __init__(self):
         super().__init__()
-        self._worker_address_to_worker: Dict[str, xo.ActorRefType["WorkerActor"]] = {}
-        self._worker_status: Dict[str, WorkerStatus] = {}
-        self._replica_model_uid_to_worker: Dict[
+        self._worker_address_to_worker: Dict[str, xo.ActorRefType["WorkerActor"]] = {}  # type: ignore
+        self._worker_status: Dict[str, WorkerStatus] = {}  # type: ignore
+        self._replica_model_uid_to_worker: Dict[  # type: ignore
             str, xo.ActorRefType["WorkerActor"]
         ] = {}
-        self._model_uid_to_replica_info: Dict[str, ReplicaInfo] = {}
+        self._model_uid_to_replica_info: Dict[str, ReplicaInfo] = {}  # type: ignore
         self._uptime = None
         self._lock = asyncio.Lock()
 
@@ -117,12 +117,12 @@ class SupervisorActor(xo.StatelessActor):
         from .cache_tracker import CacheTrackerActor
         from .status_guard import StatusGuardActor
 
-        self._status_guard_ref: xo.ActorRefType[
+        self._status_guard_ref: xo.ActorRefType[  # type: ignore
             "StatusGuardActor"
         ] = await xo.create_actor(
             StatusGuardActor, address=self.address, uid=StatusGuardActor.uid()
         )
-        self._cache_tracker_ref: xo.ActorRefType[
+        self._cache_tracker_ref: xo.ActorRefType[  # type: ignore
             "CacheTrackerActor"
         ] = await xo.create_actor(
             CacheTrackerActor, address=self.address, uid=CacheTrackerActor.uid()
@@ -130,7 +130,7 @@ class SupervisorActor(xo.StatelessActor):
 
         from .event import EventCollectorActor
 
-        self._event_collector_ref: xo.ActorRefType[
+        self._event_collector_ref: xo.ActorRefType[  # type: ignore
             EventCollectorActor
         ] = await xo.create_actor(
             EventCollectorActor, address=self.address, uid=EventCollectorActor.uid()
@@ -150,7 +150,13 @@ class SupervisorActor(xo.StatelessActor):
             register_embedding,
             unregister_embedding,
         )
-        from ..model.image import get_image_model_descriptions
+        from ..model.image import (
+            CustomImageModelFamilyV1,
+            generate_image_description,
+            get_image_model_descriptions,
+            register_image,
+            unregister_image,
+        )
         from ..model.llm import (
             CustomLLMFamilyV1,
             generate_llm_description,
@@ -166,7 +172,7 @@ class SupervisorActor(xo.StatelessActor):
             unregister_rerank,
         )
 
-        self._custom_register_type_to_cls: Dict[str, Tuple] = {
+        self._custom_register_type_to_cls: Dict[str, Tuple] = {  # type: ignore
             "LLM": (
                 CustomLLMFamilyV1,
                 register_llm,
@@ -185,6 +191,12 @@ class SupervisorActor(xo.StatelessActor):
                 unregister_rerank,
                 generate_rerank_description,
             ),
+            "image": (
+                CustomImageModelFamilyV1,
+                register_image,
+                unregister_image,
+                generate_image_description,
+            ),
             "audio": (
                 CustomAudioModelFamilyV1,
                 register_audio,
@@ -194,7 +206,7 @@ class SupervisorActor(xo.StatelessActor):
         }
 
         # record model version
-        model_version_infos: Dict[str, List[Dict]] = {}
+        model_version_infos: Dict[str, List[Dict]] = {}  # type: ignore
         model_version_infos.update(get_llm_model_descriptions())
         model_version_infos.update(get_embedding_model_descriptions())
         model_version_infos.update(get_rerank_model_descriptions())
@@ -272,7 +284,7 @@ class SupervisorActor(xo.StatelessActor):
         return {
             "chat": list(BUILTIN_LLM_MODEL_CHAT_FAMILIES),
             "generate": list(BUILTIN_LLM_MODEL_GENERATE_FAMILIES),
-            "tool_call": list(BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES),
+            "tools": list(BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES),
         }
 
     async def get_devices_count(self) -> int:
@@ -486,6 +498,7 @@ class SupervisorActor(xo.StatelessActor):
             return ret
         elif model_type == "image":
             from ..model.image import BUILTIN_IMAGE_MODELS
+            from ..model.image.custom import get_user_defined_images
 
             ret = []
             for model_name, family in BUILTIN_IMAGE_MODELS.items():
@@ -494,6 +507,16 @@ class SupervisorActor(xo.StatelessActor):
                 else:
                     ret.append({"model_name": model_name, "is_builtin": True})
 
+            for model_spec in get_user_defined_images():
+                if detailed:
+                    ret.append(
+                        await self._to_image_model_reg(model_spec, is_builtin=False)
+                    )
+                else:
+                    ret.append(
+                        {"model_name": model_spec.model_name, "is_builtin": False}
+                    )
+
             ret.sort(key=sort_helper)
             return ret
         elif model_type == "audio":
@@ -567,8 +590,9 @@ class SupervisorActor(xo.StatelessActor):
             raise ValueError(f"Model {model_name} not found")
         elif model_type == "image":
             from ..model.image import BUILTIN_IMAGE_MODELS
+            from ..model.image.custom import get_user_defined_images
 
-            for f in BUILTIN_IMAGE_MODELS.values():
+            for f in list(BUILTIN_IMAGE_MODELS.values()) + get_user_defined_images():
                 if f.model_name == model_name:
                     return f
             raise ValueError(f"Model {model_name} not found")
@@ -591,6 +615,24 @@ class SupervisorActor(xo.StatelessActor):
         else:
             raise ValueError(f"Unsupported model type: {model_type}")
 
+    @log_async(logger=logger)
+    async def query_engines_by_model_name(self, model_name: str):
+        from copy import deepcopy
+
+        from ..model.llm.llm_family import LLM_ENGINES
+
+        if model_name not in LLM_ENGINES:
+            raise ValueError(f"Model {model_name} not found")
+
+        # filter llm_class
+        engine_params = deepcopy(LLM_ENGINES[model_name])
+        for engine in engine_params:
+            params = engine_params[engine]
+            for param in params:
+                del param["llm_class"]
+
+        return engine_params
+
     @log_async(logger=logger)
     async def register_model(self, model_type: str, model: str, persist: bool):
         if model_type in self._custom_register_type_to_cls:
@@ -651,6 +693,7 @@ class SupervisorActor(xo.StatelessActor):
         self,
         model_uid: Optional[str],
         model_type: str,
+        model_engine: Optional[str],
         model_version: str,
         replica: int = 1,
         n_gpu: Optional[Union[int, str]] = "auto",
@@ -666,6 +709,7 @@ class SupervisorActor(xo.StatelessActor):
         return await self.launch_builtin_model(
             model_uid=model_uid,
             model_name=parse_results[0],
+            model_engine=model_engine,
             model_size_in_billions=parse_results[1] if model_type == "LLM" else None,
             model_format=parse_results[2] if model_type == "LLM" else None,
             quantization=parse_results[3] if model_type == "LLM" else None,
@@ -677,66 +721,6 @@ class SupervisorActor(xo.StatelessActor):
             **kwargs,
         )
 
-    async def launch_speculative_llm(
-        self,
-        model_uid: Optional[str],
-        model_name: str,
-        model_size_in_billions: Optional[Union[int, str]],
-        quantization: Optional[str],
-        draft_model_name: str,
-        draft_model_size_in_billions: Optional[int],
-        draft_quantization: Optional[str],
-        n_gpu: Optional[Union[int, str]] = "auto",
-    ) -> str:
-        if model_uid is None:
-            model_uid = self._gen_model_uid(model_name)
-        logger.debug(
-            (
-                f"Enter launch_speculative_llm, model_uid: %s, model_name: %s, model_size: %s, "
-                f"draft_model_name: %s, draft_model_size: %s"
-            ),
-            model_uid,
-            model_name,
-            str(model_size_in_billions) if model_size_in_billions else "",
-            draft_model_name,
-            draft_model_size_in_billions,
-        )
-
-        # TODO: the draft and target model must be on the same worker.
-        if not self.is_local_deployment():
-            raise ValueError(
-                "Speculative model is not supported in distributed deployment yet."
-            )
-
-        if model_uid in self._model_uid_to_replica_info:
-            raise ValueError(f"Model is already in the model list, uid: {model_uid}")
-
-        worker_ref = await self._choose_worker()
-        replica = 1
-        self._model_uid_to_replica_info[model_uid] = ReplicaInfo(
-            replica=replica, scheduler=itertools.cycle(range(replica))
-        )
-
-        try:
-            rep_model_uid = f"{model_uid}-{1}-{0}"
-            await worker_ref.launch_speculative_model(
-                model_uid=rep_model_uid,
-                model_name=model_name,
-                model_size_in_billions=model_size_in_billions,
-                quantization=quantization,
-                draft_model_name=draft_model_name,
-                draft_model_size_in_billions=draft_model_size_in_billions,
-                draft_quantization=draft_quantization,
-                n_gpu=n_gpu,
-            )
-            self._replica_model_uid_to_worker[rep_model_uid] = worker_ref
-
-        except Exception:
-            # terminate_model will remove the replica info.
-            await self.terminate_model(model_uid, suppress_exception=True)
-            raise
-        return model_uid
-
     async def launch_builtin_model(
         self,
         model_uid: Optional[str],
@@ -744,6 +728,7 @@ class SupervisorActor(xo.StatelessActor):
         model_size_in_billions: Optional[Union[int, str]],
         model_format: Optional[str],
         quantization: Optional[str],
+        model_engine: Optional[str],
         model_type: Optional[str],
         replica: int = 1,
         n_gpu: Optional[Union[int, str]] = "auto",
@@ -799,6 +784,7 @@ class SupervisorActor(xo.StatelessActor):
                 model_size_in_billions=model_size_in_billions,
                 model_format=model_format,
                 quantization=quantization,
+                model_engine=model_engine,
                 model_type=model_type,
                 n_gpu=n_gpu,
                 request_limits=request_limits,

xinference/core/worker.py

@@ -34,7 +34,7 @@ from ..constants import (
 )
 from ..core import ModelActor
 from ..core.status_guard import LaunchStatus
-from ..device_utils import gpu_count
+from ..device_utils import get_available_device_env_name, gpu_count
 from ..model.core import ModelDescription, create_model_instance
 from ..types import PeftModelConfig
 from .event import Event, EventCollectorActor, EventType
@@ -80,7 +80,7 @@ class WorkerActor(xo.StatelessActor):
             int, Set[Tuple[str, str]]
         ] = defaultdict(set)
         self._model_uid_to_addr: Dict[str, str] = {}
-        self._model_uid_to_recover_count: Dict[str, int] = {}
+        self._model_uid_to_recover_count: Dict[str, Optional[int]] = {}
         self._model_uid_to_launch_args: Dict[str, Dict] = {}
 
         # metrics export server.
@@ -137,14 +137,19 @@ class WorkerActor(xo.StatelessActor):
                                 recover_count - 1,
                             )
                             event_model_uid, _, __ = parse_replica_model_uid(model_uid)
-                            await self._event_collector_ref.report_event(
-                                event_model_uid,
-                                Event(
-                                    event_type=EventType.WARNING,
-                                    event_ts=int(time.time()),
-                                    event_content="Recreate model",
-                                ),
-                            )
+                            try:
+                                await self._event_collector_ref.report_event(
+                                    event_model_uid,
+                                    Event(
+                                        event_type=EventType.WARNING,
+                                        event_ts=int(time.time()),
+                                        event_content="Recreate model",
+                                    ),
+                                )
+                            except Exception as e:
+                                # Report callback error can be log and ignore, should not interrupt the Process
+                                logger.error("report_event error: %s" % (e))
+
                             self._model_uid_to_recover_count[model_uid] = (
                                 recover_count - 1
                             )
@@ -166,22 +171,22 @@ class WorkerActor(xo.StatelessActor):
         from .status_guard import StatusGuardActor
         from .supervisor import SupervisorActor
 
-        self._status_guard_ref: xo.ActorRefType[
+        self._status_guard_ref: xo.ActorRefType[  # type: ignore
             "StatusGuardActor"
         ] = await xo.actor_ref(
             address=self._supervisor_address, uid=StatusGuardActor.uid()
         )
-        self._event_collector_ref: xo.ActorRefType[
+        self._event_collector_ref: xo.ActorRefType[  # type: ignore
             EventCollectorActor
         ] = await xo.actor_ref(
             address=self._supervisor_address, uid=EventCollectorActor.uid()
         )
-        self._cache_tracker_ref: xo.ActorRefType[
+        self._cache_tracker_ref: xo.ActorRefType[  # type: ignore
             "CacheTrackerActor"
         ] = await xo.actor_ref(
             address=self._supervisor_address, uid=CacheTrackerActor.uid()
         )
-        self._supervisor_ref: xo.ActorRefType["SupervisorActor"] = await xo.actor_ref(
+        self._supervisor_ref: xo.ActorRefType["SupervisorActor"] = await xo.actor_ref(  # type: ignore
             address=self._supervisor_address, uid=SupervisorActor.uid()
         )
         await self._supervisor_ref.add_worker(self.address)
@@ -208,7 +213,12 @@ class WorkerActor(xo.StatelessActor):
             register_embedding,
             unregister_embedding,
         )
-        from ..model.image import get_image_model_descriptions
+        from ..model.image import (
+            CustomImageModelFamilyV1,
+            get_image_model_descriptions,
+            register_image,
+            unregister_image,
+        )
         from ..model.llm import (
             CustomLLMFamilyV1,
             get_llm_model_descriptions,
@@ -222,7 +232,7 @@ class WorkerActor(xo.StatelessActor):
             unregister_rerank,
         )
 
-        self._custom_register_type_to_cls: Dict[str, Tuple] = {
+        self._custom_register_type_to_cls: Dict[str, Tuple] = {  # type: ignore
             "LLM": (CustomLLMFamilyV1, register_llm, unregister_llm),
             "embedding": (
                 CustomEmbeddingModelSpec,
@@ -231,10 +241,15 @@ class WorkerActor(xo.StatelessActor):
             ),
             "rerank": (CustomRerankModelSpec, register_rerank, unregister_rerank),
             "audio": (CustomAudioModelFamilyV1, register_audio, unregister_audio),
+            "image": (
+                CustomImageModelFamilyV1,
+                register_image,
+                unregister_image,
+            ),
         }
 
         # record model version
-        model_version_infos: Dict[str, List[Dict]] = {}
+        model_version_infos: Dict[str, List[Dict]] = {}  # type: ignore
         model_version_infos.update(get_llm_model_descriptions())
         model_version_infos.update(get_embedding_model_descriptions())
         model_version_infos.update(get_rerank_model_descriptions())
@@ -248,7 +263,11 @@ class WorkerActor(xo.StatelessActor):
         if os.name != "nt":
 
             async def signal_handler():
-                await self._supervisor_ref.remove_worker(self.address)
+                try:
+                    await self._supervisor_ref.remove_worker(self.address)
+                except Exception as e:
+                    # Ignore the error of rpc, anyway we are exiting
+                    logger.exception("remove worker rpc error: %s", e)
                 os._exit(0)
 
             loop = asyncio.get_running_loop()
@@ -437,6 +456,7 @@ class WorkerActor(xo.StatelessActor):
     ) -> Tuple[str, List[str]]:
         env = {}
         devices = []
+        env_name = get_available_device_env_name()
         if gpu_idx is None:
             if isinstance(n_gpu, int) or (n_gpu == "auto" and gpu_count() > 0):
                 # Currently, n_gpu=auto means using 1 GPU
@@ -446,17 +466,17 @@ class WorkerActor(xo.StatelessActor):
                     if model_type in ["embedding", "rerank"]
                     else self.allocate_devices(model_uid=model_uid, n_gpu=gpu_cnt)
                 )
-                env["CUDA_VISIBLE_DEVICES"] = ",".join([str(dev) for dev in devices])
+                env[env_name] = ",".join([str(dev) for dev in devices])
                 logger.debug(f"GPU selected: {devices} for model {model_uid}")
             if n_gpu is None:
-                env["CUDA_VISIBLE_DEVICES"] = "-1"
+                env[env_name] = "-1"
                 logger.debug(f"GPU disabled for model {model_uid}")
         else:
             assert isinstance(gpu_idx, list)
             devices = await self.allocate_devices_with_gpu_idx(
                 model_uid, model_type, gpu_idx  # type: ignore
             )
-            env["CUDA_VISIBLE_DEVICES"] = ",".join([str(dev) for dev in devices])
+            env[env_name] = ",".join([str(dev) for dev in devices])
 
         if os.name != "nt" and platform.system() != "Darwin":
             # Linux
@@ -503,67 +523,6 @@ class WorkerActor(xo.StatelessActor):
         else:
             raise ValueError(f"Unsupported model type: {model_type}")
 
-    @log_async(logger=logger)
-    async def launch_speculative_model(
-        self,
-        model_uid: str,
-        model_name: str,
-        model_size_in_billions: Optional[int],
-        quantization: Optional[str],
-        draft_model_name: str,
-        draft_model_size_in_billions: Optional[int],
-        draft_quantization: Optional[str],
-        n_gpu: Optional[Union[int, str]] = "auto",
-    ):
-        if n_gpu is not None:
-            if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > gpu_count()):
-                raise ValueError(
-                    f"The parameter `n_gpu` must be greater than 0 and "
-                    f"not greater than the number of GPUs: {gpu_count()} on the machine."
-                )
-            if isinstance(n_gpu, str) and n_gpu != "auto":
-                raise ValueError("Currently `n_gpu` only supports `auto`.")
-
-        from ..model.llm.core import create_speculative_llm_model_instance
-
-        subpool_address, devices = await self._create_subpool(model_uid, n_gpu=n_gpu)
-
-        model, model_description = await asyncio.to_thread(
-            create_speculative_llm_model_instance,
-            subpool_addr=subpool_address,
-            devices=devices,
-            model_uid=model_uid,
-            model_name=model_name,
-            model_size_in_billions=model_size_in_billions,
-            quantization=quantization,
-            draft_model_name=draft_model_name,
-            draft_model_size_in_billions=draft_model_size_in_billions,
-            draft_quantization=draft_quantization,
-            is_local_deployment=True,
-        )
-
-        try:
-            model_ref = await xo.create_actor(
-                ModelActor,
-                address=subpool_address,
-                uid=model_uid,
-                worker_address=self.address,
-                model=model,
-                model_description=model_description,
-            )
-            await model_ref.load()
-        except:
-            logger.error(f"Failed to load model {model_uid}", exc_info=True)
-            self.release_devices(model_uid=model_uid)
-            await self._main_pool.remove_sub_pool(subpool_address)
-            raise
-
-        self._model_uid_to_model[model_uid] = model_ref
-        self._model_uid_to_model_spec[model_uid] = model_description
-        for dev in devices:
-            self._gpu_to_model_uid[int(dev)] = model_uid
-        self._model_uid_to_addr[model_uid] = subpool_address
-
     async def _get_model_ability(self, model: Any, model_type: str) -> List[str]:
         from ..model.llm.core import LLM
 
@@ -605,6 +564,7 @@ class WorkerActor(xo.StatelessActor):
         model_size_in_billions: Optional[Union[int, str]],
         model_format: Optional[str],
         quantization: Optional[str],
+        model_engine: Optional[str],
         model_type: str = "LLM",
         n_gpu: Optional[Union[int, str]] = "auto",
         peft_model_config: Optional[PeftModelConfig] = None,
@@ -612,20 +572,28 @@ class WorkerActor(xo.StatelessActor):
         gpu_idx: Optional[Union[int, List[int]]] = None,
         **kwargs,
     ):
-        event_model_uid, _, __ = parse_replica_model_uid(model_uid)
-        await self._event_collector_ref.report_event(
-            event_model_uid,
-            Event(
-                event_type=EventType.INFO,
-                event_ts=int(time.time()),
-                event_content="Launch model",
-            ),
-        )
+        # !!! Note that The following code must be placed at the very beginning of this function,
+        # or there will be problems with auto-recovery.
+        # Because `locals()` will collect all the local parameters of this function and pass to this function again.
         launch_args = locals()
         launch_args.pop("self")
         launch_args.pop("kwargs")
         launch_args.update(kwargs)
 
+        event_model_uid, _, __ = parse_replica_model_uid(model_uid)
+        try:
+            await self._event_collector_ref.report_event(
+                event_model_uid,
+                Event(
+                    event_type=EventType.INFO,
+                    event_ts=int(time.time()),
+                    event_content="Launch model",
+                ),
+            )
+        except Exception as e:
+            # Report callback error can be log and ignore, should not interrupt the Process
+            logger.error("report_event error: %s" % (e))
+
         if gpu_idx is not None:
             logger.info(
                 f"You specify to launch the model: {model_name} on GPU index: {gpu_idx} "
@@ -657,8 +625,6 @@ class WorkerActor(xo.StatelessActor):
 
         assert model_uid not in self._model_uid_to_model
         self._check_model_is_valid(model_name, model_format)
-        assert self._supervisor_ref is not None
-        is_local_deployment = await self._supervisor_ref.is_local_deployment()
 
         subpool_address, devices = await self._create_subpool(
             model_uid, model_type, n_gpu=n_gpu, gpu_idx=gpu_idx
@@ -673,11 +639,11 @@ class WorkerActor(xo.StatelessActor):
                 model_uid,
                 model_type,
                 model_name,
+                model_engine,
                 model_format,
                 model_size_in_billions,
                 quantization,
                 peft_model_config,
-                is_local_deployment,
                 **kwargs,
             )
             await self.update_cache_status(model_name, model_description)
@@ -715,14 +681,19 @@ class WorkerActor(xo.StatelessActor):
     @log_async(logger=logger)
     async def terminate_model(self, model_uid: str):
         event_model_uid, _, __ = parse_replica_model_uid(model_uid)
-        await self._event_collector_ref.report_event(
-            event_model_uid,
-            Event(
-                event_type=EventType.INFO,
-                event_ts=int(time.time()),
-                event_content="Terminate model",
-            ),
-        )
+        try:
+            await self._event_collector_ref.report_event(
+                event_model_uid,
+                Event(
+                    event_type=EventType.INFO,
+                    event_ts=int(time.time()),
+                    event_content="Terminate model",
+                ),
+            )
+        except Exception as e:
+            # Report callback error can be log and ignore, should not interrupt the Process
+            logger.error("report_event error: %s" % (e))
+
         origin_uid, _, _ = parse_replica_model_uid(model_uid)
         await self._status_guard_ref.update_instance_info(
             origin_uid, {"status": LaunchStatus.TERMINATING.name}