From 0088c75f0c58a8fec1b147678756213e44007617 Mon Sep 17 00:00:00 2001
From: Frost Ming <me@frostming.com>
Date: Wed, 6 Nov 2024 16:59:58 +0800
Subject: [PATCH] feat: allow extending bentoml with extra commands (#5064)

* feat: allow extending bentoml with extra commands

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

* fix: ensure newline at EOF

Signed-off-by: Frost Ming <me@frostming.com>
---
 src/bentoml/_internal/cloud/deployment.py |  2 +-
 src/bentoml_cli/cli.py                    |  4 ++++
 src/bentoml_cli/utils.py                  | 12 ++++++++++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/bentoml/_internal/cloud/deployment.py b/src/bentoml/_internal/cloud/deployment.py
index 4b34aac6d73..f1b24a9c6ac 100644
--- a/src/bentoml/_internal/cloud/deployment.py
+++ b/src/bentoml/_internal/cloud/deployment.py
@@ -1465,7 +1465,7 @@ def _build_requirements_txt(bento_dir: str) -> bytes:
     content = b""
     if filename and os.path.exists(fullpath := os.path.join(bento_dir, filename)):
         with open(fullpath, "rb") as f:
-            content = f.read()
+            content = f.read().rstrip(b"\n") + b"\n"
     for package in config.python.packages or []:
         content += f"{package}\n".encode()
     bentoml_requirement = get_bentoml_requirement()
diff --git a/src/bentoml_cli/cli.py b/src/bentoml_cli/cli.py
index 517513503f9..40927a250cf 100644
--- a/src/bentoml_cli/cli.py
+++ b/src/bentoml_cli/cli.py
@@ -19,6 +19,7 @@ def create_bentoml_cli() -> click.Command:
     from bentoml_cli.serve import serve_command
     from bentoml_cli.start import start_command
     from bentoml_cli.utils import BentoMLCommandGroup
+    from bentoml_cli.utils import get_entry_points
 
     server_context.service_type = "cli"
 
@@ -49,6 +50,9 @@ def bentoml_cli():
     bentoml_cli.add_command(develop_command)
     bentoml_cli.add_command(deployment_command)
     bentoml_cli.add_command(secret_command)
+    # Load commands from extensions
+    for ep in get_entry_points("bentoml.commands"):
+        bentoml_cli.add_command(ep.load())
 
     if psutil.WINDOWS:
         import sys
diff --git a/src/bentoml_cli/utils.py b/src/bentoml_cli/utils.py
index 3056e6af9c8..824f00327c9 100644
--- a/src/bentoml_cli/utils.py
+++ b/src/bentoml_cli/utils.py
@@ -7,6 +7,7 @@
 import re
 import time
 import typing as t
+from importlib.metadata import entry_points
 
 import click
 import click_option_group as cog
@@ -14,6 +15,8 @@
 from click.exceptions import UsageError
 
 if t.TYPE_CHECKING:
+    from importlib.metadata import EntryPoint
+
     from click import Command
     from click import Context
     from click import HelpFormatter
@@ -474,3 +477,12 @@ def is_valid_bento_tag(value: str) -> bool:
 
 def is_valid_bento_name(value: str) -> bool:
     return re.match(r"^[A-Za-z_0-9]*$", value) is not None
+
+
+def get_entry_points(group: str) -> t.Iterable[EntryPoint]:
+    """A compatible version of importlib.metadata.entry_points for Python < 3.10"""
+    try:
+        return entry_points(group=group)
+    except TypeError:
+        # For Python < 3.10, entry_points() does not accept group argument
+        return entry_points().get(group, [])