feat: tpu_type support (#4493)

* feat: tpu and tpu_type support Signed-off-by: Frost Ming <[email protected]> * fix: tpu_type only Signed-off-by: Frost Ming <[email protected]> * fix schema Signed-off-by: Frost Ming <[email protected]> --------- Signed-off-by: Frost Ming <[email protected]>
bentoml · Feb 7, 2024 · c572e5b · c572e5b
1 parent d86e523
commit c572e5b
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git a/src/_bentoml_sdk/service/config.py b/src/_bentoml_sdk/service/config.py
@@ -44,6 +44,21 @@ class ResourceSchema(TypedDict, total=False):
         "nvidia-tesla-k80",
         "nvidia-tesla-p4",
     ]
+    tpu_type: Literal[
+        "v4-2x2x1",
+        "v4-2x2x2",
+        "v4-2x2x4",
+        "v4-2x4x4",
+        "v5p-2x2x1",
+        "v5p-2x2x2",
+        "v5p-2x2x4",
+        "v5p-2x4x4",
+        "v5e-1x1",
+        "v5e-2x2",
+        "v5e-2x4",
+        "v5e-4x4",
+        "v5e-4x8",
+    ]
 
 
 WorkerSchema = Union[Posint, Literal["cpu_count"]]

diff --git a/src/_bentoml_sdk/service/factory.py b/src/_bentoml_sdk/service/factory.py
@@ -25,7 +25,6 @@
 
 from ..api import APIMethod
 from .config import ServiceConfig as Config
-from .config import validate
 
 logger = logging.getLogger("bentoml.io")
 
@@ -345,7 +344,7 @@ class InferenceService:
             def predict(self, input: str) -> str:
                 return input
     """
-    config = validate(kwargs)
+    config = kwargs
 
     def decorator(inner: type[T]) -> Service[T]:
         if isinstance(inner, Service):

diff --git a/src/bentoml/_internal/configuration/v2/__init__.py b/src/bentoml/_internal/configuration/v2/__init__.py
@@ -72,8 +72,15 @@
     # NOTE: there is a distinction between being unset and None here; if set to 'None'
     # in configuration for a specific runner, it will override the global configuration.
     s.Optional("resources"): s.Or(
-        {s.Optional(str): object}, lambda s: s == "system", None
-    ),  # type: ignore (incomplete schema typing)
+        {
+            s.Optional("cpu"): str,
+            s.Optional("memory"): str,
+            s.Optional("gpu"): s.And(Real, ensure_larger_than_zero),
+            s.Optional("gpu_type"): str,
+            s.Optional("tpu_type"): str,
+        },
+        None,
+    ),
     s.Optional("workers"): s.Or(
         lambda s: s == "cpu_count",
         s.And(int, ensure_larger_than_zero),

diff --git a/typings/schema.pyi b/typings/schema.pyi