-
Notifications
You must be signed in to change notification settings - Fork 757
/
metrics.py
55 lines (41 loc) · 1.55 KB
/
metrics.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from __future__ import annotations
import typing as t
import logging
from typing import TYPE_CHECKING
from simple_di import inject
from simple_di import Provide
if TYPE_CHECKING:
from ._internal.server.metrics.prometheus import PrometheusClient
logger = logging.getLogger(__name__)
from ._internal.configuration.containers import BentoMLContainer
def __dir__() -> list[str]:
metrics_client = BentoMLContainer.metrics_client.get()
return dir(metrics_client.prometheus_client)
class _LazyMetric:
def __init__(self, attr: t.Any):
self._attr = attr
self._proxy = None
self._initialized = False
self._args: tuple[t.Any, ...] = ()
self._kwargs: dict[str, t.Any] = {}
def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
# This is where we lazy load the proxy object.
self._args = args
self._kwargs = kwargs
return self
def __getattr__(self, item: t.Any) -> t.Any:
if item in ("_attr", "_proxy", "_initialized", "_args", "_kwargs"):
raise AttributeError(f"Attribute {item} is private to {self}.")
if self._proxy is None:
self._load_proxy()
assert self._initialized
return getattr(self._proxy, item)
@inject
def _load_proxy(
self,
metrics_client: PrometheusClient = Provide[BentoMLContainer.metrics_client],
) -> None:
self._proxy = getattr(metrics_client, self._attr)(*self._args, **self._kwargs)
self._initialized = True
def __getattr__(item: t.Any):
return _LazyMetric(item)