feat: bento client (#3028)

Co-authored-by: Aaron Pham <[email protected]>

src/bentoml/_internal/io_descriptors/__init__.py

@@ -1,4 +1,8 @@
+from __future__ import annotations
+
+from .base import from_spec
 from .base import IODescriptor
+from .base import IO_DESCRIPTOR_REGISTRY
 from .file import File
 from .json import JSON
 from .text import Text
@@ -9,6 +13,7 @@
 from .multipart import Multipart
 
 __all__ = [
+    "IO_DESCRIPTOR_REGISTRY",
     "File",
     "Image",
     "IODescriptor",
@@ -18,4 +23,5 @@
     "PandasDataFrame",
     "PandasSeries",
     "Text",
+    "from_spec",
 ]

src/bentoml/_internal/io_descriptors/base.py

@@ -5,9 +5,12 @@
 from abc import abstractmethod
 from typing import TYPE_CHECKING
 
+from ...exceptions import InvalidArgument
+
 if TYPE_CHECKING:
     from types import UnionType
 
+    from typing_extensions import Self
     from starlette.requests import Request
     from starlette.responses import Response
 
@@ -28,9 +31,17 @@
     )
 
 
+IO_DESCRIPTOR_REGISTRY: dict[str, type[IODescriptor[t.Any]]] = {}
+
 IOType = t.TypeVar("IOType")
 
 
+def from_spec(spec: dict[str, str]) -> IODescriptor[t.Any]:
+    if "id" not in spec:
+        raise InvalidArgument(f"IO descriptor spec ({spec}) missing ID.")
+    return IO_DESCRIPTOR_REGISTRY[spec["id"]].from_spec(spec)
+
+
 class IODescriptor(ABC, t.Generic[IOType]):
     """
     IODescriptor describes the input/output data format of an InferenceAPI defined
@@ -43,13 +54,32 @@ class IODescriptor(ABC, t.Generic[IOType]):
     _mime_type: str
     _rpc_content_type: str = "application/grpc"
     _proto_fields: tuple[ProtoField]
+    descriptor_id: str | None
+
+    def __init_subclass__(cls, *, descriptor_id: str | None = None):
+        if descriptor_id is not None:
+            if descriptor_id in IO_DESCRIPTOR_REGISTRY:
+                raise ValueError(
+                    f"Descriptor ID {descriptor_id} already registered to {IO_DESCRIPTOR_REGISTRY[descriptor_id]}."
+                )
+            IO_DESCRIPTOR_REGISTRY[descriptor_id] = cls
+        cls.descriptor_id = descriptor_id
+
+    @abstractmethod
+    def to_spec(self) -> dict[str, t.Any]:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_spec(cls, spec: dict[str, t.Any]) -> Self:
+        raise NotImplementedError
 
     def __repr__(self) -> str:
         return self.__class__.__qualname__
 
     @abstractmethod
     def input_type(self) -> InputType:
-        ...
+        raise NotImplementedError
 
     @abstractmethod
     def openapi_schema(self) -> Schema | Reference:

src/bentoml/_internal/io_descriptors/file.py

@@ -37,7 +37,7 @@
 FileType = t.Union[io.IOBase, t.IO[bytes], FileLike[bytes]]
 
 
-class File(IODescriptor[FileType]):
+class File(IODescriptor[FileType], descriptor_id="bentoml.io.File"):
     """
     :obj:`File` defines API specification for the inputs/outputs of a Service, where either
     inputs will be converted to or outputs will be converted from file-like objects as
