refactor(cli): make CLI commands available as modules (#4487)

* feat(cli): make CLI commands avaiable as modules Signed-off-by: Frost Ming <[email protected]> * fix: handle bentoml group Signed-off-by: Frost Ming <[email protected]> * fix: simply aliases Signed-off-by: Frost Ming <[email protected]> * fix: rename bento_command Signed-off-by: Frost Ming <[email protected]> --------- Signed-off-by: Frost Ming <[email protected]>
bentoml · Feb 23, 2024 · 7b0b0e6 · 7b0b0e6
1 parent ed91f8a
commit 7b0b0e6
Show file tree

Hide file tree

Showing 11 changed files with 1,539 additions and 1,445 deletions.
diff --git a/Makefile b/Makefile
@@ -41,6 +41,8 @@ clean: ## Clean all generated files
 	@find . -type f -name '*.py[co]' -delete -o -type d -name __pycache__ -delete
 
 # Docs
+watch-docs: ## Build and watch documentation
+	pdm run sphinx-autobuild docs/source docs/build/html --watch $(GIT_ROOT)/src/ --ignore "bazel-*"
 spellcheck-docs: ## Spell check documentation
 	pdm run sphinx-build -b spelling ./docs/source ./docs/build || (echo "Error running spellchecker.. You may need to run 'make install-spellchecker-deps'"; exit 1)
 OS := $(shell uname)

diff --git a/src/bentoml_cli/bentos.py b/src/bentoml_cli/bentos.py
@@ -18,7 +18,6 @@
 
 if t.TYPE_CHECKING:
     from click import Context
-    from click import Group
     from click import Parameter
 
     from bentoml._internal.bento import BentoStore
@@ -60,7 +59,7 @@ def parse_delete_targets_argument_callback(
     return delete_targets
 
 
-def add_bento_management_commands(cli: Group):
+def bento_management_commands() -> click.Group:
     import bentoml
     from bentoml import Tag
     from bentoml._internal.bento.bento import DEFAULT_BENTO_BUILD_FILE
@@ -72,19 +71,27 @@ def add_bento_management_commands(cli: Group):
     from bentoml._internal.utils import resolve_user_filepath
     from bentoml._internal.utils import rich_console as console
     from bentoml.bentos import import_bento
+    from bentoml_cli.utils import BentoMLCommandGroup
 
-    bento_store = BentoMLContainer.bento_store.get()
-    cloud_client = BentoMLContainer.bentocloud_client.get()
+    @click.group(cls=BentoMLCommandGroup)
+    def bentos():
+        """Commands for managing Bento bundles."""
+        pass
 
-    @cli.command()
+    @bentos.command()
     @click.argument("bento_tag", type=click.STRING)
     @click.option(
         "-o",
         "--output",
         type=click.Choice(["json", "yaml", "path"]),
         default="yaml",
     )
-    def get(bento_tag: str, output: str) -> None:  # type: ignore (not accessed)
+    @inject
+    def get(
+        bento_tag: str,
+        output: str,
+        bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
+    ) -> None:  # type: ignore (not accessed)
         """Print Bento details by providing the bento_tag.
 
         \b
@@ -102,15 +109,20 @@ def get(bento_tag: str, output: str) -> None:  # type: ignore (not accessed)
             info = yaml.dump(bento.info, indent=2, sort_keys=False)
             console.print(Syntax(info, "yaml", background_color="default"))
 
-    @cli.command(name="list")
+    @bentos.command(name="list")
     @click.argument("bento_name", type=click.STRING, required=False)
     @click.option(
         "-o",
         "--output",
         type=click.Choice(["json", "yaml", "table"]),
         default="table",
     )
-    def list_bentos(bento_name: str, output: str) -> None:  # type: ignore (not accessed)
+    @inject
+    def list_bentos(
+        bento_name: str,
+        output: str,
+        bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
+    ) -> None:  # type: ignore (not accessed)
         """List Bentos in local store
 
         \b
@@ -141,7 +153,7 @@ def list_bentos(bento_name: str, output: str) -> None:  # type: ignore (not acce
             info = json.dumps(res, indent=2)
             console.print(info)
         elif output == "yaml":
-            info = yaml.safe_dump(res, indent=2)
+            info = t.cast(str, yaml.safe_dump(res, indent=2))
             console.print(Syntax(info, "yaml", background_color="default"))
         else:
             table = Table(box=None)
@@ -158,7 +170,7 @@ def list_bentos(bento_name: str, output: str) -> None:  # type: ignore (not acce
                 )
             console.print(table)
 
-    @cli.command()
+    @bentos.command()
     @click.argument(
         "delete_targets",
         nargs=-1,
@@ -172,7 +184,12 @@ def list_bentos(bento_name: str, output: str) -> None:  # type: ignore (not acce
         is_flag=True,
         help="Skip confirmation when deleting a specific bento bundle",
     )
-    def delete(delete_targets: list[str], yes: bool) -> None:  # type: ignore (not accessed)
+    @inject
+    def delete(
+        delete_targets: list[str],
+        yes: bool,
+        bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
+    ) -> None:  # type: ignore (not accessed)
         """Delete Bento in local bento store.
 
         \b
@@ -205,15 +222,20 @@ def delete_target(target: str) -> None:
         for target in delete_targets:
             delete_target(target)
 
-    @cli.command()
+    @bentos.command()
     @click.argument("bento_tag", type=click.STRING)
     @click.argument(
         "out_path",
         type=click.STRING,
         default="",
         required=False,
     )
-    def export(bento_tag: str, out_path: str) -> None:  # type: ignore (not accessed)
+    @inject
+    def export(
+        bento_tag: str,
+        out_path: str,
+        bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
+    ) -> None:  # type: ignore (not accessed)
         """Export a Bento to an external file archive
 
         \b
@@ -235,7 +257,7 @@ def export(bento_tag: str, out_path: str) -> None:  # type: ignore (not accessed
         out_path = bento.export(out_path)
         click.echo(f"{bento} exported to {out_path}.")
 
-    @cli.command(name="import")
+    @bentos.command(name="import")
     @click.argument("bento_path", type=click.STRING)
     def import_bento_(bento_path: str) -> None:  # type: ignore (not accessed)
         """Import a previously exported Bento archive file
@@ -252,7 +274,7 @@ def import_bento_(bento_path: str) -> None:  # type: ignore (not accessed)
         bento = import_bento(bento_path)
         click.echo(f"{bento} imported.")
 
-    @cli.command()
+    @bentos.command()
     @click.argument("bento_tag", type=click.STRING)
     @click.option(
         "-f",
@@ -262,13 +284,19 @@ def import_bento_(bento_path: str) -> None:  # type: ignore (not accessed)
         help="Force pull from remote Bento store to local and overwrite even if it already exists in local",
     )
     @click.pass_obj
-    def pull(shared_options: SharedOptions, bento_tag: str, force: bool) -> None:  # type: ignore (not accessed)
+    @inject
+    def pull(
+        shared_options: SharedOptions,
+        bento_tag: str,
+        force: bool,
+        cloud_client: BentoCloudClient = Provide[BentoMLContainer.bentocloud_client],
+    ) -> None:  # type: ignore (not accessed)
         """Pull Bento from a remote Bento store server."""
         cloud_client.pull_bento(
             bento_tag, force=force, context=shared_options.cloud_context
         )
 
-    @cli.command()
+    @bentos.command()
     @click.argument("bento_tag", type=click.STRING)
     @click.option(
         "-f",
@@ -284,8 +312,14 @@ def pull(shared_options: SharedOptions, bento_tag: str, force: bool) -> None:  #
         help="Number of threads to use for upload",
     )
     @click.pass_obj
+    @inject
     def push(
-        shared_options: SharedOptions, bento_tag: str, force: bool, threads: int
+        shared_options: SharedOptions,
+        bento_tag: str,
+        force: bool,
+        threads: int,
+        bento_store: BentoStore = Provide[BentoMLContainer.bento_store],
+        cloud_client: BentoCloudClient = Provide[BentoMLContainer.bentocloud_client],
     ) -> None:  # type: ignore (not accessed)
         """Push Bento to a remote Bento store server."""
         bento_obj = bento_store.get(bento_tag)
@@ -298,7 +332,7 @@ def push(
             context=shared_options.cloud_context,
         )
 
-    @cli.command()
+    @bentos.command()
     @click.argument("build_ctx", type=click.Path(), default=".")
     @click.option(
         "-f",
@@ -429,3 +463,8 @@ def build(  # type: ignore (not accessed)
             bentoml.container.build(bento.tag, backend=backend)
 
         return bento
+
+    return bentos
+
+
+bento_command = bento_management_commands()
diff --git a/src/bentoml_cli/cli.py b/src/bentoml_cli/cli.py
@@ -5,19 +5,19 @@
 import click
 import psutil
 
-from bentoml_cli.bentos import add_bento_management_commands
-from bentoml_cli.cloud import add_cloud_command
-from bentoml_cli.containerize import add_containerize_command
-from bentoml_cli.deployment import add_deployment_command
-from bentoml_cli.env import add_env_command
-from bentoml_cli.models import add_model_management_commands
-from bentoml_cli.serve import add_serve_command
-from bentoml_cli.start import add_start_command
-from bentoml_cli.utils import BentoMLCommandGroup
-
-
-def create_bentoml_cli() -> click.Group:
+
+def create_bentoml_cli() -> click.Command:
     from bentoml._internal.context import component_context
+    from bentoml_cli.bentos import bento_command
+    from bentoml_cli.cloud import cloud_command
+    from bentoml_cli.containerize import containerize_command
+    from bentoml_cli.deployment import deploy_command
+    from bentoml_cli.deployment import deployment_command
+    from bentoml_cli.env import env_command
+    from bentoml_cli.models import model_command
+    from bentoml_cli.serve import serve_command
+    from bentoml_cli.start import start_command
+    from bentoml_cli.utils import BentoMLCommandGroup
 
     component_context.component_type = "cli"
 
@@ -37,14 +37,15 @@ def bentoml_cli():
         """
 
     # Add top-level CLI commands
-    add_env_command(bentoml_cli)
-    add_cloud_command(bentoml_cli)
-    add_bento_management_commands(bentoml_cli)
-    add_model_management_commands(bentoml_cli)
-    add_start_command(bentoml_cli)
-    add_serve_command(bentoml_cli)
-    add_containerize_command(bentoml_cli)
-    add_deployment_command(bentoml_cli)
+    bentoml_cli.add_command(env_command)
+    bentoml_cli.add_command(cloud_command)
+    bentoml_cli.add_command(model_command)
+    bentoml_cli.add_subcommands(bento_command)
+    bentoml_cli.add_subcommands(start_command)
+    bentoml_cli.add_subcommands(serve_command)
+    bentoml_cli.add_command(containerize_command)
+    bentoml_cli.add_command(deploy_command)
+    bentoml_cli.add_command(deployment_command)
 
     if psutil.WINDOWS:
         import sys