Repository: xorbitsai/inference

Issue:
ENH: RESTful client supports custom model APIs
### Is your feature request related to a problem? Please describe
Including `register_model`, `unregister_model`, and `list_model_registrations`.




Discussion:

Resolving patch:
diff --git a/xinference/client.py b/xinference/client.py
--- a/xinference/client.py
+++ b/xinference/client.py
@@ -481,20 +481,69 @@ def _gen_model_uid(cls) -> str:
         return str(uuid.uuid1())
 
     def register_model(self, model_type: str, model: str, persist: bool):
+        """
+        Register a custom model.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of model.
+        model: str
+            The model definition. (refer to: https://inference.readthedocs.io/en/latest/models/custom.html)
+        persist: bool
+        """
         coro = self._supervisor_ref.register_model(model_type, model, persist)
         self._isolation.call(coro)
 
     def unregister_model(self, model_type: str, model_name: str):
+        """
+        Unregister a custom model.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of model.
+        model_name: str
+            The name of the model
+        """
         coro = self._supervisor_ref.unregister_model(model_type, model_name)
         self._isolation.call(coro)
 
     def list_model_registrations(self, model_type: str) -> List[Dict[str, Any]]:
+        """
+        List models registered on the server.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of the model.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            The collection of registered models on the server.
+        """
         coro = self._supervisor_ref.list_model_registrations(model_type)
         return self._isolation.call(coro)
 
     def get_model_registration(
         self, model_type: str, model_name: str
     ) -> Dict[str, Any]:
+        """
+        Get the model with the model type and model name registered on the server.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of the model.
+
+        model_name: str
+            The name of the model.
+        Returns
+        -------
+        List[Dict[str, Any]]
+            The collection of registered models on the server.
+        """
         coro = self._supervisor_ref.get_model_registration(model_type, model_name)
         return self._isolation.call(coro)
 
@@ -763,3 +812,163 @@ def get_model(self, model_uid: str) -> RESTfulModelHandle:
             return RESTfulChatModelHandle(model_uid, self.base_url)
         else:
             raise ValueError(f"Unrecognized model ability: {desc['model_ability']}")
