wehub-resource-sync
167 lines
4.3 KiB
Python
167 lines
4.3 KiB
Python
from abc import ABC, abstractmethod
|
|
from enum import Enum
|
|
from typing import Dict, Generic, Optional, Tuple, TypeVar
|
|
|
|
from ray.util.metrics import Gauge
|
|
|
|
RUN_NAME_TAG_KEY = "ray_train_run_name"
|
|
RUN_ID_TAG_KEY = "ray_train_run_id"
|
|
|
|
T = TypeVar("T")
|
|
E = TypeVar("E", bound=Enum)
|
|
|
|
|
|
class Metric(ABC):
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
default: T,
|
|
description: str,
|
|
base_tags: Dict[str, str],
|
|
):
|
|
"""
|
|
Initialize a new metric.
|
|
|
|
Args:
|
|
name: The name of the metric.
|
|
default: The default value of the metric.
|
|
description: The description of the metric.
|
|
base_tags: The base tags for the metric.
|
|
"""
|
|
self._default = default
|
|
self._base_tags = base_tags
|
|
self._gauge = Gauge(
|
|
name,
|
|
description=description,
|
|
tag_keys=self._get_tag_keys(),
|
|
)
|
|
|
|
@abstractmethod
|
|
def record(self, value: T):
|
|
"""Update the metric value.
|
|
|
|
Args:
|
|
value: The value to update the metric with.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_value(self) -> T:
|
|
"""Get the value of the metric.
|
|
|
|
Returns:
|
|
The value of the metric. If the metric has not been recorded,
|
|
the default value is returned.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def reset(self):
|
|
"""Reset values and clean up resources."""
|
|
pass
|
|
|
|
def _get_tag_keys(self) -> Tuple[str, ...]:
|
|
return tuple(self._base_tags.keys())
|
|
|
|
|
|
class TimeMetric(Metric):
|
|
"""A metric for tracking elapsed time."""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
description: str,
|
|
base_tags: Dict[str, str],
|
|
):
|
|
self._current_value = 0.0
|
|
super().__init__(
|
|
name=name,
|
|
default=0.0,
|
|
description=description,
|
|
base_tags=base_tags,
|
|
)
|
|
|
|
def record(self, value: float):
|
|
"""Update the time metric value by accumulating the time.
|
|
|
|
Args:
|
|
value: The time value to increment the metric by.
|
|
"""
|
|
self._current_value += value
|
|
self._gauge.set(self._current_value, self._base_tags)
|
|
|
|
def get_value(self) -> float:
|
|
return self._current_value
|
|
|
|
def reset(self):
|
|
self._current_value = self._default
|
|
self._gauge.set(self._default, self._base_tags)
|
|
|
|
|
|
class EnumMetric(Metric, Generic[E]):
|
|
"""A metric for tracking enum values."""
|
|
|
|
DEFAULT_VALUE = 0
|
|
RECORDED_VALUE = 1
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
description: str,
|
|
base_tags: Dict[str, str],
|
|
enum_tag_key: str,
|
|
):
|
|
self._enum_tag_key = enum_tag_key
|
|
self._current_value: Optional[E] = None
|
|
super().__init__(
|
|
name=name,
|
|
default=self.DEFAULT_VALUE,
|
|
description=description,
|
|
base_tags=base_tags,
|
|
)
|
|
|
|
def record(self, enum_value: E) -> None:
|
|
"""Record a specific enum value.
|
|
|
|
The metric will be reset to 0 for the previous value and set to 1 for the new value.
|
|
|
|
Args:
|
|
enum_value: The enum value to record for.
|
|
"""
|
|
if enum_value == self._current_value:
|
|
return
|
|
|
|
if self._current_value is not None:
|
|
previous_tags = self._get_tags(self._current_value)
|
|
self._gauge.set(self._default, previous_tags)
|
|
|
|
current_tags = self._get_tags(enum_value)
|
|
self._gauge.set(self.RECORDED_VALUE, current_tags)
|
|
|
|
self._current_value = enum_value
|
|
|
|
def get_value(self, enum_value: E) -> int:
|
|
"""Get the value for a specific enum value.
|
|
|
|
Args:
|
|
enum_value: The enum value to get the value for
|
|
|
|
Returns:
|
|
The value for the enum value
|
|
"""
|
|
return int(enum_value == self._current_value)
|
|
|
|
def reset(self):
|
|
if self._current_value is not None:
|
|
tags = self._get_tags(self._current_value)
|
|
self._gauge.set(self._default, tags)
|
|
self._current_value = None
|
|
|
|
def _get_tag_keys(self) -> Tuple[str, ...]:
|
|
return tuple(self._base_tags.keys()) + (self._enum_tag_key,)
|
|
|
|
def _get_tags(self, enum_value: E) -> Dict[str, str]:
|
|
tags = self._base_tags.copy()
|
|
tags[self._enum_tag_key] = enum_value.name
|
|
return tags
|