feat: reorganize the new SDK package (#4337)

* feat: reorganize the new SDK package

Signed-off-by: Frost Ming <[email protected]>

* feat: export types

Signed-off-by: Frost Ming <[email protected]>

* fix: sort imports

Signed-off-by: Frost Ming <[email protected]>

* fix: use real types

Signed-off-by: Frost Ming <[email protected]>

* fixup

Signed-off-by: Frost Ming <[email protected]>

* feat: expose clients

Signed-off-by: Frost Ming <[email protected]>

examples/quickstart/model_pipeline.py

@@ -1,8 +1,9 @@
-import bentoml_io as bentoml
 import joblib
 from sklearn import datasets
 from sklearn import svm
 
+import bentoml
+
 if __name__ == "__main__":
     # Load training data
     iris = datasets.load_iris()
@@ -13,8 +14,6 @@
     clf.fit(X, y)
 
     # Save model to BentoML local model store
-    with bentoml.models.create(
-        "iris_clf",
-    ) as bento_model:
+    with bentoml.models.create("iris_clf") as bento_model:
         joblib.dump(clf, bento_model.path_of("model.pkl"))
     print(f"Model saved: {bento_model}")

examples/quickstart/service.py

@@ -1,6 +1,7 @@
-import bentoml_io as bentoml
 import numpy as np
 
+import bentoml
+
 
 @bentoml.service(resources={"cpu": "200m", "memory": "512Mi"})
 class Preprocessing:

examples/sentence-embedding/service.py

@@ -2,18 +2,14 @@
 
 import os
 
-import bentoml_io
 import numpy as np
 import torch
-from pydantic import Field
 
 import bentoml
+from bentoml import Field
 
 
-@bentoml_io.service(
-    resources={"memory": "500MiB"},
-    traffic={"timeout": 1},
-)
+@bentoml.service(resources={"memory": "500MiB"}, traffic={"timeout": 1})
 class SentenceEmbedding:
     model_ref = bentoml.models.get("all-MiniLM-L6-v2")
 
@@ -43,7 +39,7 @@ def mean_pooling(model_output, attention_mask):
             input_mask_expanded.sum(1), min=1e-9
         )
 
