xinference 0.6.4__py3-none-any.whl → 0.7.0__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-11-24T12:37:09+0800",
+ "date": "2023-12-08T13:45:18+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "8fd2e3bf97418bc99f65087c34c932e0318969bc",
- "version": "0.6.4"
+ "full-revisionid": "b5a5f0a270f85e451591eba34fe615a0fc8ce4bf",
+ "version": "0.7.0"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -21,7 +21,7 @@ import os
 import pprint
 import sys
 import warnings
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, List, Optional, Union
 
 import gradio as gr
 import xoscar as xo
@@ -43,28 +43,19 @@ from pydantic import BaseModel, Field
 from sse_starlette.sse import EventSourceResponse
 from starlette.responses import JSONResponse as StarletteJSONResponse
 from starlette.responses import RedirectResponse
-from typing_extensions import NotRequired, TypedDict
 from uvicorn import Config, Server
 from xoscar.utils import get_next_port
 
 from ..constants import XINFERENCE_DEFAULT_ENDPOINT_PORT
 from ..core.supervisor import SupervisorActor
 from ..core.utils import json_dumps
-from ..fields import (
-    frequency_penalty_field,
-    max_tokens_field,
-    mirostat_eta_field,
-    mirostat_mode_field,
-    mirostat_tau_field,
-    presence_penalty_field,
-    repeat_penalty_field,
-    stop_field,
-    stream_field,
-    temperature_field,
-    top_k_field,
-    top_p_field,
+from ..types import (
+    ChatCompletion,
+    Completion,
+    CreateChatCompletion,
+    CreateCompletion,
+    ImageList,
 )
-from ..types import ChatCompletion, Completion, CreateCompletion, ImageList
 
 logger = logging.getLogger(__name__)
 
@@ -115,50 +106,6 @@ class TextToImageRequest(BaseModel):
     user: Optional[str] = None
 
 
-class ChatCompletionRequestMessage(TypedDict):
-    role: Literal["assistant", "user", "system"]
-    content: str
-    user: NotRequired[str]
-
-
-class CreateChatCompletionRequest(BaseModel):
-    messages: List[ChatCompletionRequestMessage] = Field(
-        default=[], description="A list of messages to generate completions for."
-    )
-    max_tokens: int = max_tokens_field
-    temperature: float = temperature_field
-    top_p: float = top_p_field
-    mirostat_mode: int = mirostat_mode_field
-    mirostat_tau: float = mirostat_tau_field
-    mirostat_eta: float = mirostat_eta_field
-    stop: Optional[Union[str, List[str]]] = stop_field
-    stream: bool = stream_field
-    presence_penalty: Optional[float] = presence_penalty_field
-    frequency_penalty: Optional[float] = frequency_penalty_field
-    logit_bias: Optional[Dict[str, float]] = Field(None)
-
-    model: str
-    n: Optional[int] = 1
-    user: Optional[str] = Field(None)
-
-    # llama.cpp specific parameters
-    top_k: int = top_k_field
-    repeat_penalty: Optional[float] = repeat_penalty_field
-    logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None)
-    grammar: Optional[str] = Field(None)
-
-    class Config:
-        schema_extra = {
-            "example": {
-                "messages": [
-                    {"role": "system", "content": "you are a helpful AI assistant"},
-                    {"role": "user", "content": "Hello!"},
-                    {"role": "assistant", "content": "Hi what can I help you?"},
-                ]
-            }
-        }
-
-
 class RegisterModelRequest(BaseModel):
     model: str
     persist: bool
@@ -208,6 +155,12 @@ class RESTfulAPI:
         )
         self._router.add_api_route("/status", self.get_status, methods=["GET"])
         self._router.add_api_route("/v1/models", self.list_models, methods=["GET"])
+        self._router.add_api_route(
+            "/v1/models/prompts", self._get_builtin_prompts, methods=["GET"]
+        )
+        self._router.add_api_route(
+            "/v1/cluster/devices", self._get_devices_count, methods=["GET"]
+        )
         self._router.add_api_route(
             "/v1/models/{model_uid}", self.describe_model, methods=["GET"]
         )
@@ -305,6 +258,9 @@ class RESTfulAPI:
                 f"{pprint.pformat(invalid_routes)}"
             )
 
