diff --git a/xinference/api/restful_api.py b/xinference/api/restful_api.py index 4d6620b0bc..ed0071ee84 100644 --- a/xinference/api/restful_api.py +++ b/xinference/api/restful_api.py @@ -62,6 +62,7 @@ XINFERENCE_SSE_PING_ATTEMPTS_SECONDS, ) from ..core.event import Event, EventCollectorActor, EventType +from ..core.exceptions import ModelNotReadyError from ..core.supervisor import SupervisorActor from ..core.utils import CancelMixin from ..types import ( @@ -743,6 +744,12 @@ async def describe_model(self, model_uid: str) -> JSONResponse: try: data = await (await self._get_supervisor_ref()).describe_model(model_uid) return JSONResponse(content=data) + except ModelNotReadyError as e: + raise HTTPException( + status_code=503, + detail=f"Model is loading, please retry later: {e}", + headers={"Retry-After": "30"}, + ) except ValueError as ve: logger.error(str(ve), exc_info=True) raise HTTPException(status_code=400, detail=str(ve)) @@ -1676,6 +1683,12 @@ async def sdapi_options(self, request: Request) -> Response: raise ValueError("Unknown model") await (await self._get_supervisor_ref()).get_model(model_uid) return Response() + except ModelNotReadyError as e: + raise HTTPException( + status_code=503, + detail=f"Model is loading, please retry later: {e}", + headers={"Retry-After": "30"}, + ) except ValueError as ve: logger.error(str(ve), exc_info=True) await self._report_error_event(model_uid, str(ve)) @@ -2381,6 +2394,12 @@ async def create_chat_completion(self, request: Request) -> Response: try: desc = await (await self._get_supervisor_ref()).describe_model(model_uid) + except ModelNotReadyError as e: + raise HTTPException( + status_code=503, + detail=f"Model is loading, please retry later: {e}", + headers={"Retry-After": "30"}, + ) except ValueError as ve: logger.error(str(ve), exc_info=True) await self._report_error_event(model_uid, str(ve)) diff --git a/xinference/api/utils.py b/xinference/api/utils.py index 4307d2e8a5..77523280f5 100644 --- a/xinference/api/utils.py +++ b/xinference/api/utils.py @@ -24,6 +24,8 @@ from fastapi import HTTPException +from ..core.exceptions import ModelNotReadyError + logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -135,6 +137,12 @@ async def require_model( return await supervisor.get_model(model_uid) except HTTPException: raise + except ModelNotReadyError as e: + raise HTTPException( + status_code=503, + detail=f"Model is loading, please retry later: {e}", + headers={"Retry-After": "30"}, + ) except ValueError as ve: logger.error(str(ve), exc_info=True) if report_error_event: diff --git a/xinference/core/exceptions.py b/xinference/core/exceptions.py new file mode 100644 index 0000000000..40732cebf7 --- /dev/null +++ b/xinference/core/exceptions.py @@ -0,0 +1,19 @@ +# Copyright 2022-2026 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. + + +class ModelNotReadyError(Exception): + """Raised when a model is still loading and cannot serve requests yet.""" + + pass diff --git a/xinference/core/model.py b/xinference/core/model.py index 79cba41349..2a766bce70 100644 --- a/xinference/core/model.py +++ b/xinference/core/model.py @@ -53,6 +53,7 @@ logger = logging.getLogger(__name__) from ..device_utils import empty_cache +from .exceptions import ModelNotReadyError from .utils import CancelMixin, json_dumps, log_async try: @@ -298,6 +299,7 @@ def __init__( ), } self._loop: Optional[asyncio.AbstractEventLoop] = None + self._model_state: str = "registering" # model across workers self._n_worker = n_worker self._shard = shard @@ -330,6 +332,15 @@ def get_pid(self) -> int: return os.getpid() + def _require_ready(self): + """Guard for all inference methods: reject if not in ready state.""" + if self._model_state != "ready": + if self._model_state in ("registering", "loading"): + raise ModelNotReadyError( + f"Model is {self._model_state}, not ready for inference" + ) + raise RuntimeError(f"Model is in {self._model_state} state") + def __repr__(self) -> str: return f"ModelActor({self._replica_model_uid})" @@ -440,6 +451,7 @@ def is_sglang_backend(self) -> bool: return isinstance(self._model, SGLANGModel) async def load(self): + self._model_state = "loading" try: # Change process title for model import setproctitle @@ -477,6 +489,7 @@ def _load_with_redirect(self): async def wait_for_load(self): if hasattr(self._model, "wait_for_load"): await asyncio.to_thread(self._model.wait_for_load) + self._model_state = "ready" def need_create_pools(self): return getattr(self._model, "need_create_pools", False) @@ -512,6 +525,7 @@ def get_driver_info(self): return self._driver_info async def stop(self): + self._model_state = "stopping" if hasattr(self._model, "stop"): await asyncio.to_thread(self._model.stop) elif hasattr(self._model, "close"): @@ -783,6 +797,7 @@ async def _stream_out_async_gen(): @xo.generator @log_async(logger=logger) async def generate(self, prompt: str, *args, **kwargs): + self._require_ready() # Directly delegate to model, let model decide how to handle (batching or not) kwargs.pop("raw_params", None) if hasattr(self._model, "generate"): @@ -809,6 +824,7 @@ async def generate(self, prompt: str, *args, **kwargs): @xo.generator @log_async(logger=logger) async def chat(self, messages: List[Dict], *args, **kwargs): + self._require_ready() start_time = time.time() response = None try: @@ -884,6 +900,7 @@ async def abort_request( @request_limit @log_async(logger=logger) async def create_embedding(self, input: Union[str, List[str]], *args, **kwargs): + self._require_ready() kwargs.pop("request_id", None) if hasattr(self._model, "create_embedding"): return await self._call_wrapper_json( @@ -920,6 +937,7 @@ async def rerank( *args, **kwargs, ): + self._require_ready() kwargs.pop("request_id", None) if hasattr(self._model, "rerank"): return await self._call_wrapper_json( @@ -947,6 +965,7 @@ async def transcriptions( timestamp_granularities: Optional[List[str]] = None, **kwargs, ): + self._require_ready() kwargs.pop("request_id", None) if hasattr(self._model, "transcriptions"): return await self._call_wrapper_json( @@ -975,6 +994,7 @@ async def translations( timestamp_granularities: Optional[List[str]] = None, **kwargs, ): + self._require_ready() kwargs.pop("request_id", None) if hasattr(self._model, "translations"): return await self._call_wrapper_json( @@ -1002,6 +1022,7 @@ async def speech( stream: bool = False, **kwargs, ): + self._require_ready() kwargs.pop("request_id", None) if hasattr(self._model, "speech"): return await self._call_wrapper_binary( @@ -1028,6 +1049,7 @@ async def text_to_image( *args, **kwargs, ): + self._require_ready() if hasattr(self._model, "text_to_image"): # Get progressor (don't pop request_id, let _call_wrapper handle cancellation) request_id = kwargs.get("request_id") @@ -1052,6 +1074,7 @@ async def txt2img( self, **kwargs, ): + self._require_ready() if hasattr(self._model, "txt2img"): progressor = kwargs["progressor"] = await self._get_progressor( kwargs.pop("request_id", None) @@ -1079,6 +1102,7 @@ async def image_to_image( *args, **kwargs, ): + self._require_ready() kwargs["negative_prompt"] = negative_prompt if hasattr(self._model, "image_to_image"): progressor = kwargs["progressor"] = await self._get_progressor( @@ -1105,6 +1129,7 @@ async def img2img( self, **kwargs, ): + self._require_ready() if hasattr(self._model, "img2img"): progressor = kwargs["progressor"] = await self._get_progressor( kwargs.pop("request_id", None) @@ -1133,6 +1158,7 @@ async def inpainting( *args, **kwargs, ): + self._require_ready() kwargs["negative_prompt"] = negative_prompt if hasattr(self._model, "inpainting"): progressor = kwargs["progressor"] = await self._get_progressor( @@ -1165,6 +1191,7 @@ async def ocr( *args, **kwargs, ): + self._require_ready() if hasattr(self._model, "ocr"): return await self._call_wrapper_json( self._model.ocr, @@ -1181,6 +1208,7 @@ async def infer( *args, **kwargs, ): + self._require_ready() kwargs.pop("request_id", None) if hasattr(self._model, "infer"): return await self._call_wrapper_json( @@ -1201,6 +1229,7 @@ async def text_to_video( *args, **kwargs, ): + self._require_ready() progressor = kwargs["progressor"] = await self._get_progressor( kwargs.pop("request_id", None) ) @@ -1228,6 +1257,7 @@ async def image_to_video( *args, **kwargs, ): + self._require_ready() kwargs["negative_prompt"] = negative_prompt progressor = kwargs["progressor"] = await self._get_progressor( kwargs.pop("request_id", None) @@ -1258,6 +1288,7 @@ async def flf_to_video( *args, **kwargs, ): + self._require_ready() kwargs["negative_prompt"] = negative_prompt progressor = kwargs["progressor"] = await self._get_progressor( kwargs.pop("request_id", None) diff --git a/xinference/core/status_guard.py b/xinference/core/status_guard.py index 3f02e49f71..8e2760c5a5 100644 --- a/xinference/core/status_guard.py +++ b/xinference/core/status_guard.py @@ -29,6 +29,7 @@ class LaunchStatus(Enum): TERMINATED = 4 READY = 5 ERROR = 6 + LOADING = 7 class ReplicaStatus(BaseModel): @@ -37,7 +38,8 @@ class ReplicaStatus(BaseModel): replica_id: int replica_model_uid: str worker_address: str - status: str # CREATING, READY, ERROR + status: str # CREATING, LOADING, READY, ERROR, TERMINATING, TERMINATED + model_state: str = "" # registering/loading/ready/error/stopping/stopped created_ts: int error_message: Optional[str] = None diff --git a/xinference/core/supervisor.py b/xinference/core/supervisor.py index 27e7228ab6..8cd1f69fec 100644 --- a/xinference/core/supervisor.py +++ b/xinference/core/supervisor.py @@ -56,6 +56,7 @@ get_engine_params_by_name_with_virtual_env, ) from ..types import PeftModelConfig +from .exceptions import ModelNotReadyError from .launch_strategy import IdleFirstLaunchStrategy from .metrics import record_metrics from .resource import GPUStatus, ResourceStatus @@ -1748,7 +1749,6 @@ async def _launch_one_model( 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 @@ -1756,6 +1756,7 @@ async def _launch_one_model( await self._status_guard_ref.update_replica_status( model_uid, replica_id, {"status": LaunchStatus.READY.name} ) + self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref return rank0_address replica_gpu_idx = ( @@ -1794,8 +1795,10 @@ async def _launch_one_model( xavier_config=xavier_config, **kwargs, ) - self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref + # Wait for engine to be ready BEFORE adding to route table, + # so requests are never routed to a still-loading model. await worker_ref.wait_for_load(_replica_model_uid) + self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref # Update replica status to READY await self._status_guard_ref.update_replica_status( @@ -2711,8 +2714,13 @@ async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]: 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}" + # replica_info exists but worker route not yet registered — + # model is being launched (wait_ready=False), not missing. + # Raise ModelNotReadyError (503) instead of ValueError (404) + # so the negative cache is not poisoned and clients get a + # retry-friendly response. + raise ModelNotReadyError( + f"Model {model_uid} is launching, not yet ready for inference" ) if isinstance(worker_ref, (list, tuple)): # get first worker to fetch information if model across workers @@ -2720,10 +2728,15 @@ async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]: assert not isinstance( worker_ref, (list, tuple) ), "worker_ref must be a single worker" - return await xo.wait_for( - worker_ref.get_model(model_uid=replica_model_uid), - XINFERENCE_GET_MODEL_RPC_TIMEOUT, - ) + try: + return await xo.wait_for( + worker_ref.get_model(model_uid=replica_model_uid), + XINFERENCE_GET_MODEL_RPC_TIMEOUT, + ) + except ModelNotReadyError: + raise ModelNotReadyError( + f"Model {model_uid} is loading, please retry later" + ) @log_async(logger=logger) async def get_model_status(self, replica_model_uid: str): @@ -2753,8 +2766,8 @@ async def describe_model(self, model_uid: str) -> Dict[str, Any]: ) 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}" + raise ModelNotReadyError( + f"Model {model_uid} is launching, not yet ready for inference" ) if isinstance(worker_ref, (list, tuple)): # get status from first shard if model has multiple shards across workers diff --git a/xinference/core/tests/test_model.py b/xinference/core/tests/test_model.py index e849b6140f..eb7ff5f535 100644 --- a/xinference/core/tests/test_model.py +++ b/xinference/core/tests/test_model.py @@ -56,6 +56,10 @@ def __init__( model=MockModel(), # type: ignore replica_model_uid=replica_model_uid, ) + # This actor is constructed directly with a ready-to-use MockModel + # (no launch/load path), so mark it ready to bypass the + # _require_ready() guard added by the model state machine. + self._model_state = "ready" self._lock = asyncio.locks.Lock() async def __pre_destroy__(self): diff --git a/xinference/core/worker.py b/xinference/core/worker.py index 523aeb36fd..baf21540f4 100644 --- a/xinference/core/worker.py +++ b/xinference/core/worker.py @@ -78,6 +78,7 @@ from ..utils import get_pip_config_args, get_real_path from .cache_tracker import CacheTrackerActor from .event import Event, EventCollectorActor, EventType +from .exceptions import ModelNotReadyError from .metrics import ( launch_metrics_export_server, record_metrics, @@ -158,6 +159,7 @@ def _exclusive_venv_path_lock(env_path: str): @dataclass class ModelStatus: last_error: str = "" + model_state: str = "" @dataclass @@ -2283,6 +2285,7 @@ def check_cancel(): raise except Exception: logger.error(f"Failed to load model {model_uid}", exc_info=True) + await self._update_model_state(model_uid, "error") self.release_devices(model_uid=model_uid) for addr in all_subpool_addresses: try: @@ -2310,7 +2313,9 @@ def check_cancel(): self._model_uid_to_subpool_pids[model_uid] = subpool_pids model_spec = model.model_family.to_description() self._model_uid_to_model_spec[model_uid] = model_spec - self._model_uid_to_model_status[model_uid] = ModelStatus() + self._model_uid_to_model_status[model_uid] = ModelStatus( + model_state="loading" + ) self._model_uid_to_addr[model_uid] = subpool_address self._model_uid_to_recover_count.setdefault( model_uid, MODEL_ACTOR_AUTO_RECOVER_LIMIT @@ -2373,6 +2378,7 @@ def check_cancel(): async def wait_for_load(self, model_uid: str): model_ref = self._model_uid_to_model[model_uid] await model_ref.wait_for_load() + await self._update_model_state(model_uid, "ready") @log_sync(logger=logger, level=logging.INFO) async def cancel_launch_model(self, model_uid: str): @@ -2412,6 +2418,7 @@ async def terminate_model(self, model_uid: str, is_model_die=False): # Terminate model while its launching is not allow if model_uid in self._model_uid_launching_guard: raise ValueError(f"{model_uid} is launching") + await self._update_model_state(model_uid, "stopping") # In special cases, if the suffix is `-rank0`, this is the Xavier's rank 0 model actor. if model_uid.endswith("-rank0"): origin_uid = model_uid.removesuffix("-rank0") @@ -2509,6 +2516,7 @@ async def terminate_model(self, model_uid: str, is_model_die=False): status = LaunchStatus.ERROR.name else: status = LaunchStatus.TERMINATED.name + await self._update_model_state(model_uid, "stopped") self._model_uid_to_model_status.pop(model_uid, None) if self._status_guard_ref is None: @@ -2537,11 +2545,54 @@ def get_model_launch_status(self, model_uid: str) -> Optional[str]: async def list_models(self) -> Dict[str, Dict[str, Any]]: return {k: v for k, v in self._model_uid_to_model_spec.items()} - @log_sync(logger=logger) + @log_async(logger=logger) + async def _update_model_state(self, model_uid: str, state: str): + """Update ModelStatus.model_state and sync to StatusGuard.""" + ms = self._model_uid_to_model_status.get(model_uid) + if ms is None: + ms = ModelStatus() + self._model_uid_to_model_status[model_uid] = ms + ms.model_state = state + status_map = { + "registering": LaunchStatus.CREATING.name, + "loading": LaunchStatus.LOADING.name, + "ready": LaunchStatus.READY.name, + "error": LaunchStatus.ERROR.name, + "stopping": LaunchStatus.TERMINATING.name, + "stopped": LaunchStatus.TERMINATED.name, + } + if self._status_guard_ref is not None: + try: + origin_uid, rank_suffix = parse_replica_model_uid(model_uid) + replica_id = rank_suffix - 1 if rank_suffix > 0 else 0 + await self._status_guard_ref.update_replica_status( + origin_uid, + replica_id, + { + "status": status_map.get(state, LaunchStatus.CREATING.name), + "model_state": state, + }, + ) + except Exception: + logger.debug( + "Failed to sync model_state to StatusGuard for %s", + model_uid, + exc_info=True, + ) + def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]: model_status = self._model_uid_to_model_status.get(model_uid) - if model_status and model_status.last_error: - raise Exception(model_status.last_error) + if model_status: + if model_status.model_state in ("registering", "loading"): + raise ModelNotReadyError( + f"Model {model_uid} is {model_status.model_state}" + ) + if model_status.model_state in ("error", "stopping", "stopped"): + raise RuntimeError( + f"Model {model_uid} is in {model_status.model_state} state" + ) + if model_status.last_error: + raise Exception(model_status.last_error) model_ref = self._model_uid_to_model.get(model_uid, None) if model_ref is None: raise ValueError(f"Model not found, uid: {model_uid}")