xinference 0.7.5__py3-none-any.whl → 0.8.1__py3-none-any.whl

xinference/client/oscar/actor_client.py

@@ -14,12 +14,12 @@
 
 import asyncio
 import re
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Union
 
 import orjson
 import xoscar as xo
 
-from ...core.model import IteratorWrapper, ModelActor
+from ...core.model import ModelActor
 from ...core.supervisor import SupervisorActor
 from ...isolation import Isolation
 from ..restful.restful_client import Client
@@ -97,18 +97,18 @@ class ModelHandle:
         self._isolation = isolation
 
 
-class ClientIteratorWrapper(IteratorWrapper):
+class ClientIteratorWrapper(AsyncIterator):
+    def __init__(self, iterator_wrapper):
+        self._iw = iterator_wrapper
+
+    def __aiter__(self):
+        return self
+
     async def __anext__(self):
-        r = await super().__anext__()
+        r = await self._iw.__anext__()
         text = r.decode("utf-8")
         return orjson.loads(SSEEvent.parse(text).data)
 
-    @classmethod
-    def wrap(cls, iterator_wrapper):
-        c = cls.__new__(cls)
-        c.__dict__.update(iterator_wrapper.__dict__)
-        return c
-
 
 class EmbeddingModelHandle(ModelHandle):
     def create_embedding(self, input: Union[str, List[str]]) -> bytes:
@@ -171,14 +171,14 @@ class RerankModelHandle(ModelHandle):
         return results
 
 
-class GenerateModelHandle(EmbeddingModelHandle):
+class GenerateModelHandle(ModelHandle):
     def generate(
         self,
         prompt: str,
         generate_config: Optional[
             Union["LlamaCppGenerateConfig", "PytorchGenerateConfig"]
         ] = None,
-    ) -> Union["Completion", Iterator["CompletionChunk"]]:
+    ) -> Union["Completion", AsyncIterator["CompletionChunk"]]:
         """
         Creates a completion for the provided prompt and parameters.
 
@@ -204,7 +204,7 @@ class GenerateModelHandle(EmbeddingModelHandle):
         r = self._isolation.call(coro)
         if isinstance(r, bytes):
             return orjson.loads(r)
-        return ClientIteratorWrapper.wrap(r)
+        return ClientIteratorWrapper(r)
 
 
 class ChatModelHandle(GenerateModelHandle):
@@ -216,7 +216,7 @@ class ChatModelHandle(GenerateModelHandle):
         generate_config: Optional[
             Union["LlamaCppGenerateConfig", "PytorchGenerateConfig"]
         ] = None,
-    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+    ) -> Union["ChatCompletion", AsyncIterator["ChatCompletionChunk"]]:
         """
         Given a list of messages comprising a conversation, the model will return a response.
 
@@ -252,16 +252,16 @@ class ChatModelHandle(GenerateModelHandle):
         r = self._isolation.call(coro)
         if isinstance(r, bytes):
             return orjson.loads(r)
-        return ClientIteratorWrapper.wrap(r)
+        return ClientIteratorWrapper(r)
 
 
-class ChatglmCppChatModelHandle(EmbeddingModelHandle):
+class ChatglmCppChatModelHandle(ModelHandle):
     def chat(
         self,
         prompt: str,
         chat_history: Optional[List["ChatCompletionMessage"]] = None,
         generate_config: Optional["ChatglmCppGenerateConfig"] = None,
-    ) -> Union["ChatCompletion", Iterator["ChatCompletionChunk"]]:
+    ) -> Union["ChatCompletion", AsyncIterator["ChatCompletionChunk"]]:
         """
         Given a list of messages comprising a conversation, the ChatGLM model will return a response.
 
@@ -287,7 +287,7 @@ class ChatglmCppChatModelHandle(EmbeddingModelHandle):
         r = self._isolation.call(coro)
         if isinstance(r, bytes):
             return orjson.loads(r)
-        return ClientIteratorWrapper.wrap(r)
+        return ClientIteratorWrapper(r)
 
 
 class ImageModelHandle(ModelHandle):

xinference/client/restful/restful_client.py

@@ -53,9 +53,10 @@ class RESTfulModelHandle:
     programmatically.
     """
 
-    def __init__(self, model_uid: str, base_url: str):
+    def __init__(self, model_uid: str, base_url: str, auth_headers: Dict):
         self._model_uid = model_uid
         self._base_url = base_url
+        self.auth_headers = auth_headers
 
 
 class RESTfulEmbeddingModelHandle(RESTfulModelHandle):
