enhance(client): new impl of model serving (#2978)

client/Makefile

@@ -9,6 +9,7 @@ check:
 clean:
 	rm -rf dist/*
 
+SKIP_UI_BUILD ?= 0
 build-ui:
 	[ "${SKIP_UI_BUILD}" = "1" ] || make -C ../console install-dev-tools
 	[ "${SKIP_UI_BUILD}" = "1" ] || make -C ../console build-all

client/setup.py

@@ -36,6 +36,7 @@
     "GitPython>=3.1.24",
     "filelock",
     "fastapi",
+    "orjson",  # for web server, e.g. sw job info --web
     "protobuf>=3.19.0",
     "types-protobuf>=3.19.0",
     "lz4>=3.1.10",
@@ -51,8 +52,7 @@
 extras_require = {
     "image": ["pillow"],
     "audio": ["soundfile"],
-    "serve": ["gradio", "uvicorn"],
-    "online-serve": ["gradio~=3.15.0"],
+    "serve": ["uvicorn"],
 }
 
 all_requires = list(

client/starwhale/api/_impl/evaluation/pipeline.py

@@ -16,7 +16,7 @@
 from starwhale.consts import RunStatus, CURRENT_FNAME
 from starwhale.utils.fs import ensure_dir, ensure_file
 from starwhale.base.type import RunSubDirType, PredictLogMode
-from starwhale.api.service import Input, Output, Service
+from starwhale.api.service import Service
 from starwhale.utils.error import ParameterError, FieldTypeOrValueError
 from starwhale.base.context import Context
 from starwhale.core.job.store import JobStorage
@@ -389,11 +389,6 @@ def _update_status(self, status: str) -> None:
         fpath = self.status_dir / CURRENT_FNAME
         ensure_file(fpath, status)
 
-    def add_api(
-        self, input: Input, output: Output, func: t.Callable, name: str
-    ) -> None:
-        self.svc.add_api(input, output, func, name)
-
     def serve(self, addr: str, port: int) -> None:
         self.svc.serve(addr, port)
 

client/starwhale/api/_impl/service/service.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import os
+import typing as t
+import functools
+
+import pkg_resources
+from fastapi import FastAPI
+from pydantic import BaseModel
+from starlette.responses import FileResponse
+from starlette.staticfiles import StaticFiles
+
+from starwhale.base.models.base import SwBaseModel
+
+from .types import ServiceType
+
+STATIC_DIR_DEV = os.getenv("SW_SERVE_STATIC_DIR") or pkg_resources.resource_filename(
+    "starwhale", "web/ui"
+)
+
+
+class Query(BaseModel):
+    content: str
+
+
+class Api(SwBaseModel):
+    func: t.Callable
+    uri: str
+    inference_type: ServiceType
+
+    @staticmethod
+    def question_answering(func: t.Callable) -> t.Callable:
+        def inter(query: Query) -> str:
+            return func(query.content)  # type: ignore
+
+        return inter
+
+    def view_func(self, ins: t.Any = None) -> t.Callable:
+        func = self.func
+        if ins is not None:
+            func = functools.partial(func, ins)
+        return getattr(self, self.inference_type.value)(func)  # type: ignore
+
+
+class ServiceSpec(SwBaseModel):
+    title: t.Optional[str]
+    description: t.Optional[str]
+    version: str
+    apis: t.List[Api]
+
+
+class Service:
+    def __init__(self) -> None:
+        self.apis: t.Dict[str, Api] = {}
+        self.api_within_instance_map: t.Dict[str, t.Any] = {}
+
+    def api(self, inference_type: ServiceType) -> t.Any:
+        def decorator(func: t.Any) -> t.Any:
+            self.add_api(func, func.__name__, inference_type=inference_type)
+            return func
+
+        return decorator
+
+    def get_spec(self) -> ServiceSpec:
+        return ServiceSpec(version="0.0.1", apis=list(self.apis.values()))
+
+    def add_api(self, func: t.Callable, uri: str, inference_type: ServiceType) -> None:
+        if uri in self.apis:
+            raise ValueError(f"Duplicate api uri: {uri}")
+
+        _api = Api(func=func, uri=uri, inference_type=inference_type)
+        self.apis[uri] = _api
+
+    def add_api_instance(self, _api: Api) -> None:
+        self.apis[_api.uri] = _api
+
+    def serve(self, addr: str, port: int, title: t.Optional[str] = None) -> None:
+        """
+        Default serve implementation, users can override this method
+        :param addr
+        :param port
+        :param title webpage title
+        :return: None
+        """
+        app = FastAPI(title=title or "Starwhale Model Serving")
+
+        @app.get("/api/spec")
+        def spec() -> ServiceSpec:
+            return self.get_spec()
+
+        for _api in self.apis.values():
+            app.add_api_route(
+                f"/api/{_api.uri}",
+                _api.view_func(self.api_within_instance_map.get(_api.uri)),
+                methods=["POST"],
+            )
+
+        def index(opt: t.Any) -> FileResponse:
+            return FileResponse(os.path.join(STATIC_DIR_DEV, "client/index.html"))
+
+        app.add_route("/", index, methods=["GET"])
+        app.mount("/", StaticFiles(directory=STATIC_DIR_DEV), name="assets")
+
+        import uvicorn
+
+        uvicorn.run(app, host=addr, port=port)
+
+
+_svc = Service()
+
+
+def api(inference_type: ServiceType) -> t.Any:
+    return _svc.api(inference_type=inference_type)
+
+
+def internal_api_list() -> t.Dict[str, Api]:
+    return _svc.apis

client/starwhale/api/_impl/service/types.py

@@ -0,0 +1,11 @@
+from enum import Enum
+
+
+class ServiceType(Enum):
+    """Enumeration of service types."""
+
+    TEXT_TO_TEXT = "text_to_text"
+    TEXT_TO_IMAGE = "text_to_image"
+    TEXT_TO_AUDIO = "text_to_audio"
+    TEXT_TO_VIDEO = "text_to_video"
+    QUESTION_ANSWERING = "question_answering"

client/starwhale/api/service.py

@@ -1,9 +1,7 @@
-from ._impl.service import Api, api, Input, Output, Service
+from starwhale.api._impl.service.service import api, Service, ServiceType
 
 __all__ = [
     "Service",
-    "Api",
     "api",
-    "Input",
-    "Output",
+    "ServiceType",
 ]