+        for tp in [CreateChatCompletion, CreateCompletion]:
+            logger.debug("Dump request model fields:\n%s", tp.__fields__)
+
         class SPAStaticFiles(StaticFiles):
             async def get_response(self, path: str, scope):
                 response = await super().get_response(path, scope)
@@ -346,6 +302,28 @@ class RESTfulAPI:
         server = Server(config)
         server.run()
 
+    async def _get_builtin_prompts(self) -> JSONResponse:
+        """
+        For internal usage
+        """
+        try:
+            data = await (await self._get_supervisor_ref()).get_builtin_prompts()
+            return JSONResponse(content=data)
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
+    async def _get_devices_count(self) -> JSONResponse:
+        """
+        For internal usage
+        """
+        try:
+            data = await (await self._get_supervisor_ref()).get_devices_count()
+            return JSONResponse(content=data)
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def get_status(self) -> JSONResponse:
         try:
             data = await (await self._get_supervisor_ref()).get_status()
@@ -725,7 +703,7 @@ class RESTfulAPI:
     async def create_chat_completion(
         self,
         request: Request,
-        body: CreateChatCompletionRequest,
+        body: CreateChatCompletion,
     ) -> Response:
         exclude = {
             "prompt",
@@ -736,7 +714,7 @@ class RESTfulAPI:
             "logit_bias_type",
             "user",
         }
-        kwargs = body.dict(exclude=exclude)
+        kwargs = body.dict(exclude_unset=True, exclude=exclude)
 
         if body.logit_bias is not None:
             raise HTTPException(status_code=501, detail="Not implemented")
@@ -795,6 +773,7 @@ class RESTfulAPI:
         is_chatglm_ggml = desc.get(
             "model_format"
         ) == "ggmlv3" and "chatglm" in desc.get("model_name", "")
+        is_chatglm3 = "chatglm3" == desc.get("model_name", "")
 
         is_qwen = desc.get("model_format") == "ggmlv3" and "qwen" in desc.get(
             "model_name", ""
@@ -804,6 +783,14 @@ class RESTfulAPI:
             raise HTTPException(
                 status_code=400, detail="ChatGLM ggml does not have system prompt"
             )
+        if is_chatglm3 and body.tools and body.stream:
+            raise HTTPException(
+                status_code=400, detail="ChatGLM3 tool calls does not support stream"
+            )
+        if body.tools and not is_chatglm3:
+            raise HTTPException(
+                status_code=400, detail="Only ChatGLM3 support tool calls"
+            )
 
         if body.stream:
 

xinference/client/restful/restful_client.py

@@ -320,6 +320,7 @@ class RESTfulChatModelHandle(RESTfulGenerateModelHandle):
         prompt: str,
         system_prompt: Optional[str] = None,
         chat_history: Optional[List["ChatCompletionMessage"]] = None,
+        tools: Optional[List[Dict]] = None,
         generate_config: Optional[
             Union["LlamaCppGenerateConfig", "PytorchGenerateConfig"]
         ] = None,
@@ -335,6 +336,8 @@ class RESTfulChatModelHandle(RESTfulGenerateModelHandle):
             The system context provide to Model prior to any chats.
         chat_history: Optional[List["ChatCompletionMessage"]]
             A list of messages comprising the conversation so far.
+        tools: Optional[List[Dict]]
+            A tool list.
         generate_config: Optional[Union["LlamaCppGenerateConfig", "PytorchGenerateConfig"]]
             Additional configuration for the chat generation.
             "LlamaCppGenerateConfig" -> configuration for ggml model
@@ -373,6 +376,8 @@ class RESTfulChatModelHandle(RESTfulGenerateModelHandle):
             "model": self._model_uid,
             "messages": chat_history,
         }
+        if tools is not None:
+            request_body["tools"] = tools
         if generate_config is not None:
             for key, value in generate_config.items():
                 request_body[key] = value
@@ -397,6 +402,7 @@ class RESTfulChatglmCppChatModelHandle(RESTfulEmbeddingModelHandle):
         self,
         prompt: str,
         chat_history: Optional[List["ChatCompletionMessage"]] = None,
