Add usage stats

bentoctl/cli/__init__.py

@@ -3,6 +3,11 @@
 import click
 
 from bentoctl import __version__
+from bentoctl.utils.terraform import (
+    terraform_apply,
+    terraform_destroy,
+    is_terraform_applied,
+)
 from bentoctl.cli.interactive import deployment_config_builder
 from bentoctl.cli.operator_management import get_operator_management_subcommands
 from bentoctl.cli.utils import BentoctlCommandGroup, handle_bentoctl_exceptions
@@ -17,8 +22,8 @@
     push_docker_image_to_repository,
     tag_docker_image,
 )
-from bentoctl.terraform import is_terraform_applied, terraform_apply, terraform_destroy
-from bentoctl.utils import TempDirectory, get_debug_mode
+from bentoctl.utils import get_debug_mode
+from bentoctl.utils.temp_dir import TempDirectory
 
 CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
 
@@ -79,6 +84,7 @@ def init(save_path, do_not_generate):
     if not do_not_generate:
         generated_files = deployment_config.generate(destination_dir=save_path)
         print_generated_files_list(generated_files)
+    return deployment_config
 
 
 @bentoctl.command()
@@ -110,6 +116,7 @@ def generate(deployment_config_file, values_only, save_path):
         destination_dir=save_path, values_only=values_only
     )
     print_generated_files_list(generated_files)
+    return deployment_config
 
 
 @bentoctl.command()
@@ -171,6 +178,7 @@ def build(
         print_generated_files_list(generated_files)
     else:
         console.print(f"[green]Created docker image: {local_docker_tag}[/]")
+    return deployment_config
 
 
 @bentoctl.command()
@@ -193,6 +201,7 @@ def destroy(deployment_config_file):
         terraform_destroy()
     deployment_config.delete_repository()
     console.print(f"Deleted the repository {deployment_config.repository_name}")
+    return deployment_config
 
 
 @bentoctl.command()
@@ -210,6 +219,7 @@ def apply(deployment_config_file):
     deployment_config = DeploymentConfig.from_file(deployment_config_file)
     if deployment_config.template_type.startswith("terraform"):
         terraform_apply()
+    return deployment_config
 
 
 # subcommands

bentoctl/cli/utils.py

@@ -1,10 +1,18 @@
 import functools
 import sys
+import time
+import os
 
 import click
 
 from bentoctl.exceptions import BentoctlException
 from bentoctl.utils import set_debug_mode
+from bentoctl.utils.usage_stats import (
+    BENTOML_DO_NOT_TRACK,
+    cli_events_map,
+    CliEvent,
+    track,
+)
 
 DEBUG_ENV_VAR = "BENTOCTL_DEBUG"
 
@@ -22,7 +30,7 @@ def wrapper(*args, **kwargs):
 
 
 class BentoctlCommandGroup(click.Group):
-    NUMBER_OF_COMMON_PARAMS = 1
+    NUMBER_OF_COMMON_PARAMS = 2
 
     @staticmethod
     def bentoctl_common_params(func):
@@ -33,6 +41,13 @@ def bentoctl_common_params(func):
             default=False,
             help="Show debug logs when running the command",
         )
+        @click.option(
+            "--do-not-track",
+            is_flag=True,
+            default=False,
+            envvar=BENTOML_DO_NOT_TRACK,
+            help="Do not send uage info",
+        )
         @functools.wraps(func)
         def wrapper(verbose, *args, **kwargs):
             set_debug_mode(verbose)
@@ -41,10 +56,59 @@ def wrapper(verbose, *args, **kwargs):
 
         return wrapper
 
+    @staticmethod
+    def bentoctl_track_usage(func, cmd_group, **kwargs):
+        command_name = kwargs.get("name", func.__name__)
+
+        @functools.wraps(func)
+        def wrapper(do_not_track: bool, *args, **kwargs):
+            if do_not_track:
+                os.environs["BENTOCTL_DO_NOT_TRACK"] = str(True)
+                return func(*args, **kwargs)
+            start_time = time.time_ns()
+            if cmd_group.name in cli_events_map:
+                # If cli command is build or operator related, we will add
+                # additoinal properties
+                get_tracking_event = functools.partial(
+                    cli_events_map[cmd_group.name],
+                    cmd_group.name,
+                    command_name,
+                )
+            elif cmd_group.name == "operator":
+
+                def get_tracking_event(return_value):  # pylint: disable=unused-argument
+                    return CliEvent(
+                        cmd_group.name, command_name, operator=kwargs.get("name", None)
+                    )
+
+            else:
+
+                def get_tracking_event(ret):  # pylint: disable=unused-argument
+                    return CliEvent(cmd_group.name, command_name)
+
+            try:
+                return_value = func(*args, **kwargs)
+                event = get_tracking_event(return_value=return_value)
+                duration_in_ms = time.time_ns() - start_time
+                event.duration_in_ms = duration_in_ms / 1e6
+                track(event)
+                return return_value
+            except Exception as e:
+                event = get_tracking_event(None)
+                duration_in_ms = time.time_ns() - start_time
+                event.duration_in_ms = duration_in_ms / 1e6
+                event.error_type = type(e).__name__
+                event.return_code = 2 if isinstance(e, KeyboardInterrupt) else 1
+                track(event)
+                raise
+
+        return wrapper
+
     def command(self, *args, **kwargs):
         def wrapper(func):
             # add common parameters to command
             func = BentoctlCommandGroup.bentoctl_common_params(func)
