From 6ea81ff7cdb9e3b57985b5e81b4c2d77aecd79a1 Mon Sep 17 00:00:00 2001 From: Aaron Pham <29749331+aarnphm@users.noreply.github.com> Date: Fri, 25 Nov 2022 18:55:27 -0800 Subject: [PATCH] fix: make sure descriptor_id is parsed correctly from servicer Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com> --- src/bentoml/_internal/io_descriptors/base.py | 8 +++-- src/bentoml/_internal/io_descriptors/file.py | 3 +- src/bentoml/_internal/io_descriptors/image.py | 3 +- src/bentoml/_internal/io_descriptors/json.py | 3 +- .../_internal/io_descriptors/multipart.py | 3 +- src/bentoml/_internal/io_descriptors/numpy.py | 3 +- .../_internal/io_descriptors/pandas.py | 5 +-- src/bentoml/_internal/io_descriptors/text.py | 3 +- src/bentoml/_internal/server/grpc/servicer.py | 35 +++++++++++-------- 9 files changed, 41 insertions(+), 25 deletions(-) diff --git a/src/bentoml/_internal/io_descriptors/base.py b/src/bentoml/_internal/io_descriptors/base.py index 870144c865a..21fd9c73cd2 100644 --- a/src/bentoml/_internal/io_descriptors/base.py +++ b/src/bentoml/_internal/io_descriptors/base.py @@ -30,13 +30,17 @@ ) OpenAPIResponse = dict[str, str | dict[str, MediaType] | dict[str, t.Any]] + class SpecDict(t.TypedDict): + id: str + args: dict[str, t.Any] + IO_DESCRIPTOR_REGISTRY: dict[str, type[IODescriptor[t.Any]]] = {} IOType = t.TypeVar("IOType") -def from_spec(spec: dict[str, str]) -> IODescriptor[t.Any]: +def from_spec(spec: SpecDict) -> 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) @@ -129,7 +133,7 @@ def to_spec(self) -> dict[str, t.Any]: @classmethod @abstractmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: raise NotImplementedError @abstractmethod diff --git a/src/bentoml/_internal/io_descriptors/file.py b/src/bentoml/_internal/io_descriptors/file.py index ed4ccefe01c..90f48487254 100644 --- a/src/bentoml/_internal/io_descriptors/file.py +++ b/src/bentoml/_internal/io_descriptors/file.py @@ -30,6 +30,7 @@ from bentoml.grpc.v1 import service_pb2 as pb + from .base import SpecDict from .base import OpenAPIResponse from ..context import InferenceApiContext as Context @@ -143,7 +144,7 @@ def _from_sample(self, sample: FileType | str) -> FileType: return sample @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in File spec: {spec}") return cls(**spec["args"]) diff --git a/src/bentoml/_internal/io_descriptors/image.py b/src/bentoml/_internal/io_descriptors/image.py index ffdd55369e0..d9f78c9d7e7 100644 --- a/src/bentoml/_internal/io_descriptors/image.py +++ b/src/bentoml/_internal/io_descriptors/image.py @@ -36,6 +36,7 @@ from bentoml.grpc.v1 import service_pb2 as pb from .. import external_typing as ext + from .base import SpecDict from .base import OpenAPIResponse from ..context import InferenceApiContext as Context @@ -245,7 +246,7 @@ def to_spec(self) -> dict[str, t.Any]: } @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in Image spec: {spec}") diff --git a/src/bentoml/_internal/io_descriptors/json.py b/src/bentoml/_internal/io_descriptors/json.py index 08c1dc18b59..f32dbda8020 100644 --- a/src/bentoml/_internal/io_descriptors/json.py +++ b/src/bentoml/_internal/io_descriptors/json.py @@ -35,6 +35,7 @@ from typing_extensions import Self from .. import external_typing as ext + from .base import SpecDict from .base import OpenAPIResponse from ..context import InferenceApiContext as Context @@ -235,7 +236,7 @@ def to_spec(self) -> dict[str, t.Any]: } @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> 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"]: diff --git a/src/bentoml/_internal/io_descriptors/multipart.py b/src/bentoml/_internal/io_descriptors/multipart.py index 2f59866df56..108261b1144 100644 --- a/src/bentoml/_internal/io_descriptors/multipart.py +++ b/src/bentoml/_internal/io_descriptors/multipart.py @@ -25,6 +25,7 @@ from bentoml.grpc.v1 import service_pb2 as pb + from .base import SpecDict from .base import OpenAPIResponse from ..types import LazyType from ..context import InferenceApiContext as Context @@ -202,7 +203,7 @@ def to_spec(self) -> dict[str, t.Any]: } @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in Multipart spec: {spec}") return Multipart( diff --git a/src/bentoml/_internal/io_descriptors/numpy.py b/src/bentoml/_internal/io_descriptors/numpy.py index a928213eafe..8508bb142b0 100644 --- a/src/bentoml/_internal/io_descriptors/numpy.py +++ b/src/bentoml/_internal/io_descriptors/numpy.py @@ -28,6 +28,7 @@ from bentoml.grpc.v1 import service_pb2 as pb from .. import external_typing as ext + from .base import SpecDict from .base import OpenAPIResponse from ..context import InferenceApiContext as Context else: @@ -252,7 +253,7 @@ def to_spec(self) -> dict[str, t.Any]: } @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in NumpyNdarray spec: {spec}") res = NumpyNdarray(**spec["args"]) diff --git a/src/bentoml/_internal/io_descriptors/pandas.py b/src/bentoml/_internal/io_descriptors/pandas.py index 28c7e0f08e9..921a3f71bf7 100644 --- a/src/bentoml/_internal/io_descriptors/pandas.py +++ b/src/bentoml/_internal/io_descriptors/pandas.py @@ -35,6 +35,7 @@ from bentoml.grpc.v1 import service_pb2 as pb from .. import external_typing as ext + from .base import SpecDict from .base import OpenAPIResponse from ..context import InferenceApiContext as Context @@ -453,7 +454,7 @@ def to_spec(self) -> dict[str, t.Any]: } @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in PandasDataFrame spec: {spec}") res = PandasDataFrame(**spec["args"]) @@ -900,7 +901,7 @@ def to_spec(self) -> dict[str, t.Any]: } @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: if "args" not in spec: raise InvalidArgument(f"Missing args key in PandasSeries spec: {spec}") res = PandasSeries(**spec["args"]) diff --git a/src/bentoml/_internal/io_descriptors/text.py b/src/bentoml/_internal/io_descriptors/text.py index ed9127069a2..944840b11ff 100644 --- a/src/bentoml/_internal/io_descriptors/text.py +++ b/src/bentoml/_internal/io_descriptors/text.py @@ -19,6 +19,7 @@ from google.protobuf import wrappers_pb2 from typing_extensions import Self + from .base import SpecDict from .base import OpenAPIResponse from ..context import InferenceApiContext as Context else: @@ -111,7 +112,7 @@ def to_spec(self) -> dict[str, t.Any]: return {"id": self.descriptor_id} @classmethod - def from_spec(cls, spec: dict[str, t.Any]) -> Self: + def from_spec(cls, spec: SpecDict) -> Self: return cls() def openapi_schema(self) -> Schema: diff --git a/src/bentoml/_internal/server/grpc/servicer.py b/src/bentoml/_internal/server/grpc/servicer.py index f3dd28a2ebb..dea20e1c407 100644 --- a/src/bentoml/_internal/server/grpc/servicer.py +++ b/src/bentoml/_internal/server/grpc/servicer.py @@ -35,14 +35,17 @@ from bentoml.grpc.types import BentoServicerContext from ...service.service import Service - from ...io_descriptors.base import IODescriptor + from ...io_descriptors.base import SpecDict if LATEST_PROTOCOL_VERSION == "v1": from bentoml.grpc.v1 import service_pb2 as pb from bentoml.grpc.v1 import service_pb2_grpc as services + from bentoml.grpc.v1.service_pb2 import ServiceMetadataRequest + from bentoml.grpc.v1.service_pb2 import ServiceMetadataResponse else: from bentoml.grpc.v1alpha1 import service_pb2 as pb from bentoml.grpc.v1alpha1 import service_pb2_grpc as services + else: grpc, aio = import_grpc() health = LazyLoader( @@ -194,7 +197,6 @@ async def Call( if protocol_version == "v1": # "v1" introduces ServiceMetadata to send in bentoml.Service information. - from bentoml.grpc.v1.service_pb2 import ServiceMetadataRequest from bentoml.grpc.v1.service_pb2 import ServiceMetadataResponse async def ServiceMetadata( @@ -209,13 +211,11 @@ async def ServiceMetadata( ServiceMetadataResponse.InferenceAPI( name=api.name, docs=api.doc, - input=ServiceMetadataResponse.DescriptorMetadata( - descriptor_id=api.input.descriptor_id, - attributes=make_attributes_struct(api.input), + input=make_descriptor_spec( + api.input.to_spec(), ServiceMetadataResponse ), - output=ServiceMetadataResponse.DescriptorMetadata( - descriptor_id=api.output.descriptor_id, - attributes=make_attributes_struct(api.output), + output=make_descriptor_spec( + api.output.to_spec(), ServiceMetadataResponse ), ) for api in service.apis.values() @@ -253,13 +253,18 @@ def _tuple_converter(d: NestedDictStrAny | None) -> NestedDictStrAny | None: return d -def make_attributes_struct(io: IODescriptor[t.Any]) -> struct_pb2.Struct: +def make_descriptor_spec( + spec: SpecDict, pb: type[ServiceMetadataResponse] +) -> ServiceMetadataResponse.DescriptorMetadata: from ...io_descriptors.json import parse_dict_to_proto - return struct_pb2.Struct( - fields={ - "args": parse_dict_to_proto( - _tuple_converter(io.to_spec().get("args", None)), struct_pb2.Value() - ) - } + return pb.DescriptorMetadata( + descriptor_id=spec["id"], + attributes=struct_pb2.Struct( + fields={ + "args": parse_dict_to_proto( + _tuple_converter(spec["args"]), struct_pb2.Value() + ) + } + ), )