+        tools: Optional[List[Dict]] = None,
         generate_config: Optional["ChatglmCppGenerateConfig"] = None,
     ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
         """
@@ -408,6 +414,8 @@ class RESTfulChatglmCppChatModelHandle(RESTfulEmbeddingModelHandle):
             The user's input.
         chat_history: Optional[List["ChatCompletionMessage"]]
             A list of messages comprising the conversation so far.
+        tools: Optional[List[Dict]]
+            A tool list.
         generate_config: Optional["ChatglmCppGenerateConfig"]
             Additional configuration for ChatGLM chat generation.
 
@@ -436,7 +444,8 @@ class RESTfulChatglmCppChatModelHandle(RESTfulEmbeddingModelHandle):
             "model": self._model_uid,
             "messages": chat_history,
         }
-
+        if tools is not None:
+            request_body["tools"] = tools
         if generate_config is not None:
             for key, value in generate_config.items():
                 request_body[key] = value

xinference/conftest.py

@@ -184,7 +184,7 @@ def setup():
     local_cluster_proc = run_test_cluster_in_subprocess(
         supervisor_addr, TEST_LOGGING_CONF
     )
-    if not cluster_health_check(supervisor_addr, max_attempts=3, sleep_interval=3):
+    if not cluster_health_check(supervisor_addr, max_attempts=10, sleep_interval=3):
         raise RuntimeError("Cluster is not available after multiple attempts")
 
     port = xo.utils.get_next_port()
@@ -195,7 +195,7 @@ def setup():
         logging_conf=TEST_LOGGING_CONF,
     )
     endpoint = f"http://localhost:{port}"
-    if not api_health_check(endpoint, max_attempts=3, sleep_interval=5):
+    if not api_health_check(endpoint, max_attempts=10, sleep_interval=5):
         raise RuntimeError("Endpoint is not available after multiple attempts")
 
     yield f"http://localhost:{port}", supervisor_addr

xinference/constants.py

@@ -18,6 +18,9 @@ from pathlib import Path
 XINFERENCE_ENV_ENDPOINT = "XINFERENCE_ENDPOINT"
 XINFERENCE_ENV_MODEL_SRC = "XINFERENCE_MODEL_SRC"
 XINFERENCE_ENV_HOME_PATH = "XINFERENCE_HOME"
+XINFERENCE_ENV_HEALTH_CHECK_ATTEMPTS = "XINFERENCE_HEALTH_CHECK_ATTEMPTS"
+XINFERENCE_ENV_HEALTH_CHECK_INTERVAL = "XINFERENCE_HEALTH_CHECK_INTERVAL"
+XINFERENCE_ENV_DISABLE_VLLM = "XINFERENCE_DISABLE_VLLM"
 
 
 def get_xinference_home():
@@ -36,3 +39,10 @@ XINFERENCE_DEFAULT_ENDPOINT_PORT = 9997
 XINFERENCE_DEFAULT_LOG_FILE_NAME = "xinference.log"
 XINFERENCE_LOG_MAX_BYTES = 100 * 1024 * 1024
 XINFERENCE_LOG_BACKUP_COUNT = 30
+XINFERENCE_HEALTH_CHECK_ATTEMPTS = int(
+    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_ATTEMPTS, 3)
+)
+XINFERENCE_HEALTH_CHECK_INTERVAL = int(
+    os.environ.get(XINFERENCE_ENV_HEALTH_CHECK_INTERVAL, 3)
+)
+XINFERENCE_DISABLE_VLLM = bool(int(os.environ.get(XINFERENCE_ENV_DISABLE_VLLM, 0)))

xinference/core/model.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import inspect
+import os
 import uuid
 from typing import (
     TYPE_CHECKING,
@@ -44,6 +45,15 @@ from .utils import json_dumps, log_async
 
 T = TypeVar("T")
 
+try:
+    from torch.cuda import OutOfMemoryError
+except ImportError:
+
+    class _OutOfMemoryError(Exception):
+        pass
+
+    OutOfMemoryError = _OutOfMemoryError
+
 
 def request_limit(fn):
     """
@@ -192,18 +202,30 @@ class ModelActor(xo.StatelessActor):
             return ret
 
     async def _call_wrapper(self, _wrapper: Callable):
-        assert not (
-            inspect.iscoroutinefunction(_wrapper)
-            or inspect.isasyncgenfunction(_wrapper)
-        )
-        if self._lock is None:
-            return await asyncio.to_thread(_wrapper)
-        else:
-            async with self._lock:
+        try:
+            assert not (
+                inspect.iscoroutinefunction(_wrapper)
+                or inspect.isasyncgenfunction(_wrapper)
+            )
+            if self._lock is None:
                 return await asyncio.to_thread(_wrapper)
