diff --git a/src/bentoml/_internal/bento/bento.py b/src/bentoml/_internal/bento/bento.py index d287dc1e92f..50b85ae70be 100644 --- a/src/bentoml/_internal/bento/bento.py +++ b/src/bentoml/_internal/bento/bento.py @@ -23,6 +23,7 @@ from ...exceptions import BentoMLException from ...exceptions import InvalidArgument +from ...exceptions import NotFound from ..configuration import BENTOML_VERSION from ..configuration.containers import BentoMLContainer from ..models import ModelStore @@ -352,12 +353,22 @@ def export( finally: self._fs.removetree("models") - @property - def path(self) -> str: - return self.path_of("/") - - def path_of(self, item: str) -> str: - return self._fs.getsyspath(item) + @inject + def total_size( + self, model_store: ModelStore = Provide[BentoMLContainer.model_store] + ) -> int: + total_size = self.file_size + local_model_store = self._model_store + for model in self.info.models: + if local_model_store is not None: + try: + local_model_store.get(model.tag) + continue + except NotFound: + pass + global_model = model_store.get(model.tag) + total_size += global_model.file_size + return total_size def flush_info(self): with self._fs.open(BENTO_YAML_FILENAME, "w") as bento_yaml: diff --git a/src/bentoml/_internal/cloud/bentocloud.py b/src/bentoml/_internal/cloud/bentocloud.py index 169c6469e34..0cdd38034cf 100644 --- a/src/bentoml/_internal/cloud/bentocloud.py +++ b/src/bentoml/_internal/cloud/bentocloud.py @@ -167,7 +167,7 @@ def push_model(model: Model) -> None: apis=apis, models=models, runners=runners, - size_bytes=calc_dir_size(bento.path), + size_bytes=bento.total_size(), ) if not remote_bento: with self.spin( diff --git a/src/bentoml/_internal/cloud/yatai.py b/src/bentoml/_internal/cloud/yatai.py index 8f0c8b36a22..eb9ecd0c055 100644 --- a/src/bentoml/_internal/cloud/yatai.py +++ b/src/bentoml/_internal/cloud/yatai.py @@ -157,7 +157,7 @@ def push_model(model: Model) -> None: apis=apis, models=models, runners=runners, - size_bytes=calc_dir_size(bento.path), + size_bytes=bento.total_size(), ) if not remote_bento: with self.spin(text=f'Registering Bento "{bento.tag}" with Yatai..'): diff --git a/src/bentoml/_internal/models/model.py b/src/bentoml/_internal/models/model.py index e3b0455353c..7b66e9ff039 100644 --- a/src/bentoml/_internal/models/model.py +++ b/src/bentoml/_internal/models/model.py @@ -233,13 +233,6 @@ def from_fs(cls: t.Type[Model], item_fs: FS) -> Model: return res - @property - def path(self) -> str: - return self.path_of("/") - - def path_of(self, item: str) -> str: - return self._fs.getsyspath(item) - @classmethod def enter_cloudpickle_context( cls, diff --git a/src/bentoml/_internal/store.py b/src/bentoml/_internal/store.py index fc19bf77546..c80a02da6b7 100644 --- a/src/bentoml/_internal/store.py +++ b/src/bentoml/_internal/store.py @@ -16,6 +16,7 @@ from .exportable import Exportable from .tag import Tag from .types import PathType +from .utils import calc_dir_size T = t.TypeVar("T") @@ -43,6 +44,17 @@ def _export_name(self) -> str: def creation_time(self) -> datetime.datetime: raise NotImplementedError + @property + def path(self) -> str: + return self.path_of("/") + + def path_of(self, item: str) -> str: + return self._fs.getsyspath(item) + + @property + def file_size(self) -> int: + return calc_dir_size(self.path) + def __repr__(self): return f'{self.get_typename()}(tag="{self.tag}")' diff --git a/src/bentoml_cli/bentos.py b/src/bentoml_cli/bentos.py index ec1af7d538c..0a44d5ef935 100644 --- a/src/bentoml_cli/bentos.py +++ b/src/bentoml_cli/bentos.py @@ -66,7 +66,6 @@ def add_bento_management_commands(cli: Group): from bentoml._internal.bento.build_config import BentoBuildConfig from bentoml._internal.configuration import get_quiet_mode from bentoml._internal.configuration.containers import BentoMLContainer - 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 rich_console as console @@ -121,18 +120,20 @@ def list_bentos(bento_name: str, output: str) -> None: # type: ignore (not acce $ bentoml list FraudDetector """ bentos = bento_store.list(bento_name) - res = [ - { - "tag": str(bento.tag), - "size": human_readable_size(calc_dir_size(bento.path)), - "creation_time": bento.info.creation_time.astimezone().strftime( - "%Y-%m-%d %H:%M:%S" - ), - } - for bento in sorted( - bentos, key=lambda x: x.info.creation_time, reverse=True + res: list[dict[str, str]] = [] + for bento in sorted(bentos, key=lambda x: x.info.creation_time, reverse=True): + bento_size = bento.file_size + model_size = bento.total_size() - bento_size + res.append( + { + "tag": str(bento.tag), + "size": human_readable_size(bento_size), + "model_size": human_readable_size(model_size), + "creation_time": bento.info.creation_time.astimezone().strftime( + "%Y-%m-%d %H:%M:%S" + ), + } ) - ] if output == "json": info = json.dumps(res, indent=2) @@ -144,11 +145,13 @@ def list_bentos(bento_name: str, output: str) -> None: # type: ignore (not acce table = Table(box=None) table.add_column("Tag") table.add_column("Size") + table.add_column("Model Size") table.add_column("Creation Time") for bento in res: table.add_row( bento["tag"], bento["size"], + bento["model_size"], bento["creation_time"], ) console.print(table)