diff --git a/openllm_next/__main__.py b/openllm_next/__main__.py index 6c870792..513f6d2a 100644 --- a/openllm_next/__main__.py +++ b/openllm_next/__main__.py @@ -12,17 +12,11 @@ from openllm_next.accelerator_spec import ( can_run, get_local_machine_spec, ) +from openllm_next.analytic import DO_NOT_TRACK, OpenLLMTyper from openllm_next.clean import app as clean_app from openllm_next.cloud import deploy as cloud_deploy from openllm_next.cloud import ensure_cloud_context, get_cloud_machine_spec -from openllm_next.common import ( - CHECKED, - DO_NOT_TRACK, - INTERACTIVE, - VERBOSE_LEVEL, - OpenLLMTyper, - output, -) +from openllm_next.common import CHECKED, INTERACTIVE, VERBOSE_LEVEL, output from openllm_next.local import run as local_run from openllm_next.local import serve as local_serve from openllm_next.model import app as model_app diff --git a/openllm_next/analytic.py b/openllm_next/analytic.py new file mode 100644 index 00000000..0beb4482 --- /dev/null +++ b/openllm_next/analytic.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import functools +import os +import re +import time +import typing +from abc import ABC + +import attr +import click +import typer +import typer.core + +DO_NOT_TRACK = "BENTOML_DO_NOT_TRACK" + + +class EventMeta(ABC): + @property + def event_name(self): + # camel case to snake case + event_name = re.sub(r"(? typing.Iterable[str]: + return list(self.commands) + + +class OpenLLMTyper(typer.Typer): + def __init__(self, *args: typing.Any, **kwargs: typing.Any): + no_args_is_help = kwargs.pop("no_args_is_help", True) + context_settings = kwargs.pop("context_settings", {}) + if "help_option_names" not in context_settings: + context_settings["help_option_names"] = ("-h", "--help") + if "max_content_width" not in context_settings: + context_settings["max_content_width"] = int( + os.environ.get("COLUMNS", str(120)) + ) + klass = kwargs.pop("cls", OrderedCommands) + + super().__init__( + *args, + cls=klass, + no_args_is_help=no_args_is_help, + context_settings=context_settings, + **kwargs, + ) + + def command(self, *args: typing.Any, **kwargs: typing.Any): + def decorator(f): + @functools.wraps(f) + @click.pass_context + def wrapped(ctx: click.Context, *args, **kwargs): + from bentoml._internal.utils.analytics import track + + do_not_track = ( + os.environ.get(DO_NOT_TRACK, str(False)).lower() == "true" + ) + + # so we know that the root program is openllm + command_name = ctx.info_name + if ctx.parent.parent is not None: + # openllm model list + command_group = ctx.parent.info_name + elif ctx.parent.info_name == ctx.find_root().info_name: + # openllm run + command_group = "openllm" + + if do_not_track: + return f(*args, **kwargs) + start_time = time.time_ns() + try: + return_value = f(*args, **kwargs) + duration_in_ns = time.time_ns() - start_time + track( + OpenllmCliEvent( + cmd_group=command_group, + cmd_name=command_name, + duration_in_ms=duration_in_ns / 1e6, + ) + ) + return return_value + except BaseException as e: + duration_in_ns = time.time_ns() - start_time + track( + OpenllmCliEvent( + cmd_group=command_group, + cmd_name=command_name, + duration_in_ms=duration_in_ns / 1e6, + error_type=type(e).__name__, + return_code=2 if isinstance(e, KeyboardInterrupt) else 1, + ) + ) + raise + + return typer.Typer.command(self, *args, **kwargs)(wrapped) + + return decorator diff --git a/openllm_next/clean.py b/openllm_next/clean.py index 6a77cd1c..b14f324a 100644 --- a/openllm_next/clean.py +++ b/openllm_next/clean.py @@ -3,13 +3,8 @@ import shutil import questionary -from openllm_next.common import ( - REPO_DIR, - VENV_DIR, - VERBOSE_LEVEL, - OpenLLMTyper, - output, -) +from openllm_next.analytic import OpenLLMTyper +from openllm_next.common import REPO_DIR, VENV_DIR, VERBOSE_LEVEL, output app = OpenLLMTyper(help="clean up and release disk space used by OpenLLM") diff --git a/openllm_next/cloud.py b/openllm_next/cloud.py index 39886aad..a74383b0 100644 --- a/openllm_next/cloud.py +++ b/openllm_next/cloud.py @@ -8,13 +8,13 @@ import typing import typer from openllm_next.accelerator_spec import ACCELERATOR_SPECS +from openllm_next.analytic import OpenLLMTyper from openllm_next.common import ( INTERACTIVE, BentoInfo, DeploymentTarget, output, run_command, - OpenLLMTyper, ) app = OpenLLMTyper() diff --git a/openllm_next/common.py b/openllm_next/common.py index 9f7bf5fd..6d832bdf 100644 --- a/openllm_next/common.py +++ b/openllm_next/common.py @@ -11,18 +11,12 @@ import signal import subprocess import sys import sysconfig -import time import typing from contextlib import asynccontextmanager, contextmanager from types import SimpleNamespace -import attr -import click import typer import typer.core -from bentoml._internal.utils.analytics import \ - BENTOML_DO_NOT_TRACK as DO_NOT_TRACK -from bentoml._internal.utils.analytics import CliEvent ERROR_STYLE = "red" SUCCESS_STYLE = "green" @@ -44,88 +38,6 @@ CHECKED = "☆" T = typing.TypeVar("T") -@attr.define -class OpenllmCliEvent(CliEvent): - pass - - -class OrderedCommands(typer.core.TyperGroup): - def list_commands(self, _: click.Context) -> typing.Iterable[str]: - return list(self.commands) - - -class OpenLLMTyper(typer.Typer): - def __init__(self, *args: typing.Any, **kwargs: typing.Any): - no_args_is_help = kwargs.pop("no_args_is_help", True) - context_settings = kwargs.pop("context_settings", {}) - if "help_option_names" not in context_settings: - context_settings["help_option_names"] = ("-h", "--help") - if "max_content_width" not in context_settings: - context_settings["max_content_width"] = int( - os.environ.get("COLUMNS", str(120)) - ) - klass = kwargs.pop("cls", OrderedCommands) - - super().__init__( - *args, - cls=klass, - no_args_is_help=no_args_is_help, - context_settings=context_settings, - **kwargs, - ) - - def command(self, *args: typing.Any, **kwargs: typing.Any): - def decorator(f): - @functools.wraps(f) - @click.pass_context - def wrapped(ctx: click.Context, *args, **kwargs): - from bentoml._internal.utils.analytics import track - - do_not_track = ( - os.environ.get(DO_NOT_TRACK, str(False)).lower() == "true" - ) - - # so we know that the root program is openllm - command_name = ctx.info_name - if ctx.parent.parent is not None: - # openllm model list - command_group = ctx.parent.info_name - elif ctx.parent.info_name == ctx.find_root().info_name: - # openllm run - command_group = "openllm" - - if do_not_track: - return f(*args, **kwargs) - start_time = time.time_ns() - try: - return_value = f(*args, **kwargs) - duration_in_ns = time.time_ns() - start_time - track( - OpenllmCliEvent( - cmd_group=command_group, - cmd_name=command_name, - duration_in_ms=duration_in_ns / 1e6, - ) - ) - return return_value - except BaseException as e: - duration_in_ns = time.time_ns() - start_time - track( - OpenllmCliEvent( - cmd_group=command_group, - cmd_name=command_name, - duration_in_ms=duration_in_ns / 1e6, - error_type=type(e).__name__, - return_code=2 if isinstance(e, KeyboardInterrupt) else 1, - ) - ) - raise - - return typer.Typer.command(self, *args, **kwargs)(wrapped) - - return decorator - - class ContextVar(typing.Generic[T]): def __init__(self, default: T): self._stack: list[T] = [] diff --git a/openllm_next/model.py b/openllm_next/model.py index 5f11e9a6..26d8c42c 100644 --- a/openllm_next/model.py +++ b/openllm_next/model.py @@ -5,13 +5,13 @@ import tabulate import typer from openllm_next.accelerator_spec import DeploymentTarget, can_run +from openllm_next.analytic import OpenLLMTyper from openllm_next.common import ( + FORCE, VERBOSE_LEVEL, BentoInfo, load_config, output, - FORCE, - OpenLLMTyper, ) from openllm_next.repo import ensure_repo_updated, parse_repo_url diff --git a/openllm_next/repo.py b/openllm_next/repo.py index 5137d898..c5e696ae 100644 --- a/openllm_next/repo.py +++ b/openllm_next/repo.py @@ -4,7 +4,9 @@ import shutil import pyaml import questionary +import typer +from openllm_next.analytic import OpenLLMTyper from openllm_next.common import ( INTERACTIVE, REPO_DIR, @@ -13,7 +15,6 @@ from openllm_next.common import ( load_config, output, save_config, - OpenLLMTyper, ) UPDATE_INTERVAL = datetime.timedelta(days=3)