feat: provide openapi doc

Cargo.toml

@@ -24,3 +24,6 @@ once_cell = "1.17"
 prometheus-client = "0.21.1"
 argh = "0.1"
 axum = "0.6.18"
+utoipa = "3.3.0"
+serde_json = "1.0.96"
+serde = "1.0.163"

mosec/mixin/typed_worker.py

@@ -27,32 +27,54 @@
     warnings.warn("msgpack is required for TypedMsgPackMixin", ImportWarning)
 
 
-def parse_forward_input_type(func):
+def parse_forward_input_type(func, target="parameters", index=0):
     """Parse the input type of the forward function.
 
     - single request: return the type
     - batch request: return the list item type
     """
     sig = inspect.signature(func)
-    params = list(sig.parameters.values())
-    if len(params) < 1:
-        raise TypeError("`forward` method doesn't have enough(1) parameters")
+    if target == "parameters":
+        params = list(sig.parameters.values())
+        if len(params) < 1:
+            raise TypeError("`forward` method doesn't have enough(1) parameters")
+        typ = params[index].annotation
+    else:
+        typ = sig.return_annotation
+        if typ is inspect.Signature.empty:
+            raise TypeError("`forward` method doesn't have return annotation")
+        if not inspect.isclass(typ):
+            typ = msgspec.inspect.type_info(typ).__class__
 
-    typ = params[0].annotation
     origin = getattr(typ, "__origin__", None)
     if origin is None:
         return typ
     # GenericAlias, `func` could be batch inference
     if origin is list or origin is List:
-        if not hasattr(typ, "__args__") or len(typ.__args__) != 1:
+        if not hasattr(typ, "__args__") or len(typ.__args__) != 1:  # type: ignore
             raise TypeError(
                 "`forward` with dynamic batch should use "
                 "`List[Struct]` as the input annotation"
             )
-        return typ.__args__[0]
+        return typ.__args__[0]  # type: ignore
     raise TypeError(f"unsupported type {typ}")
 
 
+def parse_instance_param_typ(func):
+    """Parse the input type of the forward function for instance."""
+    return parse_forward_input_type(func, "parameters", 0)
+
+
+def parse_cls_param_typ(func):
+    """Parse the input type of the forward function for class."""
+    return parse_forward_input_type(func, "parameters", 1)
+
+
+def parse_cls_return_typ(func):
+    """Parse the return type of the forward function for class."""
+    return parse_forward_input_type(func, "return")
+
+
 class TypedMsgPackMixin(Worker):
     """Enable request type validation with `msgspec` and serde with `msgpack`."""
 
@@ -64,7 +86,7 @@ class TypedMsgPackMixin(Worker):
     def _get_input_type(self):
         """Get the input type from annotations."""
         if self._input_type is None:
-            self._input_type = parse_forward_input_type(self.forward)
+            self._input_type = parse_instance_param_typ(self.forward)
         return self._input_type
 
     def deserialize(self, data: Any) -> bytes:

mosec/server.py