-    @bentoml_io.api(batchable=True)
+    @bentoml.api(batchable=True)
     def encode(
         self,
         sentences: list[str] = Field(["hello world"], description="input sentences"),

pyproject.toml

@@ -133,9 +133,9 @@ fallback_version = "0.0.0"
 [tool.hatch.metadata]
 allow-direct-references = true
 [tool.hatch.build.targets.sdist]
-only-include = ["src/bentoml", "src/bentoml_cli", "src/bentoml_io"]
+only-include = ["src/bentoml", "src/bentoml_cli", "src/_bentoml_sdk", "src/_bentoml_impl"]
 [tool.hatch.build.targets.wheel]
-packages = ["src/bentoml", "src/bentoml_cli", "src/bentoml_io"]
+packages = ["src/bentoml", "src/bentoml_cli", "src/_bentoml_sdk", "src/_bentoml_impl"]
 
 [[tool.pdm.source]]
 url = "https://download.pytorch.org/whl/cpu"
@@ -378,7 +378,7 @@ convention = "google"
 
 [tool.ruff.isort]
 force-single-line = true
-known-first-party = ["bentoml"]
+known-first-party = ["bentoml", "bentoml_cli", "_bentoml_sdk", "_bentoml_impl"]
 
 [tool.pyright]
 pythonVersion = "3.12"

src/bentoml_io/arrow.py → src/_bentoml_impl/arrow.py

@@ -5,7 +5,7 @@
 import pyarrow as pa
 from pydantic import BaseModel
 
-from .types import arrow_serialization
+from _bentoml_sdk.schema import arrow_serialization
 
 SchemaDict: t.TypeAlias = t.Dict[str, t.Any]
 

src/bentoml_io/client/http.py → src/_bentoml_impl/client/http.py

@@ -5,6 +5,8 @@
 import inspect
 import io
 import logging
+import pathlib
+import tempfile
 import typing as t
 from http import HTTPStatus
 from urllib.parse import urljoin
@@ -14,12 +16,12 @@
 import attr
 from pydantic import RootModel
 
+from _bentoml_sdk import IODescriptor
+from _bentoml_sdk.typing_utils import is_file_like
+from _bentoml_sdk.typing_utils import is_image_type
 from bentoml._internal.utils.uri import uri_to_path
 from bentoml.exceptions import BentoMLException
 
-from ..types import File
-from ..typing_utils import is_file_like
-from ..typing_utils import is_image_type
 from .base import AbstractClient
 
 if t.TYPE_CHECKING:
@@ -28,15 +30,18 @@
     from aiohttp import ClientSession
     from aiohttp import MultipartWriter
 
-    from ..factory import Service
-    from ..io_models import IODescriptor
+    from _bentoml_sdk import Service
 
     T = t.TypeVar("T", bound="HTTPClient")
 
 logger = logging.getLogger("bentoml.io")
 MAX_RETRIES = 3
 
 
+def _is_file(obj: t.Any) -> bool:
+    return isinstance(obj, pathlib.PurePath) or is_file_like(obj)
+
+
 @attr.define(slots=True)
 class ClientEndpoint:
     name: str
@@ -59,6 +64,14 @@ class HTTPClient(AbstractClient):
     token: str | None = None
     _client: ClientSession | None = attr.field(init=False, default=None)
     _loop: asyncio.AbstractEventLoop | None = attr.field(init=False, default=None)
+    _opened_files: list[io.BufferedReader] = attr.field(init=False, factory=list)
+    _temp_dir: tempfile.TemporaryDirectory[str] = attr.field(init=False)
+
+    @_temp_dir.default
+    def default_temp_dir(self) -> tempfile.TemporaryDirectory[str]:
+        return tempfile.TemporaryDirectory(
+            prefix="bentoml-client-", ignore_cleanup_errors=True
+        )
 
     def __init__(
         self,
@@ -201,17 +214,22 @@ async def _call(
             endpoint = self.endpoints[name]
         except KeyError:
             raise BentoMLException(f"Endpoint {name} not found") from None
-        data = await self._prepare_request(endpoint, args, kwargs)
-        resp = await self._request(endpoint.route, data, headers=headers)
-        if endpoint.stream_output:
-            return self._parse_stream_response(endpoint, resp)
-        elif (
-            endpoint.output.get("type") == "file"
-            and self.media_type == "application/json"
-        ):
-            return await self._parse_file_response(endpoint, resp)
-        else:
-            return await self._parse_response(endpoint, resp)
+        try:
+            data = await self._prepare_request(endpoint, args, kwargs)
+            resp = await self._request(endpoint.route, data, headers=headers)
+            if endpoint.stream_output:
+                return self._parse_stream_response(endpoint, resp)
+            elif (
+                endpoint.output.get("type") == "file"
+                and self.media_type == "application/json"
+            ):
+                return await self._parse_file_response(endpoint, resp)
+            else:
+                return await self._parse_response(endpoint, resp)
+        finally:
+            for f in self._opened_files:
+                f.close()
+            self._opened_files.clear()
 
     async def _request(
         self,
@@ -262,10 +280,10 @@ def is_file_field(k: str) -> bool:
             if isinstance(model, IODescriptor):
                 return k in model.multipart_fields
             return (
-                is_file_like(value := model[k])
+                _is_file(value := model[k])
                 or isinstance(value, t.Sequence)
                 and len(value) > 0
-                and is_file_like(value[0])
+                and _is_file(value[0])
             )
 
         if isinstance(model, dict):
@@ -287,16 +305,17 @@ def is_file_field(k: str) -> bool:
                                 getattr(v, "_fp", v.fp),
                                 headers={"Content-Type": f"image/{v.format.lower()}"},
                             )
+                        elif isinstance(v, pathlib.PurePath):
+                            file = open(v, "rb")
+                            part = mp.append(file)
+                            self._opened_files.append(file)
                         else:
                             part = mp.append(v)
                         part.set_content_disposition(
                             "attachment", filename=part.filename, name=name
                         )
                 return mp
         elif isinstance(model, dict):
-            for k, v in data.items():
-                if is_file_field(v) and not isinstance(v, File):
-                    data[k] = File(v)
             return self.serde.serialize(data)
         else:
             return self.serde.serialize_model(model)
@@ -309,10 +328,10 @@ async def _prepare_request(
     ) -> bytes | MultipartWriter:
         if endpoint.input_spec is not None:
             model = endpoint.input_spec.from_inputs(*args, **kwargs)
-            if not model.multipart_fields:
-                return self.serde.serialize_model(model)
-            else:
+            if model.multipart_fields:
                 return self._build_multipart(model)
+            else:
+                return self.serde.serialize_model(model)
 
         for name, value in zip(endpoint.input["properties"], args):
             if name in kwargs:
@@ -334,7 +353,7 @@ async def _prepare_request(
         multipart_fields = {
             k
             for k, v in kwargs.items()
-            if is_file_like(v) or isinstance(v, t.Sequence) and v and is_file_like(v[0])
+            if _is_file(v) or isinstance(v, t.Sequence) and v and _is_file(v[0])
         }
         if multipart_fields:
             return self._build_multipart(kwargs)
@@ -371,21 +390,21 @@ async def _parse_stream_response(
 
     async def _parse_file_response(
         self, endpoint: ClientEndpoint, resp: ClientResponse
-    ) -> File:
-        from ..types import Audio
-        from ..types import Image
-        from ..types import Video
-
-        content_type = resp.headers.get("Content-Type")
-        cls = File
-        if content_type:
-            if content_type.startswith("image/"):
-                cls = Image
-            elif content_type.startswith("audio/"):
-                cls = Audio
-            elif content_type.startswith("video/"):
-                cls = Video
-        return cls(io.BytesIO(await resp.read()), media_type=content_type)
+    ) -> pathlib.Path:
+        from multipart.multipart import parse_options_header
+
+        content_disposition = resp.headers.get("content-disposition")
+        filename: str | None = None
+        if content_disposition:
+            _, options = parse_options_header(content_disposition)
+            if b"filename" in options:
+                filename = str(options[b"filename"], "utf-8", errors="ignore")
+
+        with tempfile.NamedTemporaryFile(
+            "wb", suffix=filename, dir=self._temp_dir.name, delete=False
+        ) as f:
+            f.write(await resp.read())
+        return pathlib.Path(f.name)
 
     async def close(self) -> None:
         if self._client is not None and not self._client.closed:

src/bentoml_io/client/proxy.py → src/_bentoml_impl/client/proxy.py

@@ -6,10 +6,10 @@
 import logging
 import typing as t
 
+from _bentoml_sdk import Service
+from _bentoml_sdk.api import APIMethod
 from bentoml._internal.utils import async_gen_to_sync
 from bentoml.exceptions import BentoMLException
-from bentoml_io.api import APIMethod
-from bentoml_io.factory import Service
 
 from .http import ClientEndpoint
 from .http import HTTPClient

src/bentoml_io/serde.py → src/_bentoml_impl/serde.py

@@ -8,12 +8,12 @@
 
 from starlette.datastructures import UploadFile
 
-from .typing_utils import is_list_type
+from _bentoml_sdk.typing_utils import is_list_type
 
 if t.TYPE_CHECKING:
     from starlette.requests import Request
 
-    from .io_models import IODescriptor
+    from _bentoml_sdk import IODescriptor
 
 T = t.TypeVar("T", bound="IODescriptor")
 

src/bentoml_io/server/__init__.py → src/_bentoml_impl/server/__init__.py

@@ -1,8 +1,10 @@
 """
-    bentoml_io.server
-    ~~~~~~~~~~~~~~~~~
+    _bentoml_impl.server
+    ~~~~~~~~~~~~~~~~~~~~
 
     A reference implementation of serving a BentoML service.
     This will be eventually migrated to Rust.
 """
-from ..factory import Service as Service
+from .serving import serve_http
+
+__all__ = ["serve_http"]

src/bentoml_io/server/allocator.py → src/_bentoml_impl/server/allocator.py

@@ -5,12 +5,11 @@
 from simple_di import Provide
 from simple_di import inject
 
+from _bentoml_sdk import Service
 from bentoml._internal.configuration.containers import BentoMLContainer
 from bentoml._internal.resource import system_resources
 from bentoml.exceptions import BentoMLConfigException
 
-from ..factory import Service
-
 
 class ResourceUnavailable(BentoMLConfigException):
     pass