chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/metrics/base.py

@@ -0,0 +1,166 @@
+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