+            func = BentoctlCommandGroup.bentoctl_track_usage(func, self, **kwargs)
 
             # move common parameters to end of the parameters list
             func.__click_params__ = (

bentoctl/utils/__init__.py

@@ -0,0 +1,21 @@
+import logging
+import os
+
+
+DEBUG_ENV_VAR = "BENTOCTL_DEBUG"
+
+
+def set_debug_mode(is_enabled: bool):
+    if is_enabled or os.environ.get(DEBUG_ENV_VAR):
+        os.environ[DEBUG_ENV_VAR] = str(True)
+        logging.getLogger().setLevel(logging.DEBUG)
+        logging.getLogger("bentoml").setLevel(logging.DEBUG)
+    else:
+        logging.getLogger().setLevel(logging.WARNING)
+        logging.getLogger("bentoml").setLevel(logging.WARNING)
+
+
+def get_debug_mode():
+    if DEBUG_ENV_VAR in os.environ:
+        return os.environ[DEBUG_ENV_VAR].lower() == "true"
+    return False

bentoctl/utils.py → bentoctl/utils/temp_dir.py

@@ -1,25 +1,6 @@
 import os
 import shutil
 import tempfile
-import logging
-
-DEBUG_ENV_VAR = "BENTOCTL_DEBUG"
-
-
-def set_debug_mode(is_enabled: bool):
-    if is_enabled or os.environ.get(DEBUG_ENV_VAR):
-        os.environ[DEBUG_ENV_VAR] = str(True)
-        logging.getLogger().setLevel(logging.DEBUG)
-        logging.getLogger("bentoml").setLevel(logging.DEBUG)
-    else:
-        logging.getLogger().setLevel(logging.WARNING)
-        logging.getLogger("bentoml").setLevel(logging.WARNING)
-
-
-def get_debug_mode():
-    if DEBUG_ENV_VAR in os.environ:
-        return os.environ[DEBUG_ENV_VAR].lower() == "true"
-    return False
 
 
 class TempDirectory(object):

bentoctl/utils/usage_stats.py

@@ -0,0 +1,50 @@
+from bentoctl import __version__
+
+# These are internal apis. We will need to make sure update these when BentoML changes.
+from bentoml._internal.utils.analytics.usage_stats import ( # noqa pylint: disable=unused-import
+    do_not_track,
+    track,
+    BENTOML_DO_NOT_TRACK,
+)
+
+
+class CliEvent:
+    def __init__(
+        self,
+        cmd_group,
+        cmd_name,
+        duration_in_ms=None,
+        error_type=None,
+        return_code=None,
+        operator=None,
+        version=None,
+    ):
+        self.cmd_group = cmd_group
+        self.cmd_name = cmd_name
+        self.duration_in_ms = duration_in_ms
+        self.error_type = error_type
+        self.return_code = return_code
+        self.operator = operator
+        self.version = version
+        self.bentoctl_version = __version__
+        self.event_name = 'bentoctl_cli'
+
+
+def _bentoctl_event(cmd_group, cmd_name, return_value=None):
+    if return_value is not None:
+        deployment_config = return_value
+        version = (
+            deployment_config.bento.tag.version if deployment_config.bento else None
+        )
+
+        return CliEvent(
+            cmd_group,
+            cmd_name,
+            operator=deployment_config.operator_name,
+            version=version,
+        )
+    else:
+        return CliEvent(cmd_group, cmd_name)
+
+
+cli_events_map = {"bentoctl": _bentoctl_event}

tests/unit/cli/test_cli_basic.py

@@ -4,6 +4,7 @@
 
 import pytest
 from click.testing import CliRunner
+from unittest.mock import MagicMock
 
 import bentoctl
 from bentoctl import __version__, deployment_config
@@ -41,9 +42,14 @@ def test_bentoctl_version():
     assert __version__ in result.output
 
 
+bentomock = MagicMock()
+bentomock.tag.version = "mock_version"
+
 @dataclass
 class DeploymentConfigMock:
     repository_name: str = None
+    bento = bentomock
+    operator_name = 'mocked_operator_name'
 
     @classmethod
     def from_file(cls, file):