added support of inputs parameters that are recognised by the API.

Models are now loaded in separate endpoints for the inputs to be easier to recognise

source/model/DummyModel.py

@@ -1,19 +0,0 @@
-import json
-
-from source.model import base
-
-
-class DummyModel(base.BaseModel):
-    """
-    A dummy model, mainly used to test the API and the manager.
-    simply send back the request made to it.
-    """
-
-    def _load(self) -> None:
-        pass
-
-    def _unload(self) -> None:
-        pass
-
-    def _infer(self, payload: dict) -> str | bytes:
-        return json.dumps(payload)

source/model/PythonModel.py

@@ -1,9 +1,14 @@
 import importlib.util
 import subprocess
 import sys
+import typing
 import uuid
+import inspect
 from pathlib import Path
 
+import fastapi
+
+from source import utils
 from source.manager import ModelManager
 from source.model import base
 
@@ -16,6 +21,8 @@ class PythonModel(base.BaseModel):
     def __init__(self, manager: ModelManager, configuration: dict, path: Path):
         super().__init__(manager, configuration, path)
 
+        ## Configuration
+
         # get the name of the file containing the model code
         file = configuration.get("file")
         if file is None:
@@ -36,11 +43,28 @@ class PythonModel(base.BaseModel):
         # load the module
         module_spec.loader.exec_module(self.module)
 
+        ## Api
+
+        # load the inputs data into the inference function signature (used by FastAPI)
+        parameters = utils.parameters.load(configuration.get("inputs", {}))
+
+        # create an endpoint wrapping the inference inside a fastapi call
+        async def infer_api(*args, **kwargs):
+            return fastapi.responses.StreamingResponse(
+                content=self.infer(*args, **kwargs),
+                media_type=self.output_type,
+            )
+
+        infer_api.__signature__ = inspect.Signature(parameters=parameters)
+
+        # add the inference endpoint on the API
+        self.manager.application.add_api_route(f"/models/{self.name}/infer", infer_api, methods=["POST"])
+
     def _load(self) -> None:
         return self.module.load(self)
 
     def _unload(self) -> None:
         return self.module.unload(self)
 
-    def _infer(self, payload: dict) -> str | bytes:
-        return self.module.infer(self, payload)
+    def _infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
+        return self.module.infer(self, *args, **kwargs)

source/model/__init__.py

@@ -1,4 +1,3 @@
 from . import base
 
-from .DummyModel import DummyModel
 from .PythonModel import PythonModel

source/model/base/BaseModel.py

@@ -1,7 +1,9 @@
 import abc
 import gc
+import typing
 from pathlib import Path
 
+from source import api
 from source.manager import ModelManager
 
 
@@ -10,13 +12,13 @@ class BaseModel(abc.ABC):
     Represent a model.
     """
 
-    def __init__(self, manager: ModelManager, configuration: dict, path: Path):
+    def __init__(self, manager: ModelManager, configuration: dict[str, typing.Any], path: Path):
         # the environment directory of the model
         self.path = path
         # the model manager
         self.manager = manager
         # the mimetype of the model responses
-        self.response_mimetype: str = configuration.get("response_mimetype", "application/json")
+        self.output_type: str = configuration.get("output_type", "application/json")
 
         self._loaded = False
 
@@ -101,13 +103,11 @@ class BaseModel(abc.ABC):
         """
         Unload the model
         Do not call manually, use `unload` instead.
-        :return:
         """
 
-    def infer(self, payload: dict) -> str | bytes:
+    def infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
         """
         Infer our payload through the model within the model manager
-        :param payload: the payload to give to the model
         :return: the response of the model
         """
 
@@ -115,12 +115,11 @@ class BaseModel(abc.ABC):
         self.load()
 
         # model specific inference part
-        return self._infer(payload)
+        return self._infer(*args, **kwargs)
 
     @abc.abstractmethod
-    def _infer(self, payload: dict) -> str | bytes:
+    def _infer(self, *args, **kwargs) -> typing.Iterator[bytes]:
         """
         Infer our payload through the model
-        :param payload: the payload to give to the model
         :return: the response of the model
         """