feat: add --with-models option to bentoml pull command (#5117)

* fix: download models before other deployment hooks

Signed-off-by: Frost Ming <me@frostming.com>

* feat: add --with-models to bentoml pull command

Signed-off-by: Frost Ming <me@frostming.com>

src/_bentoml_impl/server/serving.py

@@ -135,17 +135,17 @@ def create_dependency_watcher(
 def server_on_deployment(
     svc: AnyService, result_file: str = Provide[BentoMLContainer.result_store_file]
 ) -> None:
-    for name in dir(svc.inner):
-        member = getattr(svc.inner, name)
-        if callable(member) and getattr(member, "__bentoml_deployment_hook__", False):
-            member()
     # Resolve models before server starts.
     if bento := svc.bento:
         for model in bento.info.all_models:
             model.to_model().resolve()
     else:
         for model in svc.models:
             model.resolve()
+    for name in dir(svc.inner):
+        member = getattr(svc.inner, name)
+        if callable(member) and getattr(member, "__bentoml_deployment_hook__", False):
+            member()
     if os.path.exists(result_file):
         os.remove(result_file)
 

src/bentoml/_internal/cloud/bento.py

@@ -405,6 +405,7 @@ def pull(
         tag: str | Tag,
         *,
         force: bool = False,
+        with_models: bool = False,
         bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
     ) -> Bento:
         """Pull a bento from remote bento store
@@ -425,15 +426,21 @@ def pull(
                 tag,
                 download_task_id,
                 force=force,
+                with_models=with_models,
                 bento_store=bento_store,
             )
 
+    def _pull_bento_models(self, bento: Bento) -> None:
+        for model in bento.info.all_models:
+            model.to_model().resolve()
+
     def _do_pull_bento(
         self,
         tag: str | Tag,
         download_task_id: TaskID,
         *,
         force: bool = False,
+        with_models: bool = False,
         bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
     ) -> Bento:
         rest_client = self._client
@@ -443,6 +450,8 @@ def _do_pull_bento(
                 self.spinner.log(
                     f'[bold blue]Bento "{tag}" exists in local bento store'
                 )
+                if with_models:
+                    self._pull_bento_models(bento)
                 return bento
             bento_store.delete(tag)
         except NotFound:
@@ -525,6 +534,8 @@ def _do_pull_bento(
                     bento = Bento.from_fs(temp_fs)
                     bento = bento.save(bento_store)
                     self.spinner.log(f'[bold green]Successfully pulled bento "{_tag}"')
+                    if with_models:
+                        self._pull_bento_models(bento)
                     return bento
 
     def list(self) -> BentoWithRepositoryListSchema:

src/bentoml_cli/bentos.py

@@ -300,14 +300,16 @@ def import_bento_(bento_path: str) -> None:  # type: ignore (not accessed)
         default=False,
         help="Force pull from remote Bento store to local and overwrite even if it already exists in local",
     )
+    @click.option("--with-models", is_flag=True, default=False, help="Pull models too")
     @inject
     def pull(
         bento_tag: str,
         force: bool,
+        with_models: bool,
         cloud_client: BentoCloudClient = Provide[BentoMLContainer.bentocloud_client],
     ) -> None:  # type: ignore (not accessed)
         """Pull Bento from a remote Bento store server."""
-        cloud_client.bento.pull(bento_tag, force=force)
+        cloud_client.bento.pull(bento_tag, force=force, with_models=with_models)
 
     @bentos.command()
     @click.argument("bento_tag", type=click.STRING)

src/bentoml_cli/models.py

@@ -307,6 +307,8 @@ def pull(
             bentofile = resolve_user_filepath(bentofile, None)
         except FileNotFoundError:
             raise InvalidArgument(f'file "{bentofile}" not found')
+        else:
+            build_config = BentoBuildConfig.from_file(bentofile)
     else:
         for filename in DEFAULT_BENTO_BUILD_FILES:
             if os.path.exists(filename):