@@ -82,7 +83,7 @@ class RESTfulEmbeddingModelHandle(RESTfulModelHandle):
         """
         url = f"{self._base_url}/v1/embeddings"
         request_body = {"model": self._model_uid, "input": input}
-        response = requests.post(url, json=request_body)
+        response = requests.post(url, json=request_body, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to create the embeddings, detail: {_get_error_string(response)}"
@@ -135,7 +136,7 @@ class RESTfulRerankModelHandle(RESTfulModelHandle):
             "max_chunks_per_doc": max_chunks_per_doc,
             "return_documents": return_documents,
         }
-        response = requests.post(url, json=request_body)
+        response = requests.post(url, json=request_body, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to rerank documents, detail: {response.json()['detail']}"
@@ -182,7 +183,7 @@ class RESTfulImageModelHandle(RESTfulModelHandle):
             "response_format": response_format,
             "kwargs": json.dumps(kwargs),
         }
-        response = requests.post(url, json=request_body)
+        response = requests.post(url, json=request_body, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to create the images, detail: {_get_error_string(response)}"
@@ -246,10 +247,7 @@ class RESTfulImageModelHandle(RESTfulModelHandle):
         for key, value in params.items():
             files.append((key, (None, value)))
         files.append(("image", ("image", image, "application/octet-stream")))
-        response = requests.post(
-            url,
-            files=files,
-        )
+        response = requests.post(url, files=files, headers=self.auth_headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to variants the images, detail: {_get_error_string(response)}"
@@ -259,7 +257,7 @@ class RESTfulImageModelHandle(RESTfulModelHandle):
         return response_data
 
 
-class RESTfulGenerateModelHandle(RESTfulEmbeddingModelHandle):
+class RESTfulGenerateModelHandle(RESTfulModelHandle):
     def generate(
         self,
         prompt: str,
@@ -302,7 +300,9 @@ class RESTfulGenerateModelHandle(RESTfulEmbeddingModelHandle):
 
         stream = bool(generate_config and generate_config.get("stream"))
 
-        response = requests.post(url, json=request_body, stream=stream)
+        response = requests.post(
+            url, json=request_body, stream=stream, headers=self.auth_headers
+        )
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to generate completion, detail: {_get_error_string(response)}"
@@ -384,7 +384,9 @@ class RESTfulChatModelHandle(RESTfulGenerateModelHandle):
                 request_body[key] = value
 
         stream = bool(generate_config and generate_config.get("stream"))
-        response = requests.post(url, json=request_body, stream=stream)
+        response = requests.post(
+            url, json=request_body, stream=stream, headers=self.auth_headers
+        )
 
         if response.status_code != 200:
             raise RuntimeError(
@@ -468,7 +470,9 @@ class RESTfulMultimodalModelHandle(RESTfulModelHandle):
                 request_body[key] = value
 
         stream = bool(generate_config and generate_config.get("stream"))
-        response = requests.post(url, json=request_body, stream=stream)
+        response = requests.post(
+            url, json=request_body, stream=stream, headers=self.auth_headers
+        )
 
         if response.status_code != 200:
             raise RuntimeError(
@@ -482,7 +486,7 @@ class RESTfulMultimodalModelHandle(RESTfulModelHandle):
         return response_data
 
 
-class RESTfulChatglmCppChatModelHandle(RESTfulEmbeddingModelHandle):
+class RESTfulChatglmCppChatModelHandle(RESTfulModelHandle):
     def chat(
         self,
         prompt: str,
@@ -536,7 +540,9 @@ class RESTfulChatglmCppChatModelHandle(RESTfulEmbeddingModelHandle):
                 request_body[key] = value
 
         stream = bool(generate_config and generate_config.get("stream"))
-        response = requests.post(url, json=request_body, stream=stream)
+        response = requests.post(
+            url, json=request_body, stream=stream, headers=self.auth_headers
+        )
 
         if response.status_code != 200:
             raise RuntimeError(
@@ -589,7 +595,9 @@ class RESTfulChatglmCppGenerateModelHandle(RESTfulChatglmCppChatModelHandle):
 
         stream = bool(generate_config and generate_config.get("stream"))
 
-        response = requests.post(url, json=request_body, stream=stream)
+        response = requests.post(
+            url, json=request_body, stream=stream, headers=self.auth_headers
+        )
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to generate completion, detail: {response.json()['detail']}"
@@ -605,6 +613,47 @@ class RESTfulChatglmCppGenerateModelHandle(RESTfulChatglmCppChatModelHandle):
 class Client:
     def __init__(self, base_url):
         self.base_url = base_url
+        self._headers = {}
+        self._cluster_authed = False
+        self._check_cluster_authenticated()
+
+    def _set_token(self, token: Optional[str]):
+        if not self._cluster_authed or token is None:
+            return
+        self._headers["Authorization"] = f"Bearer {token}"
+
+    def _get_token(self) -> Optional[str]:
+        return (
+            str(self._headers["Authorization"]).replace("Bearer ", "")
+            if "Authorization" in self._headers
+            else None
+        )
+
+    def _check_cluster_authenticated(self):
+        url = f"{self.base_url}/v1/cluster/auth"
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get cluster information, detail: {response.json()['detail']}"
+            )
+        response_data = response.json()
+        self._cluster_authed = bool(response_data["auth"])
+
+    def login(self, username: str, password: str):
+        if not self._cluster_authed:
+            return
+        url = f"{self.base_url}/token"
+
+        payload = {"username": username, "password": password}
+
+        response = requests.post(url, json=payload)
+        if response.status_code != 200:
+            raise RuntimeError(f"Failed to login, detail: {response.json()['detail']}")
+
+        response_data = response.json()
+        # Only bearer token for now
+        access_token = response_data["access_token"]
+        self._headers["Authorization"] = f"Bearer {access_token}"
 
     def list_models(self) -> Dict[str, Dict[str, Any]]:
         """