+            else:
+                async with self._lock:
+                    return await asyncio.to_thread(_wrapper)
+        except OutOfMemoryError:
+            logger.exception(
+                "Model actor is out of memory, model id: %s", self.model_uid()
+            )
+            os._exit(1)
 
     async def _call_async_wrapper(self, _wrapper: Callable):
-        return await asyncio.create_task(_wrapper())
+        try:
+            return await asyncio.create_task(_wrapper())
+        except OutOfMemoryError:
+            logger.exception(
+                "Model actor is out of memory, model id: %s", self.model_uid()
+            )
+            os._exit(1)
 
     @log_async(logger=logger)
     @request_limit
@@ -365,7 +387,8 @@ class ModelActor(xo.StatelessActor):
 
         async def _async_wrapper():
             try:
-                return await anext(gen)  # noqa: F821
+                # anext is only available for Python >= 3.10
+                return await gen.__anext__()  # noqa: F821
             except StopAsyncIteration:
                 return stop
 

xinference/core/resource.py

@@ -15,7 +15,7 @@
 from dataclasses import dataclass
 from typing import Dict
 
-from xorbits._mars import resource
+import psutil
 
 
 @dataclass
@@ -28,19 +28,20 @@ class ResourceStatus:
 
 def gather_node_info() -> Dict[str, ResourceStatus]:
     node_resource = dict()
-    mem_info = resource.virtual_memory()
+    mem_info = psutil.virtual_memory()
     node_resource["cpu"] = ResourceStatus(
-        available=resource.cpu_percent() / 100.0,
-        total=resource.cpu_count(),
+        available=psutil.cpu_percent() / 100.0,
+        total=psutil.cpu_count(),
         memory_available=mem_info.available,
         memory_total=mem_info.total,
     )
-    for idx, gpu_card_stat in enumerate(resource.cuda_card_stats()):
-        node_resource[f"gpu-{idx}"] = ResourceStatus(
-            available=gpu_card_stat.gpu_usage / 100.0,
-            total=1,
-            memory_available=gpu_card_stat.fb_mem_info.available,
-            memory_total=gpu_card_stat.fb_mem_info.total,
-        )
+    # TODO: record GPU stats
+    # for idx, gpu_card_stat in enumerate(resource.cuda_card_stats()):
+    #     node_resource[f"gpu-{idx}"] = ResourceStatus(
+    #         available=gpu_card_stat.gpu_usage / 100.0,
+    #         total=1,
+    #         memory_available=gpu_card_stat.fb_mem_info.available,
+    #         memory_total=gpu_card_stat.fb_mem_info.total,
+    #     )
 
     return node_resource

xinference/core/supervisor.py

@@ -85,10 +85,11 @@ class SupervisorActor(xo.StatelessActor):
             register_embedding,
             unregister_embedding,
         )
-        from ..model.llm import LLMFamilyV1, register_llm, unregister_llm
+        from ..model.llm import register_llm, unregister_llm
+        from ..model.llm.llm_family import CustomLLMFamilyV1
 
         self._custom_register_type_to_cls: Dict[str, Tuple] = {
-            "LLM": (LLMFamilyV1, register_llm, unregister_llm),
+            "LLM": (CustomLLMFamilyV1, register_llm, unregister_llm),
             "embedding": (
                 CustomEmbeddingModelSpec,
                 register_embedding,
@@ -96,6 +97,25 @@ class SupervisorActor(xo.StatelessActor):
             ),
         }
 
+    @staticmethod
+    async def get_builtin_prompts() -> Dict[str, Any]:
+        from ..model.llm.llm_family import BUILTIN_LLM_PROMPT_STYLE
+
+        data = {}
+        for k, v in BUILTIN_LLM_PROMPT_STYLE.items():
+            data[k] = v.dict()
+        return data
+
+    async def get_devices_count(self) -> int:
+        from ..utils import cuda_count
+
+        if self.is_local_deployment():
+            return cuda_count()
+        # distributed deployment, choose a worker and return its cuda_count.
+        # Assume that each worker has the same count of cards.
+        worker_ref = await self._choose_worker()
+        return await worker_ref.get_devices_count()
+
     async def _choose_worker(self) -> xo.ActorRefType["WorkerActor"]:
         # TODO: better allocation strategy.
         min_running_model_count = None

xinference/core/worker.py

