mirror of
https://github.com/bentoml/OpenLLM.git
synced 2026-01-24 15:27:51 -05:00
146 lines
4.5 KiB
Python
146 lines
4.5 KiB
Python
# Copyright 2023 BentoML Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Telemetry related for OpenLLM tracking.
|
|
|
|
Users can disable this with OPENLLM_DO_NOT_TRACK envvar.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import functools
|
|
import os
|
|
import typing as t
|
|
from datetime import datetime
|
|
|
|
import attr
|
|
import bentoml
|
|
from bentoml._internal.utils import analytics as _internal_analytics
|
|
from bentoml._internal.utils.analytics import usage_stats as _internal_usage
|
|
|
|
if t.TYPE_CHECKING:
|
|
import openllm
|
|
import click
|
|
|
|
from ..__about__ import __version__
|
|
|
|
ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
|
|
|
|
# This variable is a proxy that will control BENTOML_DO_NOT_TRACK
|
|
OPENLLM_DO_NOT_TRACK = "OPENLLM_DO_NOT_TRACK"
|
|
|
|
DO_NOT_TRACK = os.environ.get(OPENLLM_DO_NOT_TRACK, str(False)).upper()
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def do_not_track() -> bool:
|
|
return DO_NOT_TRACK in ENV_VARS_TRUE_VALUES
|
|
|
|
|
|
@_internal_usage.silent
|
|
def track(event_properties: _internal_analytics.schemas.EventMeta):
|
|
if do_not_track():
|
|
return
|
|
_internal_analytics.track(event_properties)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def set_bentoml_tracking():
|
|
original_value = os.environ.pop(_internal_analytics.BENTOML_DO_NOT_TRACK, str(False))
|
|
try:
|
|
os.environ[_internal_analytics.BENTOML_DO_NOT_TRACK] = str(do_not_track())
|
|
yield
|
|
finally:
|
|
os.environ[_internal_analytics.BENTOML_DO_NOT_TRACK] = original_value
|
|
|
|
|
|
@attr.define
|
|
class OpenllmCliEvent(_internal_analytics.schemas.EventMeta):
|
|
cmd_group: str
|
|
cmd_name: str
|
|
openllm_version: str = __version__
|
|
|
|
# NOTE: reserved for the do_not_track logics
|
|
duration_in_ms: t.Any = attr.field(default=None)
|
|
error_type: str = attr.field(default=None)
|
|
return_code: int = attr.field(default=None)
|
|
|
|
|
|
if t.TYPE_CHECKING:
|
|
T_con = t.TypeVar("T_con", contravariant=True)
|
|
|
|
class HandlerProtocol(t.Protocol[T_con]):
|
|
@staticmethod
|
|
def __call__(group: click.Group, cmd_name: str, return_value: T_con | None = None) -> OpenllmCliEvent:
|
|
...
|
|
|
|
|
|
@attr.define
|
|
class StartInitEvent(_internal_analytics.schemas.EventMeta):
|
|
model_name: str
|
|
supported_gpu: bool = attr.field(default=False)
|
|
llm_config: t.Dict[str, t.Any] = attr.field(default=None)
|
|
|
|
@staticmethod
|
|
def handler(
|
|
llm_config: openllm.LLMConfig,
|
|
supported_gpu: bool = False,
|
|
) -> StartInitEvent:
|
|
return StartInitEvent(
|
|
model_name=llm_config.__openllm_model_name__,
|
|
llm_config=llm_config.model_dump(),
|
|
supported_gpu=supported_gpu,
|
|
)
|
|
|
|
|
|
def track_start_init(
|
|
llm_config: openllm.LLMConfig,
|
|
supported_gpu: bool = False,
|
|
):
|
|
if do_not_track():
|
|
return
|
|
track(StartInitEvent.handler(llm_config, supported_gpu))
|
|
|
|
|
|
@attr.define
|
|
class BuildEvent(OpenllmCliEvent):
|
|
bento_creation_timestamp: datetime = attr.field(default=None)
|
|
bento_size_in_gb: float = attr.field(default=0)
|
|
model_size_in_gb: float = attr.field(default=0)
|
|
model_type: str = attr.field(default=None)
|
|
model_framework: str = attr.field(default=None)
|
|
|
|
@staticmethod
|
|
def handler(group: click.Group, cmd_name: str, return_value: bentoml.Bento | None = None) -> BuildEvent:
|
|
from bentoml._internal.utils import calc_dir_size
|
|
|
|
assert group.name is not None, "group name should not be None"
|
|
if return_value is not None:
|
|
bento = return_value
|
|
return BuildEvent(
|
|
group.name,
|
|
cmd_name,
|
|
bento_creation_timestamp=bento.info.creation_time,
|
|
bento_size_in_gb=calc_dir_size(bento.path) / 1024**3,
|
|
model_size_in_gb=calc_dir_size(bento.path_of("/models")) / 1024**3,
|
|
model_type=bento.info.labels["_type"],
|
|
model_framework=bento.info.labels["_framework"],
|
|
)
|
|
return BuildEvent(group.name, cmd_name)
|
|
|
|
|
|
cli_events_map: dict[str, dict[str, HandlerProtocol[t.Any]]] = {
|
|
"openllm": {"build": BuildEvent.handler, "bundle": BuildEvent.handler}
|
|
}
|