Phil/upload schemas (#68)

* update client

* try again

* fix import

* fixes

* wip

* fixes

* try fix again

* fix

* fixes

* add pydantic schema support

* add test

* fix

* fix flat schemas

* black and isort

Author: phil-scale

.pylintrc

@@ -28,3 +28,4 @@ max-line-length=79
 [MASTER]
 # Ignore anything inside launch/clientlib (since it's documentation)
 ignore=clientlib,api_client
+extension-pkg-whitelist=pydantic

launch/client.py

@@ -5,13 +5,14 @@
 import shutil
 import tempfile
 from io import StringIO
-from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
+from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
 from zipfile import ZipFile
 
 import cloudpickle
 import requests
 import yaml
 from frozendict import frozendict
+from pydantic import BaseModel
 
 from launch.api_client import ApiClient, Configuration
 from launch.api_client.apis.tags.default_api import DefaultApi
@@ -59,6 +60,7 @@
     ModelEndpoint,
     SyncEndpoint,
 )
+from launch.pydantic_schemas import get_model_definitions
 from launch.request_validation import validate_task_request
 
 DEFAULT_NETWORK_TIMEOUT_SEC = 120
@@ -224,15 +226,39 @@ def register_endpoint_auth_decorator(self, endpoint_auth_decorator_fn):
         """
         self.endpoint_auth_decorator_fn = endpoint_auth_decorator_fn
 
+    def _upload_data(self, data: bytes) -> str:
+        if self.self_hosted:
+            if self.upload_bundle_fn is None:
+                raise ValueError("Upload_bundle_fn should be registered")
+            if self.bundle_location_fn is None:
+                raise ValueError(
+                    "Need either bundle_location_fn to know where to upload bundles"
+                )
+            raw_bundle_url = self.bundle_location_fn()  # type: ignore
+            self.upload_bundle_fn(data, raw_bundle_url)  # type: ignore
+        else:
+            model_bundle_url = self.connection.post(
+                {}, MODEL_BUNDLE_SIGNED_URL_PATH
+            )
+            s3_path = model_bundle_url["signedUrl"]
+            raw_bundle_url = (
+                f"s3://{model_bundle_url['bucket']}/{model_bundle_url['key']}"
+            )
+            requests.put(s3_path, data=data)
+        return raw_bundle_url
+
     def create_model_bundle_from_dirs(
         self,
+        *,
         model_bundle_name: str,
         base_paths: List[str],
         requirements_path: str,
         env_params: Dict[str, str],
         load_predict_fn_module_path: str,
         load_model_fn_module_path: str,
         app_config: Optional[Union[Dict[str, Any], str]] = None,
+        request_schema: Optional[Type[BaseModel]] = None,
+        response_schema: Optional[Type[BaseModel]] = None,
     ) -> ModelBundle:
         """
         Packages up code from one or more local filesystem folders and uploads them as a bundle to Scale Launch.
