freeze model info and validate metadata entries (bentoml#2363)

* fix(internal): prevent updating Bento.info

* feat(internal): add validate_labels and validate_metadata

* feat(internal): allow LazyType to take tuples of types

* feat(internal): freeze model info on save

* chore(internal): format code

* chore(internal): fix typing errors

Author: sauyon

bentoml/_internal/bento/bento.py

@@ -96,13 +96,11 @@ class Bento(StoreItem):
     _model_store: ModelStore
     _doc: t.Optional[str] = None
 
-    _flushed: bool = False
-
     @staticmethod
     def _export_ext() -> str:
         return "bento"
 
-    @__fs.validator  # type: ignore
+    @__fs.validator  # type:ignore (attrs validators not supported by pyright)
     def check_fs(self, _attr: t.Any, new_fs: "FS"):
         try:
             new_fs.makedir("models", recreate=True)
@@ -116,7 +114,7 @@ def __init__(self, tag: Tag, bento_fs: "FS", info: "BentoInfo"):
         self._tag = tag
         self.__fs = bento_fs
         self.check_fs(None, bento_fs)
-        self.info = info
+        self._info = info
         self.validate()
 
     @property
@@ -129,14 +127,8 @@ def _fs(self) -> "FS":
 
     @property
     def info(self) -> "BentoInfo":
-        self._flushed = False
         return self._info
 
-    @info.setter
-    def info(self, new_info: "BentoInfo"):
-        self._flushed = False
-        self._info = new_info
-
     @classmethod
     @inject
     def create(
@@ -300,7 +292,6 @@ def from_fs(cls, item_fs: "FS") -> "Bento":
                 f"Failed to create bento because it contains an invalid '{BENTO_YAML_FILENAME}'"
             )
 
-        res._flushed = True
         return res
 
     @property
@@ -311,14 +302,9 @@ def path_of(self, item: str) -> str:
         return self._fs.getsyspath(item)
 
     def flush_info(self):
-        if self._flushed:
-            return
-
         with self._fs.open(BENTO_YAML_FILENAME, "w") as bento_yaml:
             self.info.dump(bento_yaml)
 
-        self._flushed = True
-
     @property
     def doc(self) -> str:
         if self._doc is not None:
@@ -337,8 +323,6 @@ def save(
         self,
         bento_store: "BentoStore" = Provide[BentoMLContainer.bento_store],
     ) -> "Bento":
-        self.flush_info()
-
         if not self.validate():
             logger.warning(f"Failed to create Bento for {self.tag}, not saving.")
             raise BentoMLException("Failed to save Bento because it was invalid.")
@@ -351,28 +335,6 @@ def save(
 
         return self
 
-    def export(
-        self,
-        path: str,
-        output_format: t.Optional[str] = None,
-        *,
-        protocol: t.Optional[str] = None,
-        user: t.Optional[str] = None,
-        passwd: t.Optional[str] = None,
-        params: t.Optional[t.Dict[str, str]] = None,
-        subpath: t.Optional[str] = None,
-    ) -> str:
-        self.flush_info()
-        return super().export(
-            path,
-            output_format,
-            protocol=protocol,
-            user=user,
-            passwd=passwd,
-            params=params,
-            subpath=subpath,
-        )
-
     def validate(self):
         return self._fs.isfile(BENTO_YAML_FILENAME)
 
@@ -445,8 +407,6 @@ class BentoInfo:
     runners: t.List[BentoRunnerInfo] = attr.field(factory=list)
     apis: t.List[BentoApiInfo] = attr.field(factory=list)
 
-    _flushed: bool = False
-
     def __attrs_post_init__(self):
         # Direct set is not available when frozen=True
         object.__setattr__(self, "name", self.tag.name)
@@ -501,11 +461,8 @@ def validate(self):
 
 bentoml_cattr.register_unstructure_hook(
     BentoInfo,
-    # Ignore internal private state "_flushed"
     # Ignore tag, tag is saved via the name and version field
-    make_dict_unstructure_fn(
-        BentoInfo, bentoml_cattr, _flushed=override(omit=True), tag=override(omit=True)
-    ),
+    make_dict_unstructure_fn(BentoInfo, bentoml_cattr, tag=override(omit=True)),
 )
 
 

bentoml/_internal/models/model.py

@@ -18,6 +18,8 @@
 from ..tag import Tag
 from ..store import Store
 from ..store import StoreItem
+from ..utils import validate_labels
+from ..utils import validate_metadata
 from ...exceptions import NotFound
 from ...exceptions import BentoMLException
 from ..configuration import BENTOML_VERSION
@@ -42,8 +44,7 @@ class Model(StoreItem):
     _info: "ModelInfo"
     _custom_objects: t.Optional[t.Dict[str, t.Any]] = None
 
-    _info_flushed = False
-    _custom_objects_flushed = False
+    _flushed: bool = False
 
     @staticmethod
     def _export_ext() -> str:
@@ -59,24 +60,20 @@ def _fs(self) -> "FS":
 
     @property
     def info(self) -> "ModelInfo":
-        self._info_flushed = False
         return self._info
 
-    @info.setter
-    def info(self, new_info: "ModelInfo"):
-        self._info_flushed = False
-        self._info = new_info
-
     @property
     def custom_objects(self) -> t.Dict[str, t.Any]:
-        self._custom_objects_flushed = False
-
         if self._custom_objects is None:
             if self._fs.isfile(CUSTOM_OBJECTS_FILENAME):
                 with self._fs.open(CUSTOM_OBJECTS_FILENAME, "rb") as cofile:
-                    self._custom_objects = cloudpickle.load(cofile)
+                    self._custom_objects: "t.Optional[t.Dict[str, t.Any]]" = (
+                        cloudpickle.load(cofile)
+                    )
+                    if not isinstance(self._custom_objects, dict):
+                        raise ValueError("Invalid custom objects found.")
             else:
-                self._custom_objects = {}
+                self._custom_objects: "t.Optional[t.Dict[str, t.Any]]" = {}
 
         return self._custom_objects
 
@@ -128,6 +125,9 @@ def create(
         context["bentoml_version"] = BENTOML_VERSION
         context["python_version"] = PYTHON_VERSION
 
+        validate_labels(labels)
+        validate_metadata(metadata)
+
         model_fs = fs.open_fs(f"temp://bentoml_model_{name}")
 
         res = Model(
@@ -155,8 +155,7 @@ def save(
         return self
 
     def _save(self, model_store: "ModelStore") -> "Model":
-        self.flush_info()
-        self.flush_custom_objects()
+        self.flush()
 
         if not self.validate():
             logger.warning(f"Failed to create Model for {self.tag}, not saving.")
@@ -187,8 +186,7 @@ def from_fs(cls, item_fs: FS) -> "Model":
                 f"Failed to load bento model because it contains an invalid '{MODEL_YAML_FILENAME}'"
             )
 
-        res._info_flushed = True
-        res._custom_objects_flushed = True
+        res._flushed = True
         return res
 
     @property
@@ -198,26 +196,22 @@ def path(self) -> str:
     def path_of(self, item: str) -> str:
         return self._fs.getsyspath(item)
 
-    def flush_info(self):
-        if self._info_flushed:
-            return
+    def flush(self):
+        if not self._flushed:
+            self._flush_info()
+            self._flush_custom_objects()
+        self._flushed = True
 
+    def _flush_info(self):
         with self._fs.open(MODEL_YAML_FILENAME, "w") as model_yaml:
             self.info.dump(model_yaml)
 
-        self._info_flushed = True
-
-    def flush_custom_objects(self):
-        if self._custom_objects_flushed:
-            return
-
+    def _flush_custom_objects(self):
         # pickle custom_objects if it is not None and not empty
         if self.custom_objects:
             with self._fs.open(CUSTOM_OBJECTS_FILENAME, "wb") as cofile:
                 cloudpickle.dump(self.custom_objects, cofile)
 
-        self._custom_objects_flushed = True
-
     @property
     def creation_time(self) -> datetime:
         return self.info.creation_time
@@ -233,8 +227,8 @@ def export(
         params: t.Optional[t.Dict[str, str]] = None,
         subpath: t.Optional[str] = None,
     ) -> str:
-        self.flush_info()
-        self.flush_custom_objects()
+        self.flush()
+
         return super().export(
             path,
             output_format,
@@ -264,14 +258,30 @@ def __init__(self, base_path: "t.Union[PathType, FS]"):
 class ModelInfo:
     tag: Tag
     module: str
-    labels: t.Dict[str, t.Any]
+    labels: t.Dict[str, str]
     options: t.Dict[str, t.Any]
     metadata: t.Dict[str, t.Any]
     context: t.Dict[str, t.Any]
     bentoml_version: str = BENTOML_VERSION
     api_version: str = "v1"
     creation_time: datetime = attr.field(factory=lambda: datetime.now(timezone.utc))
 
+    def __eq__(self, other: "t.Union[ModelInfo, FrozenModelInfo]"):
+        if not isinstance(other, (ModelInfo, FrozenModelInfo)):
+            return False
+
+        return (
+            self.tag == other.tag
+            and self.module == other.module
+            and self.labels == other.labels
+            and self.options == other.options
+            and self.metadata == other.metadata
+            and self.context == other.context
+            and self.bentoml_version == other.bentoml_version
+            and self.api_version == other.api_version
+            and self.creation_time == other.creation_time
+        )
+
     def __attrs_post_init__(self):
         self.validate()
 
@@ -292,8 +302,8 @@ def to_dict(self) -> t.Dict[str, t.Any]:
     def dump(self, stream: t.IO[t.Any]):
         return yaml.dump(self, stream, sort_keys=False)
 
-    @classmethod
-    def from_yaml_file(cls, stream: t.IO[t.Any]):
+    @staticmethod
+    def from_yaml_file(stream: t.IO[t.Any]):
         try:
             yaml_content = yaml.safe_load(stream)
         except yaml.YAMLError as exc:  # pragma: no cover - simple error handling
@@ -311,7 +321,7 @@ def from_yaml_file(cls, stream: t.IO[t.Any]):
         del yaml_content["version"]
 
         try:
-            model_info = cls(**yaml_content)  # type: ignore
+            model_info = FrozenModelInfo(**yaml_content)  # type: ignore
         except TypeError:  # pragma: no cover - simple error handling
             raise BentoMLException(f"unexpected field in {MODEL_YAML_FILENAME}")
         return model_info
@@ -321,6 +331,15 @@ def validate(self):
         # add tests when implemented
         ...
 
+    def freeze(self) -> "ModelInfo":
+        self.__class__ = FrozenModelInfo
+        return self
+
+
+@attr.define(repr=False, frozen=True)  # type: ignore (pyright doesn't allow for a frozen subclass)
+class FrozenModelInfo(ModelInfo):
+    pass
+
 
 def copy_model(
     model_tag: t.Union[Tag, str],
@@ -346,3 +365,4 @@ def _ModelInfo_dumper(dumper: yaml.Dumper, info: ModelInfo) -> yaml.Node:
 
 
 yaml.add_representer(ModelInfo, _ModelInfo_dumper)
+yaml.add_representer(FrozenModelInfo, _ModelInfo_dumper)

bentoml/_internal/types.py

@@ -7,6 +7,10 @@
 import urllib.parse
 import urllib.request
 from typing import TYPE_CHECKING
+from datetime import date
+from datetime import time
+from datetime import datetime
+from datetime import timedelta
 from dataclasses import dataclass
 
 from .utils.dataclasses import json_serializer
@@ -31,6 +35,24 @@
 else:
     PathType = t.Union[str, os.PathLike]
 
+MetadataType = t.Union[
+    str,
+    bytes,
+    bool,
+    int,
+    float,
+    complex,
+    datetime,
+    date,
+    time,
+    timedelta,
+    t.List["MetadataType"],
+    t.Tuple["MetadataType"],
+    t.Dict[str, "MetadataType"],
+]
+
+MetadataDict = t.Dict[str, MetadataType]
+
 JSONSerializable = t.NewType("JSONSerializable", object)
 
 
@@ -135,7 +157,9 @@ def __hash__(self) -> int:
     def __repr__(self) -> str:
         return f'LazyType("{self.module}", "{self.qualname}")'
 
-    def get_class(self, import_module: bool = True) -> "t.Type[T]":
+    def get_class(
+        self, import_module: bool = True
+    ) -> "t.Union[t.Type[T], t.Tuple[t.Type[T]]]":
         if self._runtime_class is None:
             try:
                 m = sys.modules[self.module]
@@ -147,6 +171,10 @@ def get_class(self, import_module: bool = True) -> "t.Type[T]":
                 else:
                     raise ValueError(f"Module {self.module} not imported")
 
+            if isinstance(self.qualname, tuple):
+                self._runtime_class = tuple(
+                    (t.cast("t.Type[T]", getattr(m, x)) for x in self.qualname)
+                )
             self._runtime_class = t.cast("t.Type[T]", getattr(m, self.qualname))
 
         return self._runtime_class