xinference 0.9.4__py3-none-any.whl → 0.10.1__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-03-21T14:58:01+0800",
+ "date": "2024-04-11T15:35:46+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "2c9465ade7f358d57d4bc087277882d896a8de15",
- "version": "0.9.4"
+ "full-revisionid": "e3a947ebddfc53b5e8ec723c1f632c2b895edef1",
+ "version": "0.10.1"
 }
 '''  # END VERSION_JSON
 

xinference/api/oauth2/auth_service.py

@@ -11,8 +11,9 @@
 # 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.
+import re
 from datetime import timedelta
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from fastapi import Depends, HTTPException, status
 from fastapi.security import OAuth2PasswordBearer, SecurityScopes
@@ -40,13 +41,30 @@ class AuthService:
     def config(self):
         return self._config
 
+    @staticmethod
+    def is_legal_api_key(key: str) -> bool:
+        pattern = re.compile("^sk-[a-zA-Z0-9]{13}$")
+        return re.match(pattern, key) is not None
+
     def init_auth_config(self):
         if self._auth_config_file:
             config: AuthStartupConfig = parse_file_as(
                 path=self._auth_config_file, type_=AuthStartupConfig
             )
+            all_api_keys = set()
             for user in config.user_config:
                 user.password = get_password_hash(user.password)
+                for api_key in user.api_keys:
+                    if not self.is_legal_api_key(api_key):
+                        raise ValueError(
+                            "Api-Key should be a string started with 'sk-' with a total length of 16"
+                        )
+                    if api_key in all_api_keys:
+                        raise ValueError(
+                            "Duplicate api-keys exists, please check your configuration"
+                        )
+                    else:
+                        all_api_keys.add(api_key)
             return config
 
     def __call__(
@@ -67,28 +85,30 @@ class AuthService:
             headers={"WWW-Authenticate": authenticate_value},
         )
 
-        try:
-            assert self._config is not None
-            payload = jwt.decode(
-                token,
-                self._config.auth_config.secret_key,
-                algorithms=[self._config.auth_config.algorithm],
-                options={"verify_exp": False},  # TODO: supports token expiration
-            )
-            username: str = payload.get("sub")
-            if username is None:
+        if self.is_legal_api_key(token):
+            user, token_scopes = self.get_user_and_scopes_with_api_key(token)
+        else:
+            try:
+                assert self._config is not None
+                payload = jwt.decode(
+                    token,
+                    self._config.auth_config.secret_key,
+                    algorithms=[self._config.auth_config.algorithm],
+                    options={"verify_exp": False},  # TODO: supports token expiration
+                )
+                username: str = payload.get("sub")
+                if username is None:
+                    raise credentials_exception
+                token_scopes = payload.get("scopes", [])
+                user = self.get_user(username)
+            except (JWTError, ValidationError):
                 raise credentials_exception
-            token_scopes = payload.get("scopes", [])
-            token_data = TokenData(scopes=token_scopes, username=username)
-        except (JWTError, ValidationError):
-            raise credentials_exception
-        user = self.get_user(token_data.username)
         if user is None:
             raise credentials_exception
-        if "admin" in token_data.scopes:
+        if "admin" in token_scopes:
             return user
         for scope in security_scopes.scopes:
-            if scope not in token_data.scopes:
+            if scope not in token_scopes:
                 raise HTTPException(
                     status_code=status.HTTP_403_FORBIDDEN,
                     detail="Not enough permissions",
@@ -102,6 +122,15 @@ class AuthService:
                 return user
         return None
 
+    def get_user_and_scopes_with_api_key(
+        self, api_key: str
+    ) -> Tuple[Optional[User], List]:
+        for user in self._config.user_config:
+            for key in user.api_keys:
+                if api_key == key:
+                    return user, user.permissions
+        return None, []
+
     def authenticate_user(self, username: str, password: str):
         user = self.get_user(username)
         if not user:

xinference/api/oauth2/types.py

@@ -23,6 +23,7 @@ class LoginUserForm(BaseModel):
 
 class User(LoginUserForm):
     permissions: List[str]
+    api_keys: List[str]
 
 
 class AuthConfig(BaseModel):

xinference/api/restful_api.py

@@ -89,7 +89,9 @@ class CreateCompletionRequest(CreateCompletion):
 
 class CreateEmbeddingRequest(BaseModel):
     model: str
-    input: Union[str, List[str]] = Field(description="The input to embed.")
+    input: Union[str, List[str], List[int], List[List[int]]] = Field(
+        description="The input to embed."
+    )
     user: Optional[str] = None
 
     class Config:
@@ -693,6 +695,8 @@ class RESTfulAPI:
         peft_model_path = payload.get("peft_model_path", None)
         image_lora_load_kwargs = payload.get("image_lora_load_kwargs", None)
         image_lora_fuse_kwargs = payload.get("image_lora_fuse_kwargs", None)
+        worker_ip = payload.get("worker_ip", None)
+        gpu_idx = payload.get("gpu_idx", None)
 
         exclude_keys = {
             "model_uid",
@@ -707,6 +711,8 @@ class RESTfulAPI:
             "peft_model_path",
             "image_lora_load_kwargs",
             "image_lora_fuse_kwargs",
+            "worker_ip",
+            "gpu_idx",
         }
 
         kwargs = {
@@ -734,6 +740,8 @@ class RESTfulAPI:
                 peft_model_path=peft_model_path,
                 image_lora_load_kwargs=image_lora_load_kwargs,
                 image_lora_fuse_kwargs=image_lora_fuse_kwargs,
+                worker_ip=worker_ip,
+                gpu_idx=gpu_idx,
                 **kwargs,
             )
 
@@ -999,8 +1007,16 @@ class RESTfulAPI:
                 raise HTTPException(status_code=500, detail=str(e))
 
     async def create_embedding(self, request: Request) -> Response:
-        body = CreateEmbeddingRequest.parse_obj(await request.json())
+        payload = await request.json()
+        body = CreateEmbeddingRequest.parse_obj(payload)
         model_uid = body.model
+        exclude = {
+            "model",
+            "input",
+            "user",
+            "encoding_format",
+        }
+        kwargs = {key: value for key, value in payload.items() if key not in exclude}
 
         try:
             model = await (await self._get_supervisor_ref()).get_model(model_uid)
@@ -1014,7 +1030,7 @@ class RESTfulAPI:
             raise HTTPException(status_code=500, detail=str(e))
 
         try:
-            embedding = await model.create_embedding(body.input)
+            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)
@@ -1027,8 +1043,15 @@ class RESTfulAPI:
             raise HTTPException(status_code=500, detail=str(e))
 
     async def rerank(self, request: Request) -> Response:
-        body = RerankRequest.parse_obj(await request.json())
+        payload = await request.json()
+        body = RerankRequest.parse_obj(payload)
         model_uid = body.model
+        kwargs = {
+            key: value
+            for key, value in payload.items()
+            if key not in RerankRequest.__annotations__.keys()
+        }
+
         try:
             model = await (await self._get_supervisor_ref()).get_model(model_uid)
         except ValueError as ve:
@@ -1047,6 +1070,7 @@ class RESTfulAPI:
                 top_n=body.top_n,
                 max_chunks_per_doc=body.max_chunks_per_doc,
                 return_documents=body.return_documents,
+                **kwargs,
             )
             return Response(scores, media_type="application/json")
         except RuntimeError as re:
@@ -1337,9 +1361,12 @@ class RESTfulAPI:
                     detail=f"Only {function_call_models} support tool messages",
                 )
         if body.tools and body.stream:
-            raise HTTPException(
-                status_code=400, detail="Tool calls does not support stream"
-            )
+            is_vllm = await model.is_vllm_backend()
+            if not is_vllm or model_family not in ["qwen-chat", "qwen1.5-chat"]:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Streaming support for tool calls is available only when using vLLM backend and Qwen models.",
+                )
 
         if body.stream:
 

xinference/client/oscar/actor_client.py

@@ -111,7 +111,7 @@ class ClientIteratorWrapper(AsyncIterator):
 
 
 class EmbeddingModelHandle(ModelHandle):
-    def create_embedding(self, input: Union[str, List[str]]) -> bytes:
+    def create_embedding(self, input: Union[str, List[str]], **kwargs) -> bytes:
         """
         Creates an embedding vector representing the input text.
 
@@ -128,7 +128,7 @@ class EmbeddingModelHandle(ModelHandle):
             machine learning models and algorithms.
         """
 
-        coro = self._model_ref.create_embedding(input)
+        coro = self._model_ref.create_embedding(input, **kwargs)
         return orjson.loads(self._isolation.call(coro))
 
 
@@ -140,6 +140,7 @@ class RerankModelHandle(ModelHandle):
         top_n: Optional[int],
         max_chunks_per_doc: Optional[int],
         return_documents: Optional[bool],
+        **kwargs,
     ):
         """
         Returns an ordered list of documents ordered by their relevance to the provided query.
@@ -163,7 +164,7 @@ class RerankModelHandle(ModelHandle):
 
         """
         coro = self._model_ref.rerank(
-            documents, query, top_n, max_chunks_per_doc, return_documents
+            documents, query, top_n, max_chunks_per_doc, return_documents, **kwargs
         )
         results = orjson.loads(self._isolation.call(coro))
         for r in results["results"]:

xinference/client/restful/restful_client.py

@@ -80,7 +80,7 @@ class RESTfulModelHandle:
 
 
 class RESTfulEmbeddingModelHandle(RESTfulModelHandle):
-    def create_embedding(self, input: Union[str, List[str]]) -> "Embedding":
+    def create_embedding(self, input: Union[str, List[str]], **kwargs) -> "Embedding":
         """
         Create an Embedding from user input via RESTful APIs.
 
@@ -102,7 +102,11 @@ class RESTfulEmbeddingModelHandle(RESTfulModelHandle):
 
         """
         url = f"{self._base_url}/v1/embeddings"
-        request_body = {"model": self._model_uid, "input": input}
+        request_body = {
+            "model": self._model_uid,
+            "input": input,
+        }
+        request_body.update(kwargs)
         response = requests.post(url, json=request_body, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
@@ -121,6 +125,7 @@ class RESTfulRerankModelHandle(RESTfulModelHandle):
         top_n: Optional[int] = None,
         max_chunks_per_doc: Optional[int] = None,
         return_documents: Optional[bool] = None,
+        **kwargs,
     ):
         """
         Returns an ordered list of documents ordered by their relevance to the provided query.
@@ -156,6 +161,7 @@ class RESTfulRerankModelHandle(RESTfulModelHandle):
             "max_chunks_per_doc": max_chunks_per_doc,
             "return_documents": return_documents,
         }
+        request_body.update(kwargs)
         response = requests.post(url, json=request_body, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
@@ -651,11 +657,13 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
 
 
 class Client:
-    def __init__(self, base_url):
+    def __init__(self, base_url, api_key: Optional[str] = None):
         self.base_url = base_url
-        self._headers = {}
+        self._headers: Dict[str, str] = {}
         self._cluster_authed = False
         self._check_cluster_authenticated()
+        if api_key is not None and self._cluster_authed:
+            self._headers["Authorization"] = f"Bearer {api_key}"
 
     def _set_token(self, token: Optional[str]):
         if not self._cluster_authed or token is None:
@@ -795,6 +803,8 @@ class Client:
         peft_model_path: Optional[str] = None,
         image_lora_load_kwargs: Optional[Dict] = None,
         image_lora_fuse_kwargs: Optional[Dict] = None,
+        worker_ip: Optional[str] = None,
+        gpu_idx: Optional[Union[int, List[int]]] = None,
         **kwargs,
     ) -> str:
         """
@@ -828,6 +838,10 @@ class Client:
             lora load parameters for image model
         image_lora_fuse_kwargs: Optional[Dict]
             lora fuse parameters for image model
+        worker_ip: Optional[str]
+            Specify the worker ip where the model is located in a distributed scenario.
+        gpu_idx: Optional[Union[int, List[int]]]
+            Specify the GPU index where the model is located.
         **kwargs:
             Any other parameters been specified.
 
@@ -853,6 +867,8 @@ class Client:
             "peft_model_path": peft_model_path,
             "image_lora_load_kwargs": image_lora_load_kwargs,
             "image_lora_fuse_kwargs": image_lora_fuse_kwargs,
+            "worker_ip": worker_ip,
+            "gpu_idx": gpu_idx,
         }
 
         for key, value in kwargs.items():

xinference/conftest.py

@@ -261,12 +261,23 @@ def setup_with_auth():
     if not cluster_health_check(supervisor_addr, max_attempts=10, sleep_interval=3):
         raise RuntimeError("Cluster is not available after multiple attempts")
 
-    user1 = User(username="user1", password="pass1", permissions=["admin"])
-    user2 = User(username="user2", password="pass2", permissions=["models:list"])
+    user1 = User(
+        username="user1",
+        password="pass1",
+        permissions=["admin"],
+        api_keys=["sk-3sjLbdwqAhhAF", "sk-0HCRO1rauFQDL"],
+    )
+    user2 = User(
+        username="user2",
+        password="pass2",
+        permissions=["models:list"],
+        api_keys=["sk-72tkvudyGLPMi"],
+    )
     user3 = User(
         username="user3",
         password="pass3",
         permissions=["models:list", "models:read", "models:start"],
+        api_keys=["sk-m6jEzEwmCc4iQ", "sk-ZOTLIY4gt9w11"],
     )
     auth_config = AuthConfig(
         algorithm="HS256",

xinference/core/supervisor.py

@@ -92,6 +92,15 @@ class SupervisorActor(xo.StatelessActor):
     def uid(cls) -> str:
         return "supervisor"
 
+    def _get_worker_ref_by_ip(
+        self, ip: str
+    ) -> Optional[xo.ActorRefType["WorkerActor"]]:
+        for addr, ref in self._worker_address_to_worker.items():
+            existing_ip = addr.split(":")[0]
+            if existing_ip == ip:
+                return ref
+        return None
+
     async def __post_create__(self):
         self._uptime = time.time()
         if not XINFERENCE_DISABLE_HEALTH_CHECK:
@@ -717,8 +726,25 @@ class SupervisorActor(xo.StatelessActor):
         peft_model_path: Optional[str] = None,
         image_lora_load_kwargs: Optional[Dict] = None,
         image_lora_fuse_kwargs: Optional[Dict] = None,
+        worker_ip: Optional[str] = None,
+        gpu_idx: Optional[Union[int, List[int]]] = None,
         **kwargs,
     ) -> str:
+        target_ip_worker_ref = (
+            self._get_worker_ref_by_ip(worker_ip) if worker_ip is not None else None
+        )
+        if (
+            worker_ip is not None
+            and not self.is_local_deployment()
+            and target_ip_worker_ref is None
+        ):
+            raise ValueError(f"Worker ip address {worker_ip} is not in the cluster.")
+        if worker_ip is not None and self.is_local_deployment():
+            logger.warning(
+                f"You specified the worker ip: {worker_ip} in local mode, "
+                f"xinference will ignore this option."
+            )
+
         if model_uid is None:
             model_uid = self._gen_model_uid(model_name)
 
@@ -735,7 +761,11 @@ class SupervisorActor(xo.StatelessActor):
                 )
 
             nonlocal model_type
-            worker_ref = await self._choose_worker()
+            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(
@@ -750,6 +780,7 @@ class SupervisorActor(xo.StatelessActor):
                 peft_model_path=peft_model_path,
                 image_lora_load_kwargs=image_lora_load_kwargs,
                 image_lora_fuse_kwargs=image_lora_fuse_kwargs,
+                gpu_idx=gpu_idx,
                 **kwargs,
             )
             self._replica_model_uid_to_worker[_replica_model_uid] = worker_ref
@@ -839,6 +870,12 @@ class SupervisorActor(xo.StatelessActor):
                             address,
                             dead_models,
                         )
+                        for replica_model_uid in dead_models:
+                            model_uid, _, _ = parse_replica_model_uid(replica_model_uid)
+                            self._model_uid_to_replica_info.pop(model_uid, None)
+                            self._replica_model_uid_to_worker.pop(
+                                replica_model_uid, None
+                            )
                         dead_nodes.append(address)
                     elif (
                         status.failure_remaining_count
@@ -948,6 +985,16 @@ class SupervisorActor(xo.StatelessActor):
 
     @log_async(logger=logger)
     async def remove_worker(self, worker_address: str):
+        uids_to_remove = []
+        for model_uid in self._replica_model_uid_to_worker:
+            if self._replica_model_uid_to_worker[model_uid].address == worker_address:
+                uids_to_remove.append(model_uid)
+
+        for replica_model_uid in uids_to_remove:
+            model_uid, _, _ = parse_replica_model_uid(replica_model_uid)
+            self._model_uid_to_replica_info.pop(model_uid, None)
+            self._replica_model_uid_to_worker.pop(replica_model_uid, None)
+
         if worker_address in self._worker_address_to_worker:
             del self._worker_address_to_worker[worker_address]
             logger.debug("Worker %s has been removed successfully", worker_address)