@@ -302,6 +328,13 @@ def create_model_bundle_from_dirs(
                 the function located at load_predict_fn_module_path.
 
             app_config: Either a Dictionary that represents a YAML file contents or a local path to a YAML file.
+
+            request_schema: A pydantic model that represents the request schema for the model
+                bundle. This is used to validate the request body for the model bundle's endpoint.
+
+            response_schema: A pydantic model that represents the request schema for the model
+                bundle. This is used to validate the response for the model bundle's endpoint.
+                Note: If request_schema is specified, then response_schema must also be specified.
         """
         with open(requirements_path, "r", encoding="utf-8") as req_f:
             requirements = req_f.read().splitlines()
@@ -315,24 +348,20 @@ def create_model_bundle_from_dirs(
         finally:
             shutil.rmtree(tmpdir)
 
-        if self.self_hosted:
-            if self.upload_bundle_fn is None:
-                raise ValueError("Upload_bundle_fn should be registered")
-            if self.bundle_location_fn is None:
-                raise ValueError(
-                    "Need either bundle_location_fn to know where to upload bundles"
-                )
-            raw_bundle_url = self.bundle_location_fn()  # type: ignore
-            self.upload_bundle_fn(data, raw_bundle_url)  # type: ignore
-        else:
-            model_bundle_url = self.connection.post(
-                {}, MODEL_BUNDLE_SIGNED_URL_PATH
+        raw_bundle_url = self._upload_data(data)
+
+        schema_location = None
+        if bool(request_schema) ^ bool(response_schema):
+            raise ValueError(
+                "If request_schema is specified, then response_schema must also be specified."
             )
-            s3_path = model_bundle_url["signedUrl"]
-            raw_bundle_url = (
-                f"s3://{model_bundle_url['bucket']}/{model_bundle_url['key']}"
+        if request_schema is not None and response_schema is not None:
+            model_definitions = get_model_definitions(
+                request_schema=request_schema,
+                response_schema=response_schema,
             )
-            requests.put(s3_path, data=data)
+            model_definitions_encoded = json.dumps(model_definitions).encode()
+            schema_location = self._upload_data(model_definitions_encoded)
 
         bundle_metadata = {
             "load_predict_fn_module_path": load_predict_fn_module_path,
@@ -350,6 +379,7 @@ def create_model_bundle_from_dirs(
             bundle_metadata=bundle_metadata,
             requirements=requirements,
             env_params=env_params,
+            schema_location=schema_location,
         )
         _add_app_config_to_bundle_create_payload(payload, app_config)
 
@@ -367,6 +397,7 @@ def create_model_bundle_from_dirs(
                 packaging_type=ModelBundlePackagingType("zip"),
                 metadata=bundle_metadata,
                 app_config=payload.get("app_config"),
+                schema_location=schema_location,
             )
             create_model_bundle_request = CreateModelBundleRequest(**payload)  # type: ignore
             api_instance.create_model_bundle_v1_model_bundles_post(
@@ -390,6 +421,8 @@ def create_model_bundle(  # pylint: disable=too-many-statements
         bundle_url: Optional[str] = None,
         app_config: Optional[Union[Dict[str, Any], str]] = None,
         globals_copy: Optional[Dict[str, Any]] = None,
+        request_schema: Optional[Type[BaseModel]] = None,
+        response_schema: Optional[Type[BaseModel]] = None,
     ) -> ModelBundle:
         """
         Uploads and registers a model bundle to Scale Launch.
@@ -466,6 +499,13 @@ def create_model_bundle(  # pylint: disable=too-many-statements
 
             bundle_url: (Only used in self-hosted mode.) The desired location of bundle.
                 Overrides any value given by ``self.bundle_location_fn``
+
+            request_schema: A pydantic model that represents the request schema for the model
+                bundle. This is used to validate the request body for the model bundle's endpoint.
+
+            response_schema: A pydantic model that represents the request schema for the model
+                bundle. This is used to validate the response for the model bundle's endpoint.
+                Note: If request_schema is specified, then response_schema must also be specified.
         """
         # TODO(ivan): remove `disable=too-many-branches` when get rid of `load_*` functions
         # pylint: disable=too-many-branches
@@ -533,29 +573,20 @@ def create_model_bundle(  # pylint: disable=too-many-statements
             )
 
         serialized_bundle = cloudpickle.dumps(bundle)
+        raw_bundle_url = self._upload_data(data=serialized_bundle)
 
-        if self.self_hosted:
-            if self.upload_bundle_fn is None:
-                raise ValueError("Upload_bundle_fn should be registered")
-            if self.bundle_location_fn is None and bundle_url is None:
-                raise ValueError(
-                    "Need either bundle_location_fn or bundle_url to know where to upload bundles"
-                )
-            if bundle_url is None:
-                bundle_url = self.bundle_location_fn()  # type: ignore
-            self.upload_bundle_fn(serialized_bundle, bundle_url)
-            raw_bundle_url = bundle_url
-        else:
-            # Grab a signed url to make upload to
-            model_bundle_s3_url = self.connection.post(
-                {}, MODEL_BUNDLE_SIGNED_URL_PATH
+        schema_location = None
+        if bool(request_schema) ^ bool(response_schema):
+            raise ValueError(
+                "If request_schema is specified, then response_schema must also be specified."
             )
-            s3_path = model_bundle_s3_url["signedUrl"]
-            raw_bundle_url = f"s3://{model_bundle_s3_url['bucket']}/{model_bundle_s3_url['key']}"
-
-            # Make bundle upload
-
-            requests.put(s3_path, data=serialized_bundle)
+        if request_schema is not None and response_schema is not None:
+            model_definitions = get_model_definitions(
+                request_schema=request_schema,
+                response_schema=response_schema,
+            )
+            model_definitions_encoded = json.dumps(model_definitions).encode()
+            schema_location = self._upload_data(model_definitions_encoded)
 
         payload = dict(
             packaging_type="cloudpickle",
@@ -564,6 +595,7 @@ def create_model_bundle(  # pylint: disable=too-many-statements
             bundle_metadata=bundle_metadata,
             requirements=requirements,
             env_params=env_params,
+            schema_location=schema_location,
         )
 
         _add_app_config_to_bundle_create_payload(payload, app_config)
@@ -581,7 +613,8 @@ def create_model_bundle(  # pylint: disable=too-many-statements
                 requirements=requirements,
                 packaging_type=ModelBundlePackagingType("cloudpickle"),
                 metadata=bundle_metadata,
-                app_config=payload.get("app_config"),
+                app_config=app_config,
+                schema_location=schema_location,
             )
             create_model_bundle_request = CreateModelBundleRequest(**payload)  # type: ignore
             api_instance.create_model_bundle_v1_model_bundles_post(

launch/pydantic_schemas.py

@@ -0,0 +1,54 @@
+from enum import Enum
+from typing import Any, Dict, Set, Type, Union
+
+from pydantic import BaseModel
+from pydantic.schema import get_flat_models_from_models, model_process_schema
+
+REF_PREFIX = "#/components/schemas/"
+
+
+def get_model_definitions(
+    request_schema: Type[BaseModel], response_schema: Type[BaseModel]
+) -> Dict[str, Any]:
+    """
+    Gets the model schemas in jsonschema format from a sequence of Pydantic BaseModels.
+    """
+    flat_models = get_flat_models_from_models(
+        [request_schema, response_schema]
+    )
+    model_name_map = {model: model.__name__ for model in flat_models}
+    model_name_map.update(
+        {request_schema: "RequestSchema", response_schema: "ResponseSchema"}
+    )
+    return get_model_definitions_from_flat_models(
+        flat_models=flat_models, model_name_map=model_name_map
+    )
+
+
+def get_model_definitions_from_flat_models(
+    *,
+    flat_models: Set[Union[Type[BaseModel], Type[Enum]]],
+    model_name_map: Dict[Union[Type[BaseModel], Type[Enum]], str],
+) -> Dict[str, Any]:
+    """
+    Gets the model schemas in jsonschema format from a set of Pydantic BaseModels (or Enums).
+    Inspired by https://github.com/tiangolo/fastapi/blob/99d8470a8e1cf76da8c5274e4e372630efc95736/fastapi/utils.py#L38
+
+    Args:
+        flat_models (Set[Union[Type[BaseModel], Type[Enum]]]): The models.
+        model_name_map (Dict[Union[Type[BaseModel], Type[Enum]], str]): The map from model to name.
+
+    Returns:
+        Dict[str, Any]: OpenAPI-compatible schema of model definitions.
+    """
+    definitions: Dict[str, Dict[str, Any]] = {}
+    for model in flat_models:
+        m_schema, m_definitions, _ = model_process_schema(
+            model, model_name_map=model_name_map, ref_prefix=REF_PREFIX
+        )
+        definitions.update(m_definitions)
+        model_name = model_name_map[model]
+        if "description" in m_schema:
+            m_schema["description"] = m_schema["description"].split("\f")[0]
+        definitions[model_name] = m_schema
+    return definitions

poetry.lock

@@ -47,6 +47,7 @@ pyyaml = ">=5.3.1,<7.0.0"
 typing-extensions = "^4.1.1"
 click = "^8.0.0"
 frozendict = "^2.3.4"
+pydantic = "^1.10.4"
 types-frozendict = "^2.0.9"
 
 [tool.poetry.dev-dependencies]