@@ -20,12 +20,12 @@ from logging import getLogger
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import xoscar as xo
-from xorbits._mars.resource import cuda_count
 from xoscar import MainActorPoolType
 
 from ..constants import XINFERENCE_CACHE_DIR
 from ..core import ModelActor
 from ..model.core import ModelDescription, create_model_instance
+from ..utils import cuda_count
 from .resource import gather_node_info
 from .utils import log_async, log_sync, parse_replica_model_uid, purge_dir
 
@@ -48,6 +48,7 @@ class WorkerActor(xo.StatelessActor):
         self._supervisor_address = supervisor_address
         self._supervisor_ref = None
         self._main_pool = main_pool
+        self._main_pool.recover_sub_pool = self.recover_sub_pool
 
         # internal states.
         self._model_uid_to_model: Dict[str, xo.ActorRefType["ModelActor"]] = {}
@@ -55,9 +56,22 @@ class WorkerActor(xo.StatelessActor):
         self._gpu_to_model_uid: Dict[int, str] = {}
         self._gpu_to_embedding_model_uids: Dict[int, Set[str]] = defaultdict(set)
         self._model_uid_to_addr: Dict[str, str] = {}
+        self._model_uid_to_launch_args: Dict[str, Dict] = {}
 
         self._lock = asyncio.Lock()
 
+    async def recover_sub_pool(self, address):
+        logger.warning("Process %s is down, create model.", address)
+        for model_uid, addr in self._model_uid_to_addr.items():
+            if addr == address:
+                launch_args = self._model_uid_to_launch_args.get(model_uid)
+                try:
+                    await self.terminate_model(model_uid)
+                except Exception:
+                    pass
+                await self.launch_builtin_model(**launch_args)
+                break
+
     @classmethod
     def uid(cls) -> str:
         return "worker"
@@ -79,10 +93,11 @@ class WorkerActor(xo.StatelessActor):
             register_embedding,
             unregister_embedding,
         )