+
+    def describe_model(self, model_uid: str):
+        """
+        Get model information via RESTful APIs.
+
+        Parameters
+        ----------
+        model_uid: str
+            The unique id that identify the model.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the following keys:
+            - "model_type": str
+               the type of the model determined by its function, e.g. "LLM" (Large Language Model)
+            - "model_name": str
+               the name of the specific LLM model family
+            - "model_lang": List[str]
+               the languages supported by the LLM model
+            - "model_ability": List[str]
+               the ability or capabilities of the LLM model
+            - "model_description": str
+               a detailed description of the LLM model
+            - "model_format": str
+               the format specification of the LLM model
+            - "model_size_in_billions": int
+               the size of the LLM model in billions
+            - "quantization": str
+               the quantization applied to the model
+            - "revision": str
+               the revision number of the LLM model specification
+
+        Raises
+        ------
+        RuntimeError
+            Report failure to get the wanted model with given model_uid. Provide details of failure through error message.
+
+        """
+
+        url = f"{self.base_url}/v1/models/{model_uid}"
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get the model description, detail: {response.json()['detail']}"
+            )
+        return response.json()
+
+    def register_model(self, model_type: str, model: str, persist: bool):
+        """
+        Register a custom model.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of model.
+        model: str
+            The model definition. (refer to: https://inference.readthedocs.io/en/latest/models/custom.html)
+        persist: bool
+
+
+        Raises
+        ------
+        RuntimeError
+            Report failure to register the custom model. Provide details of failure through error message.
+        """
+        url = f"{self.base_url}/v1/model_registrations/{model_type}"
+        request_body = {"model": model, "persist": persist}
+        response = requests.post(url, json=request_body)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to register model, detail: {response.json()['detail']}"
+            )
+
+        response_data = response.json()
+        return response_data
+
+    def unregister_model(self, model_type: str, model_name: str):
+        """
+        Unregister a custom model.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of model.
+        model_name: str
+            The name of the model
+
+        Raises
+        ------
+        RuntimeError
+            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)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to register model, detail: {response.json()['detail']}"
+            )
+
+        response_data = response.json()
+        return response_data
+
+    def list_model_registrations(self, model_type: str) -> List[Dict[str, Any]]:
+        """
+        List models registered on the server.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of the model.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            The collection of registered models on the server.
+
+        Raises
+        ------
+        RuntimeError
+            Report failure to list model registration. Provide details of failure through error message.
+
+        """
+        url = f"{self.base_url}/v1/model_registrations/{model_type}"
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to list model registration, detail: {response.json()['detail']}"
+            )
+
+        response_data = response.json()
+        return response_data
+
+    def get_model_registration(
+        self, model_type: str, model_name: str
+    ) -> Dict[str, Any]:
+        """
+        Get the model with the model type and model name registered on the server.
+
+        Parameters
+        ----------
+        model_type: str
+            The type of the model.
+
+        model_name: str
+            The name of the model.
+        Returns
+        -------
+        List[Dict[str, Any]]
+            The collection of registered models on the server.
+        """
+        url = f"{self.base_url}/v1/model_registrations/{model_type}/{model_name}"
+        response = requests.get(url)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to list model registration, detail: {response.json()['detail']}"
+            )
+
+        response_data = response.json()
+        return response_data
diff --git a/xinference/core/restful_api.py b/xinference/core/restful_api.py
--- a/xinference/core/restful_api.py
+++ b/xinference/core/restful_api.py
@@ -212,6 +212,11 @@ class Config:
         }
 
 
+class RegisterModelRequest(BaseModel):
+    model: str
+    persist: bool
+
+
 class RESTfulAPIActor(xo.Actor):
     def __init__(self, sockets: List[socket.socket], gradio_block: gr.Blocks):
         super().__init__()
@@ -266,6 +271,29 @@ def serve(self):
             methods=["POST"],
             response_model=ChatCompletion,
         )
+
+        # for custom models
+        self._router.add_api_route(
+            "/v1/model_registrations/{model_type}",
+            self.register_model,
+            methods=["POST"],
+        )
+        self._router.add_api_route(
+            "/v1/model_registrations/{model_type}/{model_name}",
+            self.unregister_model,
+            methods=["DELETE"],
+        )
+        self._router.add_api_route(
+            "/v1/model_registrations/{model_type}",
+            self.list_model_registrations,
+            methods=["GET"],
+        )
+        self._router.add_api_route(
+            "/v1/model_registrations/{model_type}/{model_name}",
+            self.get_model_registrations,
+            methods=["GET"],
+        )
+
         app.include_router(self._router)
         app = gr.mount_gradio_app(app, self._gradio_block, path="/")
 
@@ -555,3 +583,50 @@ async def event_publisher(inner_send_chan: MemoryObjectSendStream):
             except Exception as e:
                 logger.error(e, exc_info=True)
                 raise HTTPException(status_code=500, detail=str(e))
+
+    async def register_model(self, model_type: str, request: RegisterModelRequest):
+        model = request.model
+        persist = request.persist
+
+        try:
+            await self._supervisor_ref.register_model(model_type, model, persist)
+        except ValueError as re:
+            logger.error(re, exc_info=True)
+            raise HTTPException(status_code=400, detail=str(re))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
+    async def unregister_model(self, model_type: str, model_name: str):
+        try:
+            await self._supervisor_ref.unregister_model(model_type, model_name)
+        except ValueError as re:
+            logger.error(re, exc_info=True)
+            raise HTTPException(status_code=400, detail=str(re))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
+    async def list_model_registrations(self, model_type: str) -> List[Dict[str, Any]]:
+        try:
+            return await self._supervisor_ref.list_model_registrations(model_type)
+        except ValueError as re:
+            logger.error(re, exc_info=True)
+            raise HTTPException(status_code=400, detail=str(re))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
+    async def get_model_registrations(
+        self, model_type: str, model_name: str
+    ) -> Dict[str, Any]:
+        try:
+            return await self._supervisor_ref.get_model_registration(
+                model_type, model_name
+            )
+        except ValueError as re:
+            logger.error(re, exc_info=True)
+            raise HTTPException(status_code=400, detail=str(re))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))


Test patch:
diff --git a/xinference/core/tests/test_restful_api.py b/xinference/core/tests/test_restful_api.py
--- a/xinference/core/tests/test_restful_api.py
+++ b/xinference/core/tests/test_restful_api.py
@@ -176,3 +176,90 @@ async def test_restful_api(setup):
 
     url = f"{endpoint}/v1/models/test_restful_api2"
     response = requests.delete(url)
+
+    # list model registration
+
+    url = f"{endpoint}/v1/model_registrations/LLM"
+
+    response = requests.get(url)
+
+    assert response.status_code == 200
+    model_regs = response.json()
+    assert len(model_regs) > 0
+    for model_reg in model_regs:
+        assert model_reg["is_builtin"]
+
+    # register_model
+
+    model = """{
+  "version": 1,
+  "model_name": "custom_model",
+  "model_lang": [
+    "en", "zh"
+  ],
+  "model_ability": [
+    "embed",
+    "chat"
+  ],
+  "model_specs": [
+    {
+      "model_format": "pytorch",
+      "model_size_in_billions": 7,
+      "quantizations": [
+        "4-bit",
+        "8-bit",
+        "none"
+      ],
+      "model_id": "ziqingyang/chinese-alpaca-2-7b"
+    }
+  ],
+  "prompt_style": {
+    "style_name": "ADD_COLON_SINGLE",
+    "system_prompt": "Below is an instruction that describes a task. Write a response that appropriately completes the request.",
+    "roles": [
+      "Instruction",
+      "Response"
+    ],
+    "intra_message_sep": "\\n\\n### "
+  }
+}"""
+
+    url = f"{endpoint}/v1/model_registrations/LLM"
+
+    payload = {"model": model, "persist": False}
+
+    response = requests.post(url, json=payload)
+    assert response.status_code == 200
+
+    url = f"{endpoint}/v1/model_registrations/LLM"
+
+    response = requests.get(url)
+
+    assert response.status_code == 200
+    new_model_regs = response.json()
+    assert len(new_model_regs) == len(model_regs) + 1
+
+    # get_model_registrations
+    url = f"{endpoint}/v1/model_registrations/LLM/custom_model"
+    response = requests.get(url, json=payload)
+    assert response.status_code == 200
+    data = response.json()
+    assert "custom_model" in data["model_name"]
+
+    # unregister_model
+    url = f"{endpoint}/v1/model_registrations/LLM/custom_model"
+
+    response = requests.delete(url, json=payload)
+    assert response.status_code == 200
+
+    url = f"{endpoint}/v1/model_registrations/LLM"
+
+    response = requests.get(url)
+    assert response.status_code == 200
+    new_model_regs = response.json()
+    assert len(new_model_regs) == len(model_regs)
+    custom_model_reg = None
+    for model_reg in new_model_regs:
+        if model_reg["model_name"] == "custom_model":
+            custom_model_reg = model_reg
+    assert custom_model_reg is None
diff --git a/xinference/tests/test_client.py b/xinference/tests/test_client.py
--- a/xinference/tests/test_client.py
+++ b/xinference/tests/test_client.py
@@ -180,3 +180,67 @@ def test_RESTful_client(setup):
 
     client.terminate_model(model_uid=model_uid2)
     assert len(client.list_models()) == 0
+
+
+def test_RESTful_client_custom_model(setup):
+    endpoint, _ = setup
+    client = RESTfulClient(endpoint)
+
+    model_regs = client.list_model_registrations(model_type="LLM")
+    assert len(model_regs) > 0
+    for model_reg in model_regs:
+        assert model_reg["is_builtin"]
+
+    model = """{
+  "version": 1,
+  "model_name": "custom_model",
+  "model_lang": [
+    "en", "zh"
+  ],
+  "model_ability": [
+    "embed",
+    "chat"
+  ],
+  "model_specs": [
+    {
+      "model_format": "pytorch",
+      "model_size_in_billions": 7,
+      "quantizations": [
+        "4-bit",
+        "8-bit",
+        "none"
+      ],
+      "model_id": "ziqingyang/chinese-alpaca-2-7b"
+    }
+  ],
+  "prompt_style": {
+    "style_name": "ADD_COLON_SINGLE",
+    "system_prompt": "Below is an instruction that describes a task. Write a response that appropriately completes the request.",
+    "roles": [
+      "Instruction",
+      "Response"
+    ],
+    "intra_message_sep": "\\n\\n### "
+  }
+}"""
+    client.register_model(model_type="LLM", model=model, persist=False)
+
+    data = client.get_model_registration(model_type="LLM", model_name="custom_model")
+    assert "custom_model" in data["model_name"]
+
+    new_model_regs = client.list_model_registrations(model_type="LLM")
+    assert len(new_model_regs) == len(model_regs) + 1
+    custom_model_reg = None
+    for model_reg in new_model_regs:
+        if model_reg["model_name"] == "custom_model":
+            custom_model_reg = model_reg
+    assert custom_model_reg is not None
+
+    client.unregister_model(model_type="LLM", model_name="custom_model")
+    new_model_regs = client.list_model_registrations(model_type="LLM")
+    assert len(new_model_regs) == len(model_regs)
+    custom_model_reg = None
+    for model_reg in new_model_regs:
+        if model_reg["model_name"] == "custom_model":
+            custom_model_reg = model_reg
+    assert custom_model_reg is None