@@ -40,6 +40,8 @@
 """
 
 import multiprocessing as mp
+import os
+import pathlib
 import shutil
 import signal
 import subprocess
@@ -50,6 +52,7 @@
 from time import monotonic, sleep
 from typing import Dict, List, Optional, Type, Union
 
+import msgspec
 import pkg_resources
 
 from mosec.args import mosec_args
@@ -58,13 +61,20 @@
 from mosec.env import env_var_context
 from mosec.ipc import IPCWrapper
 from mosec.log import get_internal_logger
+from mosec.mixin.typed_worker import (
+    TypedMsgPackMixin,
+    parse_cls_param_typ,
+    parse_cls_return_typ,
+)
 from mosec.worker import Worker
 
 logger = get_internal_logger()
 
 
 GUARD_CHECK_INTERVAL = 1
 NEW_PROCESS_METHOD = {"spawn", "fork"}
+MOSEC_REF_TEMPLATE = "#/components/schemas/{name}"
+MOSEC_OPENAPI_PATH = "mosec_openapi.json"
 
 
 class Server:
@@ -363,6 +373,34 @@ def append_worker(
         self._coordinator_ctx.append(start_method)
         self._coordinator_pools.append([None] * num)
 
+    def _generate_openapi(self):
+        """Generate the OpenAPI specification."""
+        if not self._worker_cls:
+            return
+        req_w, res_w = self._worker_cls[0], self._worker_cls[-1]
+        if not (
+            issubclass(req_w, TypedMsgPackMixin)
+            and issubclass(res_w, TypedMsgPackMixin)
+        ):
+            return
+        req_typ, res_typ = parse_cls_param_typ(req_w.forward), parse_cls_return_typ(
+            res_w.forward
+        )
+        schema, schemas = msgspec.json.schema_components(
+            [req_typ, res_typ], MOSEC_REF_TEMPLATE
+        )
+        req_schema, res_schema = schema
+        schema = {
+            "req_schema": req_schema,
+            "res_schema": res_schema,
+            "schemas": schemas,
+        }
+        tmp_path = pathlib.Path(os.path.join(self._configs["path"], MOSEC_OPENAPI_PATH))
+        tmp_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path.write_text(
+            msgspec.json.encode(schema).decode("utf-8"), encoding="utf-8"
+        )
+
     def run(self):
         """Start the mosec model server."""
         self._validate_server()
@@ -375,6 +413,7 @@ def run(self):
             return
 
         self._handle_signal()
+        self._generate_openapi()
         self._start_controller()
         try:
             self._manage_coordinators()

src/apidoc.rs

@@ -0,0 +1,110 @@
+// Copyright 2023 MOSEC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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.
+
+use std::{collections::BTreeMap, str::FromStr};
+
+use hyper::StatusCode;
+use serde::Deserialize;
+use utoipa::openapi::{
+    path::Operation, request_body::RequestBody, Components, Content, ContentBuilder, OpenApi,
+    PathItemType, RefOr, ResponseBuilder, Responses, Schema,
+};
+
+#[derive(Deserialize, Default)]
+pub(crate) struct InferenceSchemas {
+    req_schema: RefOr<Schema>,
+    res_schema: RefOr<Schema>,
+    schemas: BTreeMap<String, RefOr<Schema>>,
+}
+
+impl FromStr for InferenceSchemas {
+    type Err = serde_json::Error;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        serde_json::from_str::<InferenceSchemas>(s)
+    }
+}
+
+pub(crate) struct MosecApiDoc {
+    pub rust_api: OpenApi,
+}
+
+impl MosecApiDoc {
+    fn get_operation<'a>(
+        &self,
+        api: &'a mut OpenApi,
+        route: &str,
+        method: &PathItemType,
+    ) -> Option<&'a mut Operation> {
+        let path = api.paths.paths.get_mut(route).unwrap();
+        path.operations.get_mut(method)
+    }
+
+    fn get_route_request_body<'a>(
+        &self,
+        api: &'a mut OpenApi,
+        route: &str,
+        method: &PathItemType,
+    ) -> Option<&'a mut RequestBody> {
+        let op = self.get_operation(api, route, method).unwrap();
+        op.request_body.as_mut()
+    }
+
+    fn get_route_responses<'a>(
+        &self,
+        api: &'a mut OpenApi,
+        route: &str,
+        method: &PathItemType,
+    ) -> &'a mut Responses {
+        let op = self.get_operation(api, route, method).unwrap();
+        &mut op.responses
+    }
+
+    fn merge_schemas(&self, api: &mut OpenApi, mut other_schemas: BTreeMap<String, RefOr<Schema>>) {
+        if api.components.is_none() {
+            api.components = Some(Components::default());
+        }
+        let schemas = &mut api.components.as_mut().unwrap().schemas;
+        schemas.append(&mut other_schemas);
+    }
+
+    fn merge_request(&self, req_body: &mut RequestBody, req_schema: RefOr<Schema>) {
+        let content = ContentBuilder::new().schema(req_schema).build();
+        req_body
+            .content
+            .insert("application/json".to_string(), content);
+    }
+
+    fn merge_response(&self, response: &mut Responses, res_schema: RefOr<Schema>) {
+        let ok_res = ResponseBuilder::new()
+            .content("application/json", Content::new(res_schema))
+            .build();
+        response
+            .responses
+            .insert(StatusCode::OK.as_str().to_string(), RefOr::from(ok_res));
+    }
+
+    pub fn merge(&self, python_schema: InferenceSchemas) -> OpenApi {
+        let mut api = self.rust_api.clone();
+        self.merge_schemas(&mut api, python_schema.schemas.clone());
+
+        let req_body = self
+            .get_route_request_body(&mut api, "/inference", &PathItemType::Post)
+            .unwrap();
+        self.merge_request(req_body, python_schema.req_schema);
+
+        let response = self.get_route_responses(&mut api, "/inference", &PathItemType::Post);
+        self.merge_response(response, python_schema.res_schema);
+        api
+    }
+}