@@ -121,6 +121,15 @@ def __new__(  # pylint: disable=arguments-differ # returning subclass from new
         res._mime_type = mime_type
         return res
 
+    def to_spec(self):
+        raise NotImplementedError
+
+    @classmethod
+    def from_spec(cls, spec: dict[str, t.Any]) -> Self:
+        if "args" not in spec:
+            raise InvalidArgument(f"Missing args key in File spec: {spec}")
+        return cls(**spec["args"])
+
     def input_type(self) -> t.Type[t.Any]:
         return FileLike[bytes]
 
@@ -130,17 +139,19 @@ def openapi_schema(self) -> Schema:
     def openapi_components(self) -> dict[str, t.Any] | None:
         pass
 
-    def openapi_request_body(self) -> RequestBody:
-        return RequestBody(
-            content={self._mime_type: MediaType(schema=self.openapi_schema())},
-            required=True,
-        )
+    def openapi_request_body(self) -> dict[str, t.Any]:
+        return {
+            "content": {self._mime_type: MediaType(schema=self.openapi_schema())},
+            "required": True,
+            "x-bentoml-io-descriptor": self.to_spec(),
+        }
 
     def openapi_responses(self) -> OpenAPIResponse:
-        return OpenAPIResponse(
-            description=SUCCESS_DESCRIPTION,
-            content={self._mime_type: MediaType(schema=self.openapi_schema())},
-        )
+        return {
+            "description": SUCCESS_DESCRIPTION,
+            "content": {self._mime_type: MediaType(schema=self.openapi_schema())},
+            "x-bentoml-io-descriptor": self.to_spec(),
+        }
 
     async def to_http_response(self, obj: FileType, ctx: Context | None = None):
         if isinstance(obj, bytes):
@@ -176,16 +187,23 @@ async def to_proto(self, obj: FileType) -> pb.File:
 
         return pb.File(kind=kind, content=body)
 
-    if TYPE_CHECKING:
+    async def from_proto(self, field: pb.File | bytes) -> FileLike[bytes]:
+        raise NotImplementedError
 
-        async def from_proto(self, field: pb.File | bytes) -> FileLike[bytes]:
-            ...
+    async def from_http_request(self, request: Request) -> t.IO[bytes]:
+        raise NotImplementedError
 
-        async def from_http_request(self, request: Request) -> t.IO[bytes]:
-            ...
 
+class BytesIOFile(File, descriptor_id=None):
+    def to_spec(self) -> dict[str, t.Any]:
+        return {
+            "id": super().descriptor_id,
+            "args": {
+                "kind": "binaryio",
+                "mime_type": self._mime_type,
+            },
+        }
 
-class BytesIOFile(File):
     async def from_http_request(self, request: Request) -> t.IO[bytes]:
         content_type, _ = parse_options_header(request.headers["content-type"])
         if content_type.decode("utf-8") == "multipart/form-data":

src/bentoml/_internal/io_descriptors/image.py

@@ -83,7 +83,7 @@ def initialize_pillow():
     READABLE_MIMES = {k for k, v in MIME_EXT_MAPPING.items() if v not in PIL_WRITE_ONLY_FORMATS}  # type: ignore (lazy constant)
 
 
-class Image(IODescriptor[ImageType]):
+class Image(IODescriptor[ImageType], descriptor_id="bentoml.io.Image"):
     """
     :obj:`Image` defines API specification for the inputs/outputs of a Service, where either
     inputs will be converted to or outputs will be converted from images as specified
@@ -213,6 +213,23 @@ def __init__(
         self._pilmode: _Mode | None = pilmode
         self._format: str = MIME_EXT_MAPPING[self._mime_type]
 
+    def to_spec(self) -> dict[str, t.Any]:
+        return {
+            "id": self.descriptor_id,
+            "args": {
+                "pilmode": self._pilmode,
+                "mime_type": self._mime_type,
+                "allowed_mime_types": list(self._allowed_mimes),
+            },
+        }
+
+    @classmethod
+    def from_spec(cls) -> Self:
+        if "args" not in spec:
+            raise InvalidArgument(f"Missing args key in Image spec: {spec}")
+
+        return cls(**spec["args"])
+
     def input_type(self) -> UnionType:
         return ImageType
 

src/bentoml/_internal/io_descriptors/json.py

@@ -72,7 +72,7 @@ def default(self, o: type) -> t.Any:
         return super().default(o)
 
 
-class JSON(IODescriptor[JSONType]):
+class JSON(IODescriptor[JSONType], descriptor_id="bentoml.io.JSON"):
     """
     :obj:`JSON` defines API specification for the inputs/outputs of a Service, where either
     inputs will be converted to or outputs will be converted from a JSON representation
@@ -200,6 +200,30 @@ def __init__(
                 "'validate_json' option from 'bentoml.io.JSON' has been deprecated. Use a Pydantic model to specify validation options instead."
             )
 
+    def to_spec(self) -> dict[str, t.Any]:
+        return {
+            "id": self.descriptor_id,
+            "args": {
+                "has_pydantic_model": self._pydantic_model is not None,
+                "has_json_encoder": self._json_encoder is not None,
+            },
+        }
+
+    @classmethod
+    def from_spec(cls, spec: dict[str, t.Any]) -> Self:
+        if "args" not in spec:
+            raise InvalidArgument(f"Missing args key in JSON spec: {spec}")
+        if "has_pydantic_model" in spec["args"] and spec["args"]["has_pydantic_model"]:
+            logger.warning(
+                "BentoML does not support loading pydantic models from URLs; output will be a normal dictionary."
+            )
+        if "has_json_encoder" in spec["args"] and spec["args"]["has_json_encoder"]:
+            logger.warning(
+                "BentoML does not support loading JSON encoders from URLs; output will be a normal dictionary."
+            )
+
+        return cls()
+
     def input_type(self) -> UnionType:
         return JSONType
 
@@ -227,16 +251,18 @@ def openapi_components(self) -> dict[str, t.Any] | None:
         return {"schemas": pydantic_components_schema(self._pydantic_model)}
 
     def openapi_request_body(self) -> RequestBody:
-        return RequestBody(
-            content={self._mime_type: MediaType(schema=self.openapi_schema())},
-            required=True,
-        )
+        return {
+            "content": {self._mime_type: MediaType(schema=self.openapi_schema())},
+            "required": True,
+            "x-bentoml-io-descriptor": self.to_spec(),
+        }
 
     def openapi_responses(self) -> OpenAPIResponse:
-        return OpenAPIResponse(
-            description=SUCCESS_DESCRIPTION,
-            content={self._mime_type: MediaType(schema=self.openapi_schema())},
-        )
+        return {
+            "description": SUCCESS_DESCRIPTION,
+            "content": {self._mime_type: MediaType(schema=self.openapi_schema())},
+            "x-bentoml-io-descriptor": self.to_spec(),
+        }
 
     async def from_http_request(self, request: Request) -> JSONType:
         json_str = await request.body()

src/bentoml/_internal/io_descriptors/multipart.py

@@ -8,6 +8,7 @@
 from multipart.multipart import parse_options_header
 from starlette.responses import Response
 
+from . import from_spec as io_descriptor_from_spec
 from .base import IODescriptor
 from ...exceptions import InvalidArgument
 from ...exceptions import BentoMLException
@@ -32,7 +33,7 @@
     pb, _ = import_generated_stubs()
 
 
-class Multipart(IODescriptor[t.Dict[str, t.Any]]):
+class Multipart(IODescriptor[t.Dict[str, t.Any]], descriptor_id="bentoml.io.Multipart"):
     """
     :obj:`Multipart` defines API specification for the inputs/outputs of a Service, where inputs/outputs
     of a Service can receive/send a **multipart** request/responses as specified in your API function signature.
@@ -187,6 +188,26 @@ def input_type(
 
         return res
 
+    def to_spec(self) -> dict[str, t.Any]:
+        return {
+            "id": self.descriptor_id,
+            "args": {
+                argname: descriptor.to_spec()
+                for argname, descriptor in self._inputs.items()
+            },
+        }
+
+    @classmethod
+    def from_spec(cls, spec: dict[str, t.Any]) -> Self:
+        if "args" not in spec:
+            raise InvalidArgument(f"Missing args key in Multipart spec: {spec}")
+        return Multipart(
+            **{
+                argname: io_descriptor_from_spec(spec)
+                for argname, spec in spec["args"].items()
+            }
+        )
+
     def openapi_schema(self) -> Schema:
         return Schema(
             type="object",
@@ -197,16 +218,18 @@ def openapi_components(self) -> dict[str, t.Any] | None:
         pass
 
     def openapi_request_body(self) -> RequestBody:
-        return RequestBody(
-            content={self._mime_type: MediaType(schema=self.openapi_schema())},
-            required=True,
-        )
+        return {
+            "content": {self._mime_type: MediaType(schema=self.openapi_schema())},
+            "required": True,
+            "x-bentoml-descriptor": self.to_spec(),
+        }
 
     def openapi_responses(self) -> OpenAPIResponse:
-        return OpenAPIResponse(
-            description=SUCCESS_DESCRIPTION,
-            content={self._mime_type: MediaType(schema=self.openapi_schema())},
-        )
+        return {
+            "description": SUCCESS_DESCRIPTION,
+            "content": {self._mime_type: MediaType(schema=self.openapi_schema())},
+            "x-bentoml-descriptor": self.to_spec(),
+        }
 
     async def from_http_request(self, request: Request) -> dict[str, t.Any]:
         ctype, _ = parse_options_header(request.headers["content-type"])