@@ -619,7 +668,7 @@ class Client:
 
         url = f"{self.base_url}/v1/models"
 
-        response = requests.get(url)
+        response = requests.get(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to list model, detail: {_get_error_string(response)}"
@@ -664,7 +713,7 @@ class Client:
         }
 
         url = f"{self.base_url}/experimental/speculative_llms"
-        response = requests.post(url, json=payload)
+        response = requests.post(url, json=payload, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to launch model, detail: {_get_error_string(response)}"
@@ -739,7 +788,7 @@ class Client:
         for key, value in kwargs.items():
             payload[str(key)] = value
 
-        response = requests.post(url, json=payload)
+        response = requests.post(url, json=payload, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to launch model, detail: {_get_error_string(response)}"
@@ -766,7 +815,7 @@ class Client:
 
         url = f"{self.base_url}/v1/models/{model_uid}"
 
-        response = requests.delete(url)
+        response = requests.delete(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to terminate model, detail: {_get_error_string(response)}"
@@ -774,7 +823,7 @@ class Client:
 
     def _get_supervisor_internal_address(self):
         url = f"{self.base_url}/v1/address"
-        response = requests.get(url)
+        response = requests.get(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(f"Failed to get supervisor internal address")
         response_data = response.json()
@@ -806,7 +855,7 @@ class Client:
         """
 
         url = f"{self.base_url}/v1/models/{model_uid}"
-        response = requests.get(url)
+        response = requests.get(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to get the model description, detail: {_get_error_string(response)}"
@@ -815,21 +864,35 @@ class Client:
 
         if desc["model_type"] == "LLM":
             if desc["model_format"] == "ggmlv3" and "chatglm" in desc["model_name"]:
-                return RESTfulChatglmCppGenerateModelHandle(model_uid, self.base_url)
+                return RESTfulChatglmCppGenerateModelHandle(
+                    model_uid, self.base_url, auth_headers=self._headers
+                )
             elif "chat" in desc["model_ability"]:
-                return RESTfulChatModelHandle(model_uid, self.base_url)
+                return RESTfulChatModelHandle(
+                    model_uid, self.base_url, auth_headers=self._headers
+                )
             elif "generate" in desc["model_ability"]:
-                return RESTfulGenerateModelHandle(model_uid, self.base_url)
+                return RESTfulGenerateModelHandle(
+                    model_uid, self.base_url, auth_headers=self._headers
+                )
             else:
                 raise ValueError(f"Unrecognized model ability: {desc['model_ability']}")
         elif desc["model_type"] == "embedding":
-            return RESTfulEmbeddingModelHandle(model_uid, self.base_url)
+            return RESTfulEmbeddingModelHandle(
+                model_uid, self.base_url, auth_headers=self._headers
+            )
         elif desc["model_type"] == "image":
-            return RESTfulImageModelHandle(model_uid, self.base_url)
+            return RESTfulImageModelHandle(
+                model_uid, self.base_url, auth_headers=self._headers
+            )
         elif desc["model_type"] == "rerank":
-            return RESTfulRerankModelHandle(model_uid, self.base_url)
+            return RESTfulRerankModelHandle(
+                model_uid, self.base_url, auth_headers=self._headers
+            )
         elif desc["model_type"] == "multimodal":
-            return RESTfulMultimodalModelHandle(model_uid, self.base_url)
+            return RESTfulMultimodalModelHandle(
+                model_uid, self.base_url, auth_headers=self._headers
+            )
         else:
             raise ValueError(f"Unknown model type:{desc['model_type']}")
 
@@ -876,7 +939,7 @@ class Client:
         """
 
         url = f"{self.base_url}/v1/models/{model_uid}"
-        response = requests.get(url)
+        response = requests.get(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to get the model description, detail: {_get_error_string(response)}"
@@ -903,7 +966,7 @@ class Client:
         """
         url = f"{self.base_url}/v1/model_registrations/{model_type}"
         request_body = {"model": model, "persist": persist}
-        response = requests.post(url, json=request_body)
+        response = requests.post(url, json=request_body, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to register model, detail: {_get_error_string(response)}"
@@ -929,7 +992,7 @@ class Client:
             Report failure to unregister the custom model. Provide details of failure through error message.
         """
         url = f"{self.base_url}/v1/model_registrations/{model_type}/{model_name}"
-        response = requests.delete(url)
+        response = requests.delete(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to register model, detail: {_get_error_string(response)}"
@@ -959,7 +1022,7 @@ class Client:
 
         """
         url = f"{self.base_url}/v1/model_registrations/{model_type}"
-        response = requests.get(url)
+        response = requests.get(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to list model registration, detail: {_get_error_string(response)}"
@@ -987,7 +1050,7 @@ class Client:
             The collection of registered models on the server.
         """
         url = f"{self.base_url}/v1/model_registrations/{model_type}/{model_name}"
-        response = requests.get(url)
+        response = requests.get(url, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(
                 f"Failed to list model registration, detail: {_get_error_string(response)}"

xinference/conftest.py

@@ -13,16 +13,19 @@
 # limitations under the License.
 
 import asyncio
+import json
 import logging
 import multiprocessing
 import os
 import signal
 import sys
+import tempfile
 from typing import Dict, Optional
 
 import pytest
 import xoscar as xo
 
+from .api.oauth2.types import AuthConfig, AuthStartupConfig, User
 from .constants import XINFERENCE_LOG_BACKUP_COUNT, XINFERENCE_LOG_MAX_BYTES
 from .core.supervisor import SupervisorActor
 from .deploy.utils import create_worker_actor_pool, get_log_file, get_timestamp_ms
@@ -141,7 +144,11 @@ async def _start_test_cluster(
             SupervisorActor, address=address, uid=SupervisorActor.uid()
         )
         await start_worker_components(
-            address=address, supervisor_address=address, main_pool=pool
+            address=address,
+            supervisor_address=address,
+            main_pool=pool,
+            metrics_exporter_host=None,
+            metrics_exporter_port=None,
         )
         await pool.join()
     except asyncio.CancelledError:
@@ -233,3 +240,58 @@ def setup_with_file_logging():
 
     local_cluster_proc.terminate()
     restful_api_proc.terminate()
+
+
+@pytest.fixture
+def setup_with_auth():
+    from .api.restful_api import run_in_subprocess as run_restful_api
+    from .deploy.utils import health_check as cluster_health_check
+
+    logging.config.dictConfig(TEST_LOGGING_CONF)  # type: ignore
+
+    supervisor_addr = f"localhost:{xo.utils.get_next_port()}"
+    local_cluster_proc = run_test_cluster_in_subprocess(
+        supervisor_addr, TEST_LOGGING_CONF
+    )
+    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"])
+    user3 = User(
+        username="user3",
+        password="pass3",
+        permissions=["models:list", "models:read", "models:start"],
+    )
+    auth_config = AuthConfig(
+        algorithm="HS256",
+        secret_key="09d25e094faa6ca2556c818166b7a9563b93f7099f6f0f4caa6cf63b88e8d3e7",
+        token_expire_in_minutes=30,
+    )
+    startup_config = AuthStartupConfig(
+        auth_config=auth_config, user_config=[user1, user2, user3]
+    )
+    _, auth_file = tempfile.mkstemp()
+    with open(auth_file, "w") as fd:
+        fd.write(json.dumps(startup_config.dict()))
+
+    port = xo.utils.get_next_port()
+    restful_api_proc = run_restful_api(
+        supervisor_addr,
+        host="localhost",
+        port=port,
+        logging_conf=TEST_LOGGING_CONF,
+        auth_config_file=auth_file,
+    )
+    endpoint = f"http://localhost:{port}"
+    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
+
+    local_cluster_proc.terminate()
+    restful_api_proc.terminate()
+    try:
+        os.remove(auth_file)
+    except:
+        pass

xinference/constants.py

@@ -39,6 +39,7 @@ XINFERENCE_CACHE_DIR = os.path.join(XINFERENCE_HOME, "cache")
 XINFERENCE_MODEL_DIR = os.path.join(XINFERENCE_HOME, "model")
 XINFERENCE_LOG_DIR = os.path.join(XINFERENCE_HOME, "logs")
 XINFERENCE_IMAGE_DIR = os.path.join(XINFERENCE_HOME, "image")
+XINFERENCE_AUTH_DIR = os.path.join(XINFERENCE_HOME, "auth")
 
 XINFERENCE_DEFAULT_LOCAL_HOST = "127.0.0.1"
 XINFERENCE_DEFAULT_DISTRIBUTED_HOST = "0.0.0.0"