Repository: bentoml/BentoML

Issue:
feature: Support for Dict of Tensors in production mode
### Feature request

When running `serve --production`, the invocation to `runner.async_run()` goes through an instance of `DataContainer` according to the input type. If you have a PyTorch model and passes as input a `torch.Tensor`, it will be serialized using `PyTorchTensorContainer`, sent to the runner server and then be deserialized there (and possibly having batches combined and so on). It also works if you invoke it using a numpy array or a pandas dataframe. It will be serialized and then converted into a Tensor later. However, it only works if your PyTorch model receives one or more tensors (eg. forward(tensor1: torch.Tensor, tensor2: torch.Tensor)`. If your model receives a Dict of Tensors, it doesn't work. And it's very common for models to receive this kind of input (BERT for example, that receives a dict with 'input_ids', 'token_type_ids', and 'attention_mask').

When you try to invoke the model with a dict, it will work outside of production mode. But in production mode, it will try to use the `DefaultContainer`, but it fails to deserialize the dict correctly and you get an error from the runner server like:

```
Traceback (most recent call last):
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/protocols/http/h11_impl.py", line 429, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/middleware/proxy_headers.py", line 78, in __call__
    return await self.app(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/middleware/message_logger.py", line 86, in __call__
    raise exc from None
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/uvicorn/middleware/message_logger.py", line 82, in __call__
    await self.app(scope, inner_receive, inner_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/applications.py", line 118, in __call__
    await self.middleware_stack(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/errors.py", line 184, in __call__
    raise exc
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/errors.py", line 162, in __call__
    await self.app(scope, receive, _send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/opentelemetry/instrumentation/asgi/__init__.py", line 581, in __call__
    await self.app(scope, otel_receive, otel_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/http/instruments.py", line 293, in __call__
    await self.app(scope, receive, wrapped_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/http/access.py", line 126, in __call__
    await self.app(scope, receive, wrapped_send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/exceptions.py", line 79, in __call__
    raise exc
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/middleware/exceptions.py", line 68, in __call__
    await self.app(scope, receive, sender)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/routing.py", line 706, in __call__
    await route.handle(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/routing.py", line 276, in handle
    await self.app(scope, receive, send)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/starlette/routing.py", line 66, in app
    response = await func(request)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/runner_app.py", line 276, in _request_handler
    payload = await infer(params)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/marshal/dispatcher.py", line 180, in _func
    raise r
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/marshal/dispatcher.py", line 444, in outbound_call
    outputs = await self.callback(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/server/runner_app.py", line 237, in infer_batch
    batch_ret = await runner_method.async_run(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/runner/runner.py", line 55, in async_run
    return await self.runner._runner_handle.async_run_method(self, *args, **kwargs)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/runner/runner_handle/local.py", line 59, in async_run_method
    return await anyio.to_thread.run_sync(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/anyio/to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 867, in run
    result = context.run(func, *args)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/runner/runnable.py", line 139, in method
    return self.func(obj, *args, **kwargs)
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/bentoml/_internal/frameworks/common/pytorch.py", line 107, in _run
    return getattr(self.model, method_name)(
  File "/home/fernando/.local/share/virtualenvs/page-tagging-svc-6EE2Yu73/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/fernando/bentoml/bentos/page-tagging-svc/ibwdpixeoopskldn/src/src/torch/model/transformer.py", line 106, in forward
    input_ids=batch["input_ids"],
TypeError: list indices must be integers or slices, not str
```

The dict somehow becomes a list of the keys of the dict.

As suggestion, BentoML could have a `DictContainer`, which would use `AutoContainer` for each value inside the dictionary. The `AutoContainer` would find the correct `DataContainer` for each type (`PyTorchTensorContainer` for tensors, for example), that would be able to handle the serialization and deserialization, combining batches and so on.

### Motivation

I need to use BentoML with a model that receives a `Dict[str, torch.Tensor]`.

### Other

I have the following workaround in place for now. It's worth noting that it only handles Dict of Tensors though.

```python
class PyTorchTensorDictContainer(DataContainer[Dict[str, torch.Tensor], torch.Tensor]):
    @classmethod
    def batches_to_batch(
        cls,
        batches: t.Sequence[Dict[str, torch.Tensor]],
        batch_dim: int = 0,
    ) -> t.Tuple[Dict[str, torch.Tensor], list[int]]:
        batch = {}
        for key in batches[0].keys():
            batch[key] = torch.cat(
                tuple(subbatch[key] for subbatch in batches), dim=batch_dim
            )
        indices = list(
            itertools.accumulate(subbatch[key].shape[batch_dim] for subbatch in batches)
        )
        indices = [0] + indices
        return batch, indices

    @classmethod
    def batch_to_batches(
        cls,
        batch: Dict[str, torch.Tensor],
        indices: t.Sequence[int],
        batch_dim: int = 0,
    ) -> t.List[Dict[str, torch.Tensor]]:
        sizes = [indices[i] - indices[i - 1] for i in range(1, len(indices))]
        output: list[Dict[str, torch.Tensor]] = [{} for _ in range(len(sizes))]
        for key in batch.keys():
            split_tensor = torch.split(batch[key], sizes, dim=batch_dim)
            for i, tensor in enumerate(split_tensor):
                output[i][key] = tensor
        return output

    @classmethod
    @inject
    def to_payload(  # pylint: disable=arguments-differ
        cls,
        batch: Dict[str, torch.Tensor],
        batch_dim: int = 0,
    ) -> Payload:
        batch = batch.copy()
        for key in batch.keys():
            batch[key] = batch[key].cpu().numpy()

        return cls.create_payload(
            pickle.dumps(batch),
            batch_size=batch[key].shape[batch_dim],
            meta={"plasma": False},
        )

    @classmethod
    @inject
    def from_payload(  # pylint: disable=arguments-differ
        cls,
        payload: Payload,
    ) -> Dict[str, torch.Tensor]:
        ret: Dict[str, np.ndarray] = pickle.loads(payload.data)
        output = {}
        for key in ret.keys():
            output[key] = torch.from_numpy(ret[key]).requires_grad_(False)
        return output

    @classmethod
    @inject
    def batch_to_payloads(  # pylint: disable=arguments-differ
        cls,
        batch: Dict[str, torch.Tensor],
        indices: t.Sequence[int],
        batch_dim: int = 0,
    ) -> t.List[Payload]:
        batches = cls.batch_to_batches(batch, indices, batch_dim)
        payloads = [cls.to_payload(i, batch_dim=batch_dim) for i in batches]
        return payloads

    @classmethod
    @inject
    def from_batch_payloads(  # pylint: disable=arguments-differ
        cls,
        payloads: t.Sequence[Payload],
        batch_dim: int = 0,
    ) -> t.Tuple[Dict[str, torch.Tensor], list[int]]:
        batches = [cls.from_payload(payload) for payload in payloads]
        return cls.batches_to_batch(batches, batch_dim)

def make_pytorch_runnable_method(
    method_name: str,
    partial_kwargs: dict[str, t.Any] | None = None,
) -> t.Callable[..., torch.Tensor]:
    if partial_kwargs is None:
        partial_kwargs = {}

    def _run(
        self: PytorchModelRunnable,
        *args: t.Dict[str, torch.Tensor],
        **kwargs,
    ) -> torch.Tensor:
        params = Params(*args, **kwargs)

        def _mapping(item: t.Dict[str, torch.Tensor]) -> t.Dict[str, torch.Tensor]:
            item = item.copy()
            for key, tensor in item.items():
                item[key] = tensor.to(self.device_id)

            return item

        with inference_mode_ctx():
            params = params.map(_mapping)
            return getattr(self.model, method_name)(
                *params.args,
                **dict(partial_kwargs, **params.kwargs),
            )

    return _run

def get_runnable(bento_model: Model):
    partial_kwargs: t.Dict[str, t.Any] = bento_model.info.options.partial_kwargs  # type: ignore

    runnable_class: type[PytorchModelRunnable] = partial_class(
        PytorchModelRunnable,
        bento_model=bento_model,
        loader=load_model,
    )
    for method_name, options in bento_model.info.signatures.items():
        method_partial_kwargs = partial_kwargs.get(method_name)
        runnable_class.add_method(
            make_pytorch_runnable_method(method_name, method_partial_kwargs),
            name=method_name,
            batchable=options.batchable,
            batch_dim=options.batch_dim,
            input_spec=options.input_spec,
            output_spec=options.output_spec,
        )
    return runnable_class
```

Then, instead of creating the runner with `model.to_runner()`, I use:

```python
runner = Runner(
    get_runnable(model),
    name=model.tag.name,
    models=[model],
)
```

Discussion:

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 3e342c33..a0b9ccb9 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -73,7 +73,7 @@ jobs:
         run: |
           npm install -g npm@^7 pyright
       - name: Setup bufbuild/buf
-        uses: bufbuild/buf-setup-action@v1.18.0
+        uses: bufbuild/buf-setup-action@v1.19.0
         with:
           github_token: ${{ github.token }}
       - name: Cache pip dependencies
diff --git a/docs/source/guides/batching.rst b/docs/source/guides/batching.rst
index 8623c5e0..8aca5d42 100644
--- a/docs/source/guides/batching.rst
+++ b/docs/source/guides/batching.rst
@@ -88,3 +88,52 @@ Error handling
 If adaptive batching cannot keep up with rate of the incoming requests while satisfying the max
 latency configuration, HTTP 503 Service Unavailable is returned. To workaround this error, consider
 relaxing the max latency requirement and further scaling the underlying hardware resources.
+
+Custom Batching
+---------------
+
+Currently, adaptive batching is only effective for certain types of parameters including(non-exhaustive) ``numpy.ndarray``, ``pandas.Series`` and framework-specific types such as ``torch.Tensor``.
+Batch parameters of other types are simply collected into a list and passed to the inference function. If your model accepts parameters that are not batchable by default,
+you can achieve adaptive batching by wrapping the inference function with a :ref:`Runner <concepts/runner:Custom Runner>`.
+
+We will demonstrate this with a PyTorch example which accepts a dictionary of ``torch.Tensor`` as input.
+
+.. code-block:: python
+    :caption: `service.py`
+
+    import bentoml
+    import torch
+
+    class MyRunnable(bentoml.Runnable):
+        def __init__(self):
+            super().__init__()
+            # if torch.cuda.device_count():
+            if torch.cuda.is_available():
+                self.device_id = "cuda"
+                # by default, torch.FloatTensor will be used on CPU.
+                torch.set_default_tensor_type("torch.cuda.FloatTensor")
+            else:
+                self.device_id = "cpu"
+            self.model = bentoml.pytorch.load_model("my_pytorch_model", device_id=self.device_id)
+            # We want to turn off dropout and batchnorm when running inference.
+            self.model.train(False)
+
+        @bentoml.Runnable.method(batchable=True)
+        def __call__(self, **batch):
+            # Our model accepts a dictionary of Tensor as input, but we use ``**``
+            # to deconstruct it to keyword arguments, so they can be batched rightly.
+            # It works only if every parameter is "batchable" such as torch.Tensor/np.ndarray.
+            # Move the tensors to the target device and pass it to the model as a dict.
+            batch = {k: tensor.to(self.device_id) for k, tensor in batch.items()}
+            with torch.inference_mode():
+                return self.model(batch)
+
+    # Build the runner from the runnable manually, instead of calling model.to_runner() method
+    runner = bentoml.Runner(MyRunnable)
+    svc = bentoml.Service(name="my_ml_service", runners=[runner])
+
+    # Define an API endpoint for the inference
+    @svc.api(input=bentoml.io.JSON(), output=bentoml.io.JSON())
+    async def predict(api_input):
+        # Deconstruct the input dictionary to keyword arguments for batching
+        return await runner.async_run(**api_input)
diff --git a/docs/source/guides/logging.rst b/docs/source/guides/logging.rst
index 32ab430f..ed2cbd94 100644
--- a/docs/source/guides/logging.rst
+++ b/docs/source/guides/logging.rst
@@ -132,6 +132,9 @@ When using BentoML as a library, BentoML does not configure any logs. By default
     bentoml_logger.addHandler(ch)
     bentoml_logger.setLevel(logging.DEBUG)
 
+.. note::
+
+    :bdg-info:`Important:` ``bentoml serve`` will fork the given `service.py` into the child process. This means any type of rotation FileHandler are not supported within the service definition. See `notes from Python logging module <https://docs.python.org/3/howto/logging-cookbook.html#logging-to-a-single-file-from-multiple-processes>`_ for more information.
 ----
 
 .. rubric:: Notes
diff --git a/examples/custom_runner/torch_hub_yolov5/requirements.txt b/examples/custom_runner/torch_hub_yolov5/requirements.txt
index 40cb5088..ab870bce 100644
--- a/examples/custom_runner/torch_hub_yolov5/requirements.txt
+++ b/examples/custom_runner/torch_hub_yolov5/requirements.txt
@@ -5,7 +5,8 @@ matplotlib>=3.2.2
 numpy>=1.18.5
 opencv-python>=4.1.1
 pandas>=1.1.4
-protobuf<4.21.3  # https://github.com/ultralytics/yolov5/issues/8012
+protobuf<=3.20.1 # https://github.com/ultralytics/yolov5/issues/8012
+ipython
 requests>=2.23.0
 scipy>=1.4.1  # Google Colab version
 seaborn>=0.11.0
diff --git a/examples/custom_runner/torch_hub_yolov5/service.py b/examples/custom_runner/torch_hub_yolov5/service.py
index a6b04a83..ec6d8903 100644
--- a/examples/custom_runner/torch_hub_yolov5/service.py
+++ b/examples/custom_runner/torch_hub_yolov5/service.py
@@ -10,7 +10,7 @@ class Yolov5Runnable(bentoml.Runnable):
     SUPPORTS_CPU_MULTI_THREADING = True
 
     def __init__(self):
-        self.model = torch.hub.load("ultralytics/yolov5", "yolov5s")
+        self.model = torch.hub.load("ultralytics/yolov5:v6.2", "yolov5s")
 
         if torch.cuda.is_available():
             self.model.cuda()
diff --git a/src/bentoml/_internal/bento/build_config.py b/src/bentoml/_internal/bento/build_config.py
index b538ba0b..a531bd53 100644
--- a/src/bentoml/_internal/bento/build_config.py
+++ b/src/bentoml/_internal/bento/build_config.py
@@ -8,7 +8,6 @@ import logging
 import subprocess
 from sys import version_info
 from shlex import quote
-from typing import TYPE_CHECKING
 
 import fs
 import attr
@@ -34,7 +33,7 @@ from ..container.frontend.dockerfile import ALLOWED_CUDA_VERSION_ARGS
 from ..container.frontend.dockerfile import SUPPORTED_PYTHON_VERSIONS
 from ..container.frontend.dockerfile import CONTAINER_SUPPORTED_DISTROS
 
-if TYPE_CHECKING:
+if t.TYPE_CHECKING:
     from attr import Attribute
     from fs.base import FS
 
@@ -278,7 +277,7 @@ class DockerOptions:
         return bentoml_cattr.unstructure(self)
 
 
-if TYPE_CHECKING:
+if t.TYPE_CHECKING:
     CondaPipType = dict[t.Literal["pip"], list[str]]
     DependencyType = list[str | CondaPipType]
 else:
@@ -312,7 +311,7 @@ def conda_dependencies_validator(
                 )
 
 
-if TYPE_CHECKING:
+if t.TYPE_CHECKING:
     ListStr: t.TypeAlias = list[str]
     CondaYamlDict = dict[str, DependencyType | list[str]]
 else:
@@ -689,7 +688,7 @@ def _python_options_structure_hook(d: t.Any, _: t.Type[PythonOptions]) -> Python
 bentoml_cattr.register_structure_hook(PythonOptions, _python_options_structure_hook)
 
 
-if TYPE_CHECKING:
+if t.TYPE_CHECKING:
     OptionsCls = DockerOptions | CondaOptions | PythonOptions
 
 
@@ -740,7 +739,7 @@ class BentoBuildConfig:
         converter=dict_options_converter(CondaOptions),
     )
 
-    if TYPE_CHECKING:
+    if t.TYPE_CHECKING:
         # NOTE: This is to ensure that BentoBuildConfig __init__
         # satisfies type checker. docker, python, and conda accepts
         # dict[str, t.Any] since our converter will handle the conversion.
diff --git a/src/bentoml/_internal/frameworks/transformers.py b/src/bentoml/_internal/frameworks/transformers.py
index 303c2da6..6cb1add8 100644
--- a/src/bentoml/_internal/frameworks/transformers.py
+++ b/src/bentoml/_internal/frameworks/transformers.py
@@ -25,6 +25,8 @@ from ..models.model import ModelOptions
 from ..models.model import ModelSignature
 
 if t.TYPE_CHECKING:
+    import torch
+    import tensorflow as tf
     import cloudpickle
     from transformers.models.auto.auto_factory import (
         _BaseAutoModelClass as BaseAutoModelClass,
@@ -49,6 +51,10 @@ if t.TYPE_CHECKING:
     TupleAutoModel = tuple[BaseAutoModelClass, ...]
 
     class PreTrainedProtocol(t.Protocol):
+        @property
+        def framework(self) -> str:
+            ...
+
         def save_pretrained(self, save_directory: str, **kwargs: t.Any) -> None:
             ...
 
@@ -64,6 +70,8 @@ else:
     TupleStr = TupleAutoModel = tuple
     DefaultMapping = SimpleDefaultMapping = ModelDefaultMapping = dict
 
+    torch = LazyLoader("torch", globals(), "torch")
+    tf = LazyLoader("tf", globals(), "tensorflow")
     cloudpickle = LazyLoader(
         "cloudpickle",
         globals(),
@@ -853,7 +861,13 @@ def save_model(
         if metadata is None:
             metadata = {}
 
-        metadata["_pretrained_class"] = pretrained.__class__.__name__
+        metadata.update(
+            {
+                "_pretrained_class": pretrained.__class__.__name__,
+                "_framework": pretrained.framework,
+            }
+        )
+
         with bentoml.models.create(
             name,
             module=MODULE_NAME,
@@ -886,19 +900,56 @@ def get_runnable(bento_model: bentoml.Model) -> type[bentoml.Runnable]:
             super().__init__()
 
             available_gpus = os.getenv("CUDA_VISIBLE_DEVICES", "")
+            # assign CPU resources
+            kwargs = {}
+
             if available_gpus not in ("", "-1"):
                 # assign GPU resources
                 if not available_gpus.isdigit():
                     raise ValueError(
                         f"Expecting numeric value for CUDA_VISIBLE_DEVICES, got {available_gpus}."
                     )
-                else:
-                    kwargs = {"device": int(available_gpus)}
+                if "_pretrained_class" not in bento_model.info.metadata:
+                    # NOTE: then this is a pipeline. We then pass the device to it.
+                    kwargs["device"] = int(available_gpus)
+            if "_pretrained_class" not in bento_model.info.metadata:
+                self.model = load_model(bento_model, **kwargs)
             else:
-                # assign CPU resources
-                kwargs = {}
-
-            self.model = load_model(bento_model, **kwargs)
+                if "_framework" in bento_model.info.metadata:
+                    if "torch" == bento_model.info.metadata["_framework"]:
+                        self.model = t.cast(
+                            transformers.PreTrainedModel,
+                            load_model(bento_model, **kwargs),
+                        ).to(
+                            torch.device(
+                                "cuda" if available_gpus not in ("", "-1") else "cpu"
+                            )
+                        )
+                        torch.set_default_tensor_type("torch.cuda.FloatTensor")
+                    elif "tf" == bento_model.info.metadata["_framework"]:
+                        with tf.device(
+                            "/device:CPU:0"
+                            if available_gpus in ("", "-1")
+                            else f"/device:GPU:{available_gpus}"
+                        ):
+                            self.model = t.cast(
+                                transformers.TFPreTrainedModel,
+                                load_model(bento_model, **kwargs),
+                            )
+                    else:
+                        # NOTE: we need to hide all GPU from TensorFlow, otherwise it will try to allocate
+                        # memory on the GPU and make it unavailable for JAX.
+                        tf.config.experimental.set_visible_devices([], "GPU")
+                        self.model = t.cast(
+                            transformers.FlaxPreTrainedModel,
+                            load_model(bento_model, **kwargs),
+                        )
+                else:
+                    logger.warning(
+                        "Current %s is saved with an older version of BentoML. Setting GPUs on this won't work as expected. Make sure to save it with a newer version of BentoML.",
+                        bento_model,
+                    )
+                    self.model = load_model(bento_model, **kwargs)
 
             # NOTE: backward compatibility with previous BentoML versions.
             self.pipeline = self.model
diff --git a/src/bentoml/_internal/io_descriptors/json.py b/src/bentoml/_internal/io_descriptors/json.py
index 6cd6ea23..dbd77c83 100644
--- a/src/bentoml/_internal/io_descriptors/json.py
+++ b/src/bentoml/_internal/io_descriptors/json.py
@@ -29,6 +29,10 @@ if t.TYPE_CHECKING:
 
     import pydantic
     import pydantic.schema as schema
+
+    if pkg_version_info("pydantic")[0] >= 2:
+        import pydantic.json_schema as jschema
+
     from google.protobuf import message as _message
     from google.protobuf import struct_pb2
     from typing_extensions import Self
@@ -39,6 +43,9 @@ if t.TYPE_CHECKING:
 else:
     pydantic = LazyLoader("pydantic", globals(), "pydantic", exc_msg=EXC_MSG)
     schema = LazyLoader("schema", globals(), "pydantic.schema", exc_msg=EXC_MSG)
+    jschema = LazyLoader(
+        "jschema", globals(), "pydantic.json_schema", exc_msg="Pydantic v2 is required."
+    )
     # lazy load our proto generated.
     struct_pb2 = LazyLoader("struct_pb2", globals(), "google.protobuf.struct_pb2")
     # lazy load numpy for processing ndarray.
@@ -184,10 +191,6 @@ class JSON(
         json_encoder: type[json.JSONEncoder] = DefaultJsonEncoder,
     ):
         if pydantic_model is not None:
-            if pkg_version_info("pydantic")[0] >= 2:
-                raise BadInput(
-                    "pydantic 2.x is not yet supported. Add upper bound to 'pydantic': 'pip install \"pydantic<2\"'"
-                ) from None
             assert issubclass(
                 pydantic_model, pydantic.BaseModel
             ), "'pydantic_model' must be a subclass of 'pydantic.BaseModel'."
@@ -306,15 +309,24 @@ class JSON(
             return Schema(type="object")
 
         # returns schemas from pydantic_model.
-        return Schema(
-            **schema.model_process_schema(
-                self._pydantic_model,
-                model_name_map=schema.get_model_name_map(
-                    schema.get_flat_models_from_model(self._pydantic_model)
-                ),
-                ref_prefix=REF_PREFIX,
-            )[0]
-        )
+        if pkg_version_info("pydantic")[0] >= 2:
+            json_schema = jschema.model_json_schema(
+                self._pydantic_model, ref_template=REF_PREFIX + "{model}"
+            )
+            # NOTE: we don't need def here, as these will be available in openapi.components.
+            if "$defs" in json_schema:
+                json_schema.pop("$defs", None)
+            return Schema(**json_schema)
+        else:
+            return Schema(
+                **schema.model_process_schema(
+                    self._pydantic_model,
+                    model_name_map=schema.get_model_name_map(
+                        schema.get_flat_models_from_model(self._pydantic_model)
+                    ),
+                    ref_prefix=REF_PREFIX,
+                )[0]
+            )
 
     def openapi_components(self) -> dict[str, t.Any] | None:
         if not self._pydantic_model:
@@ -329,7 +341,10 @@ class JSON(
             if LazyType["pydantic.BaseModel"]("pydantic.BaseModel").isinstance(
                 self.sample
             ):
-                return self.sample.dict()
+                if pkg_version_info("pydantic")[0] >= 2:
+                    return self.sample.model_dump()
+                else:
+                    return self.sample.dict()
             elif isinstance(self.sample, (str, list)):
                 return json.dumps(
                     self.sample,
@@ -373,7 +388,10 @@ class JSON(
 
         if self._pydantic_model:
             try:
-                pydantic_model = self._pydantic_model.parse_obj(json_obj)
+                if pkg_version_info("pydantic")[0] >= 2:
+                    pydantic_model = self._pydantic_model.model_validate(json_obj)
+                else:
+                    pydantic_model = self._pydantic_model.parse_obj(json_obj)
                 return pydantic_model
             except pydantic.ValidationError as e:
                 raise BadInput(f"Invalid JSON input received: {e}") from None
@@ -385,7 +403,10 @@ class JSON(
     ):
         # This is to prevent cases where custom JSON encoder is used.
         if LazyType["pydantic.BaseModel"]("pydantic.BaseModel").isinstance(obj):
-            obj = obj.dict()
+            if pkg_version_info("pydantic")[0] >= 2:
+                obj = obj.model_dump()
+            else:
+                obj = obj.dict()
 
         json_str = (
             json.dumps(
@@ -419,7 +440,12 @@ class JSON(
             content = field
             if self._pydantic_model:
                 try:
-                    return self._pydantic_model.parse_raw(content)
+                    if pkg_version_info("pydantic")[0] >= 2:
+                        return self._pydantic_model.model_validate_json(
+                            json.loads(content)
+                        )
+                    else:
+                        return self._pydantic_model.parse_raw(content)
                 except pydantic.ValidationError as e:
                     raise BadInput(f"Invalid JSON input received: {e}") from None
             try:
@@ -432,14 +458,20 @@ class JSON(
 
             if self._pydantic_model:
                 try:
-                    return self._pydantic_model.parse_obj(parsed)
+                    if pkg_version_info("pydantic")[0] >= 2:
+                        return self._pydantic_model.model_validate(parsed)
+                    else:
+                        return self._pydantic_model.parse_obj(parsed)
                 except pydantic.ValidationError as e:
                     raise BadInput(f"Invalid JSON input received: {e}") from None
         return parsed
 
     async def to_proto(self, obj: JSONType) -> struct_pb2.Value:
         if LazyType["pydantic.BaseModel"]("pydantic.BaseModel").isinstance(obj):
-            obj = obj.dict()
+            if pkg_version_info("pydantic")[0] >= 2:
+                obj = obj.model_dump()
+            else:
+                obj = obj.dict()
         msg = struct_pb2.Value()
         return parse_dict_to_proto(obj, msg, json_encoder=self._json_encoder)
 
diff --git a/src/bentoml/_internal/service/openapi/utils.py b/src/bentoml/_internal/service/openapi/utils.py
index d2212531..6e58d28e 100644
--- a/src/bentoml/_internal/service/openapi/utils.py
+++ b/src/bentoml/_internal/service/openapi/utils.py
@@ -9,12 +9,16 @@ from bentoml.exceptions import InvalidArgument
 from bentoml.exceptions import InternalServerError
 
 from ...utils import LazyLoader
+from ...utils.pkg import pkg_version_info
 from .specification import Schema
 
 if TYPE_CHECKING:
     import pydantic
     import pydantic.schema as schema
 
+    if pkg_version_info("pydantic")[0] >= 2:
+        import pydantic.json_schema as jschema
+
     from bentoml.exceptions import BentoMLException
 
 else:
@@ -23,11 +27,33 @@ else:
     )
     pydantic = LazyLoader("pydantic", globals(), "pydantic", exc_msg=_exc_msg)
     schema = LazyLoader("schema", globals(), "pydantic.schema", exc_msg=_exc_msg)
+    jschema = LazyLoader(
+        "jschema", globals(), "pydantic.json_schema", exc_msg="Pydantic v2 is required."
+    )
 
 REF_PREFIX = "#/components/schemas/"
 
 
-def pydantic_components_schema(pydantic_model: t.Type[pydantic.BaseModel]):
+def pydantic_components_schema(pydantic_model: type[pydantic.BaseModel]):
+    if pkg_version_info("pydantic")[0] >= 2:
+        return pydantic_v2_components_schema(pydantic_model)
+    else:
+        return pydantic_v1_components_schema(pydantic_model)
+
+
+def pydantic_v2_components_schema(pydantic_model: type[pydantic.BaseModel]):
+    json_schema = jschema.model_json_schema(
+        pydantic_model, ref_template=REF_PREFIX + "{model}"
+    )
+    defs = json_schema.pop("$defs", None)
+    components: dict[str, Schema] = {pydantic_model.__name__: Schema(**json_schema)}
+    if defs is not None:
+        # NOTE: This is a nested models, hence we will update the definitions
+        components.update({k: Schema(**v) for k, v in defs.items()})
+    return components
+
+
+def pydantic_v1_components_schema(pydantic_model: type[pydantic.BaseModel]):
     flat_models = schema.get_flat_models_from_model(pydantic_model)
     model_name_map = schema.get_model_name_map(flat_models)
 
diff --git a/src/bentoml/_internal/utils/__init__.py b/src/bentoml/_internal/utils/__init__.py
index c635c051..394b1a7b 100644
--- a/src/bentoml/_internal/utils/__init__.py
+++ b/src/bentoml/_internal/utils/__init__.py
@@ -58,7 +58,6 @@ __all__ = [
     "reserve_free_port",
     "LazyLoader",
     "validate_or_create_dir",
-    "display_path_under_home",
     "rich_console",
     "experimental",
     "compose",
@@ -162,16 +161,6 @@ def calc_dir_size(path: PathType) -> int:
     return sum(f.stat().st_size for f in Path(path).glob("**/*") if f.is_file())
 
 
-def display_path_under_home(path: str) -> str:
-    # Shorten path under home directory with leading `~`
-    # e.g. from `/Users/foo/bar` to just `~/bar`
-    try:
-        return str("~" / Path(path).relative_to(Path.home()))
-    except ValueError:
-        # when path is not under home directory, return original full path
-        return path
-
-
 def human_readable_size(size: t.Union[int, float], decimal_places: int = 2) -> str:
     for unit in ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]:
         if size < 1024.0 or unit == "PiB":
diff --git a/src/bentoml/bentos.py b/src/bentoml/bentos.py
index 3c9968ed..61a714c9 100644
--- a/src/bentoml/bentos.py
+++ b/src/bentoml/bentos.py
@@ -26,6 +26,9 @@ if t.TYPE_CHECKING:
     from .server import Server
     from ._internal.bento import BentoStore
     from ._internal.yatai_client import YataiClient
+    from ._internal.bento.build_config import CondaOptions
+    from ._internal.bento.build_config import DockerOptions
+    from ._internal.bento.build_config import PythonOptions
 
 
 logger = logging.getLogger(__name__)
@@ -268,11 +271,11 @@ def build(
     name: str | None = None,
     labels: dict[str, str] | None = None,
     description: str | None = None,
-    include: list[str] | None = None,
-    exclude: list[str] | None = None,
-    docker: dict[str, t.Any] | None = None,
-    python: dict[str, t.Any] | None = None,
-    conda: dict[str, t.Any] | None = None,
+    include: t.List[str] | None = None,
+    exclude: t.List[str] | None = None,
+    docker: DockerOptions | dict[str, t.Any] | None = None,
+    python: PythonOptions | dict[str, t.Any] | None = None,
+    conda: CondaOptions | dict[str, t.Any] | None = None,
     version: str | None = None,
     build_ctx: str | None = None,
     _bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
diff --git a/src/bentoml_cli/bentos.py b/src/bentoml_cli/bentos.py
index 57accad2..3f39df6b 100644
--- a/src/bentoml_cli/bentos.py
+++ b/src/bentoml_cli/bentos.py
@@ -55,7 +55,6 @@ def add_bento_management_commands(cli: Group):
     from bentoml._internal.utils import calc_dir_size
     from bentoml._internal.utils import human_readable_size
     from bentoml._internal.utils import resolve_user_filepath
-    from bentoml._internal.utils import display_path_under_home
     from bentoml._internal.bento.bento import Bento
     from bentoml._internal.bento.bento import DEFAULT_BENTO_BUILD_FILE
     from bentoml._internal.configuration import get_quiet_mode
@@ -114,7 +113,6 @@ def add_bento_management_commands(cli: Group):
         res = [
             {
                 "tag": str(bento.tag),
-                "path": display_path_under_home(bento.path),
                 "size": human_readable_size(calc_dir_size(bento.path)),
                 "creation_time": bento.info.creation_time.astimezone().strftime(
                     "%Y-%m-%d %H:%M:%S"
@@ -136,13 +134,11 @@ def add_bento_management_commands(cli: Group):
             table.add_column("Tag")
             table.add_column("Size")
             table.add_column("Creation Time")
-            table.add_column("Path")
             for bento in res:
                 table.add_row(
                     bento["tag"],
                     bento["size"],
                     bento["creation_time"],
-                    bento["path"],
                 )
             console.print(table)
 


Test patch:
diff --git a/requirements/tests-requirements.txt b/requirements/tests-requirements.txt
index 9910c0e3..0acba429 100644
--- a/requirements/tests-requirements.txt
+++ b/requirements/tests-requirements.txt
@@ -2,9 +2,9 @@
 pandas
 pydantic
 scikit-learn
-yamllint==1.31.0
+yamllint==1.32.0
 black[jupyter]==22.12.0
-coverage[toml]==7.2.5
+coverage[toml]==7.2.6
 setuptools>=63
 isort==5.11.4
 ruff==0.0.269
diff --git a/tests/unit/_internal/io/test_json.py b/tests/unit/_internal/io/test_json.py
index 81f576ee..63365d84 100644
--- a/tests/unit/_internal/io/test_json.py
+++ b/tests/unit/_internal/io/test_json.py
@@ -16,10 +16,8 @@ import pydantic
 
 from bentoml.io import JSON
 from bentoml.exceptions import BadInput
-from bentoml.exceptions import UnprocessableEntity
 from bentoml.grpc.utils import import_generated_stubs
 from bentoml._internal.utils import LazyLoader
-from bentoml._internal.utils.pkg import pkg_version_info
 from bentoml._internal.io_descriptors.json import DefaultJsonEncoder
 
 if TYPE_CHECKING:
@@ -80,16 +78,6 @@ dumps = partial(
 )
 
 
-@pytest.mark.skipif(
-    pkg_version_info("pydantic")[0] < 2 and pkg_version_info("bentoml")[:2] <= (1, 1),
-    reason="Pydantic 2.x is not yet supported until official releases of Pydantic.",
-)
-def test_not_yet_supported_pydantic():
-    with pytest.raises(UnprocessableEntity) as exc_info:
-        JSON(pydantic_model=Nested)
-    assert "pydantic 2.x is not yet supported" in str(exc_info.value)
-
-
 def test_invalid_init():
     with pytest.raises(AssertionError) as exc_info:
         JSON(pydantic_model=ExampleAttrsClass)  # type: ignore (testing exception)