-        from ..model.llm import LLMFamilyV1, register_llm, unregister_llm
+        from ..model.llm import register_llm, unregister_llm
+        from ..model.llm.llm_family import CustomLLMFamilyV1
 
         self._custom_register_type_to_cls: Dict[str, Tuple] = {
-            "LLM": (LLMFamilyV1, register_llm, unregister_llm),
+            "LLM": (CustomLLMFamilyV1, register_llm, unregister_llm),
             "embedding": (
                 CustomEmbeddingModelSpec,
                 register_embedding,
@@ -93,6 +108,12 @@ class WorkerActor(xo.StatelessActor):
     async def __pre_destroy__(self):
         self._upload_task.cancel()
 
+    @staticmethod
+    def get_devices_count():
+        from ..utils import cuda_count
+
+        return cuda_count()
+
     @log_sync(logger=logger)
     def get_model_count(self) -> int:
         return len(self._model_uid_to_model)
@@ -174,7 +195,7 @@ class WorkerActor(xo.StatelessActor):
             gpu_cnt = n_gpu if isinstance(n_gpu, int) else 1
             devices = (
                 [await self.allocate_devices_for_embedding(model_uid)]
-                if model_type == "embedding"
+                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])
@@ -277,7 +298,6 @@ class WorkerActor(xo.StatelessActor):
         for dev in devices:
             self._gpu_to_model_uid[int(dev)] = model_uid
         self._model_uid_to_addr[model_uid] = subpool_address
-        return model_ref
 
     @log_async(logger=logger)
     async def launch_builtin_model(
@@ -291,7 +311,9 @@ class WorkerActor(xo.StatelessActor):
         n_gpu: Optional[Union[int, str]] = "auto",
         request_limits: Optional[int] = None,
         **kwargs,
-    ) -> xo.ActorRefType["ModelActor"]:
+    ):
+        launch_args = locals()
+        launch_args.pop("self")
         if n_gpu is not None:
             if isinstance(n_gpu, int) and (n_gpu <= 0 or n_gpu > cuda_count()):
                 raise ValueError(
@@ -342,7 +364,7 @@ class WorkerActor(xo.StatelessActor):
         self._model_uid_to_model[model_uid] = model_ref
         self._model_uid_to_model_spec[model_uid] = model_description
         self._model_uid_to_addr[model_uid] = subpool_address
-        return model_ref
+        self._model_uid_to_launch_args[model_uid] = launch_args
 
     @log_async(logger=logger)
     async def terminate_model(self, model_uid: str):
@@ -350,15 +372,21 @@ class WorkerActor(xo.StatelessActor):
         if model_ref is None:
             raise ValueError(f"Model not found in the model list, uid: {model_uid}")
 
-        await xo.destroy_actor(model_ref)
-        del self._model_uid_to_model[model_uid]
-        del self._model_uid_to_model_spec[model_uid]
-
-        self.release_devices(model_uid)
-
-        subpool_address = self._model_uid_to_addr[model_uid]
-        await self._main_pool.remove_sub_pool(subpool_address)
-        del self._model_uid_to_addr[model_uid]
+        try:
+            await xo.destroy_actor(model_ref)
+        except Exception as e:
+            logger.debug(
+                "Destroy model actor failed, model uid: %s, error: %s", model_uid, e
+            )
+        try:
+            subpool_address = self._model_uid_to_addr[model_uid]
+            await self._main_pool.remove_sub_pool(subpool_address)
+        finally:
+            del self._model_uid_to_model[model_uid]
+            del self._model_uid_to_model_spec[model_uid]
+            self.release_devices(model_uid)
+            del self._model_uid_to_addr[model_uid]
+            del self._model_uid_to_launch_args[model_uid]
 
     @log_async(logger=logger)
     async def list_models(self) -> Dict[str, Dict[str, Any]]:

xinference/deploy/cmdline.py

@@ -189,13 +189,14 @@ def local(
     "-p",
     default=XINFERENCE_DEFAULT_ENDPOINT_PORT,
     type=int,
-    help="Specify the port number for the supervisor.",
+    help="Specify the port number for the Xinference web ui and service.",
 )
-def supervisor(
-    log_level: str,
-    host: str,
-    port: int,
-):
+@click.option(
+    "--supervisor-port",
+    type=int,
+    help="Specify the port number for the Xinference supervisor.",
+)
+def supervisor(log_level: str, host: str, port: int, supervisor_port: Optional[int]):
     from ..deploy.supervisor import main
 
     dict_config = get_config_dict(
@@ -206,7 +207,9 @@ def supervisor(
     )
     logging.config.dictConfig(dict_config)  # type: ignore
 
-    main(host=host, port=port, logging_conf=dict_config)
+    main(
+        host=host, port=port, supervisor_port=supervisor_port, logging_conf=dict_config
+    )
 
 
 @click.command(
@@ -227,7 +230,14 @@ def supervisor(
     type=str,
     help="Specify the host address for the worker.",
 )
-def worker(log_level: str, endpoint: Optional[str], host: str):
+@click.option(
+    "--worker-port",
+    type=int,
+    help="Specify the port number for the Xinference worker.",
+)
+def worker(
+    log_level: str, endpoint: Optional[str], host: str, worker_port: Optional[int]
+):
     from ..deploy.worker import main
 
     dict_config = get_config_dict(
@@ -243,7 +253,7 @@ def worker(log_level: str, endpoint: Optional[str], host: str):
     client = RESTfulClient(base_url=endpoint)
     supervisor_internal_addr = client._get_supervisor_internal_address()
 
-    address = f"{host}:{get_next_port()}"
+    address = f"{host}:{worker_port or get_next_port()}"
     main(
         address=address,
         supervisor_address=supervisor_internal_addr,

xinference/deploy/local.py

@@ -22,6 +22,10 @@ from typing import Dict, Optional
 import xoscar as xo
 from xoscar.utils import get_next_port
 
+from ..constants import (
+    XINFERENCE_HEALTH_CHECK_ATTEMPTS,
+    XINFERENCE_HEALTH_CHECK_INTERVAL,
+)
 from ..core.supervisor import SupervisorActor
 from .utils import health_check
 from .worker import start_worker_components
@@ -79,7 +83,11 @@ def main(host: str, port: int, logging_conf: Optional[Dict] = None):
     supervisor_address = f"{host}:{get_next_port()}"
     local_cluster = run_in_subprocess(supervisor_address, logging_conf)
 
-    if not health_check(address=supervisor_address, max_attempts=3, sleep_interval=3):
+    if not health_check(
+        address=supervisor_address,
+        max_attempts=XINFERENCE_HEALTH_CHECK_ATTEMPTS,
+        sleep_interval=XINFERENCE_HEALTH_CHECK_INTERVAL,
+    ):
         raise RuntimeError("Cluster is not available after multiple attempts")
 
     try: