Move OpenAPI schemas into Python package (#1095)

Co-authored-by: adriangonz <adriangonz@users.noreply.github.com>

Dockerfile

@@ -3,9 +3,11 @@ SHELL ["/bin/bash", "-l", "-c"]
 
 COPY ./hack/build-wheels.sh ./hack/build-wheels.sh
 COPY ./mlserver ./mlserver
+COPY ./openapi ./openapi
 COPY ./runtimes ./runtimes
 COPY \
     setup.py \
+    MANIFEST.in \
     README.md \
     .
 

MANIFEST.in

@@ -1 +1 @@
-include openapi/*.json
+include mlserver/rest/openapi/*.json

mlserver/rest/openapi/__init__.py

@@ -0,0 +1,3 @@
+from .schema import get_openapi_schema, get_model_schema, get_model_schema_uri
+
+__all__ = ["get_openapi_schema", "get_model_schema", "get_model_schema_uri"]

mlserver/rest/openapi/dataplane.json

@@ -0,0 +1 @@
+../../../openapi/dataplane.json

mlserver/rest/openapi/model_repository.json

@@ -0,0 +1 @@
+../../../openapi/model_repository.json

mlserver/rest/openapi.py → mlserver/rest/openapi/schema.py

@@ -4,15 +4,13 @@
 from typing import Optional, Tuple
 from importlib_resources import files
 
-OPENAPI_SCHEMA_RELATIVE_PATH = "../openapi/dataplane.json"
 MODEL_NAME_PARAMETER = "model_name"
 MODEL_VERSION_PARAMETER = "model_version"
 
 
 @lru_cache
 def get_openapi_schema() -> dict:
-    mlserver_package = __package__.split(".")[0]
-    openapi_schema_path = files(mlserver_package).joinpath(OPENAPI_SCHEMA_RELATIVE_PATH)
+    openapi_schema_path = files(__package__).joinpath("dataplane.json")
     return orjson.loads(openapi_schema_path.read_bytes())
 
 

tests/rest/test_openapi.py → tests/rest/openapi/test_schema.py

@@ -2,7 +2,7 @@
 
 from typing import Optional
 
-from mlserver.rest.openapi import (
+from mlserver.rest.openapi.schema import (
     get_openapi_schema,
     get_model_schema,
     MODEL_VERSION_PARAMETER,