wehub-resource-sync
294 lines
10 KiB
Python
294 lines
10 KiB
Python
import threading
|
|
import time
|
|
from abc import ABCMeta, abstractmethod
|
|
from collections import deque
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
from ray.rllib.utils.annotations import (
|
|
OverrideToImplementCustomLogic,
|
|
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
|
)
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
class StatsBase(metaclass=ABCMeta):
|
|
"""A base class for Stats.
|
|
|
|
Stats are meant to be used to log values to and then aggregate them in a tree.
|
|
|
|
Therefore, we log to stats in two different ways:
|
|
- On a leaf component, we log values directly by pushing.
|
|
- On a non-leaf component, we only aggregate incoming values.
|
|
|
|
Additionally, we pay special respect to Stats that live at the root of the tree.
|
|
These may have a different behaviour (example: a lifetime sum).
|
|
|
|
Note the tight coupling between StatsBase and MetricsLogger.
|
|
"""
|
|
|
|
# In order to restore from a checkpoint, we need to know the class of the Stats object.
|
|
# This is set in the subclass.
|
|
stats_cls_identifier: str = None
|
|
|
|
def __init__(
|
|
self,
|
|
is_root: bool = False,
|
|
is_leaf: bool = True,
|
|
):
|
|
"""Initializes a StatsBase object.
|
|
|
|
Args:
|
|
is_root: If True, the Stats object is a root stats object.
|
|
is_leaf: If True, the Stats object is a leaf stats object.
|
|
|
|
Note: A stats object can be both root and leaf at the same time.
|
|
Note: A stats object can also be neither root nor leaf ("intermediate" stats that only aggregate from other stats but are not at the root).
|
|
"""
|
|
self.is_root = is_root
|
|
self.is_leaf = is_leaf
|
|
# Used to keep track of start times when using the `with` context manager.
|
|
# This helps us measure times with threads in parallel.
|
|
self._start_times = {}
|
|
# For non-leaf stats (root or intermediate), track the latest merged values
|
|
# This is overwritten on each merge operation
|
|
if not self.is_leaf:
|
|
self.latest_merged: Union[List[Any], Any] = None
|
|
|
|
assert (
|
|
self.stats_cls_identifier is not None
|
|
), "stats_cls_identifier must be set in the subclass"
|
|
|
|
@property
|
|
def has_throughputs(self) -> bool:
|
|
"""Returns True if the Stats object has throughput tracking enabled.
|
|
|
|
Some Stats classes may have throughput tracking enabled, such as SumStats.
|
|
"""
|
|
return False
|
|
|
|
@OverrideToImplementCustomLogic
|
|
def initialize_throughput_reference_time(self, time: float) -> None:
|
|
"""Sets the reference time for this Stats object.
|
|
|
|
This is important because the component that tracks the time
|
|
between reduce cycles is not Stats, but MetricsLogger.
|
|
|
|
Args:
|
|
time: The time to establish as the reference time for this Stats object.
|
|
"""
|
|
if self.has_throughputs:
|
|
raise ValueError(
|
|
"initialize_throughput_reference_time must be overridden for stats objects that have throughputs."
|
|
)
|
|
|
|
@abstractmethod
|
|
def __len__(self) -> int:
|
|
"""Returns the length of the internal values list."""
|
|
...
|
|
|
|
def __float__(self):
|
|
value = self.peek(compile=True)
|
|
if isinstance(value, (list, tuple, deque)):
|
|
raise ValueError(f"Can not convert {self} to float.")
|
|
return float(value)
|
|
|
|
def __int__(self):
|
|
value = self.peek(compile=True)
|
|
if isinstance(value, (list, tuple, deque)):
|
|
raise ValueError(f"Can not convert {self} to int.")
|
|
return int(value)
|
|
|
|
def __eq__(self, other):
|
|
return float(self) == float(other)
|
|
|
|
def __le__(self, other):
|
|
return float(self) <= float(other)
|
|
|
|
def __ge__(self, other):
|
|
return float(self) >= float(other)
|
|
|
|
def __lt__(self, other):
|
|
return float(self) < float(other)
|
|
|
|
def __gt__(self, other):
|
|
return float(self) > float(other)
|
|
|
|
def __add__(self, other):
|
|
return float(self) + float(other)
|
|
|
|
def __sub__(self, other):
|
|
return float(self) - float(other)
|
|
|
|
def __mul__(self, other):
|
|
return float(self) * float(other)
|
|
|
|
def __format__(self, fmt):
|
|
return f"{float(self):{fmt}}"
|
|
|
|
def __enter__(self) -> "StatsBase":
|
|
"""Called when entering a context (with which users can measure a time delta).
|
|
|
|
Returns:
|
|
This stats instance.
|
|
"""
|
|
thread_id = threading.get_ident()
|
|
self._start_times[thread_id] = time.perf_counter()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, tb) -> None:
|
|
"""Called when exiting a context (with which users can measure a time delta).
|
|
|
|
This pushes the time delta since __enter__ to this Stats object.
|
|
"""
|
|
thread_id = threading.get_ident()
|
|
assert self._start_times[thread_id] is not None
|
|
time_delta_s = time.perf_counter() - self._start_times[thread_id]
|
|
self.push(time_delta_s)
|
|
|
|
del self._start_times[thread_id]
|
|
|
|
@classmethod
|
|
def from_state(cls, state: Dict[str, Any]) -> "StatsBase":
|
|
"""Creates a stats object from a state dictionary.
|
|
|
|
Any implementation of this should call this base classe's
|
|
`stats_object.set_state()` to set the state of the stats object.
|
|
|
|
Args:
|
|
state: The state to set after instantiation.
|
|
"""
|
|
init_args = cls._get_init_args(state=state)
|
|
stats = cls(**init_args)
|
|
stats.set_state(state)
|
|
return stats
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
def clone(
|
|
self,
|
|
init_overrides: Optional[Dict[str, Any]] = None,
|
|
) -> "StatsBase":
|
|
"""Returns a new stats object with the same settings as `self`.
|
|
|
|
Args:
|
|
init_overrides: Optional dict of initialization arguments to override. Can be used to change is_root, is_leaf, etc.
|
|
|
|
Returns:
|
|
A new stats object similar to `self` but missing internal values.
|
|
"""
|
|
init_args = self.__class__._get_init_args(stats_object=self)
|
|
if init_overrides:
|
|
init_args.update(init_overrides)
|
|
new_stats = self.__class__(**init_args)
|
|
return new_stats
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
def get_state(self) -> Dict[str, Any]:
|
|
"""Returns the state of the stats object."""
|
|
state = {
|
|
"stats_cls_identifier": self.stats_cls_identifier,
|
|
"is_root": self.is_root,
|
|
"is_leaf": self.is_leaf,
|
|
}
|
|
if not self.is_leaf:
|
|
state["latest_merged"] = self.latest_merged
|
|
return state
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
def set_state(self, state: Dict[str, Any]) -> None:
|
|
"""Sets the state of the stats object.
|
|
|
|
Args:
|
|
state: The state to set on this StatsBase object.
|
|
"""
|
|
|
|
# Handle legacy state that uses old attribute names
|
|
self.is_root = state["is_root"]
|
|
self.is_leaf = state["is_leaf"]
|
|
# Prevent setting a state with a different stats class identifier
|
|
assert self.stats_cls_identifier == state["stats_cls_identifier"]
|
|
if not self.is_leaf:
|
|
# Handle legacy state that doesn't have latest_merged
|
|
self.latest_merged = state["latest_merged"]
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
@staticmethod
|
|
def _get_init_args(stats_object=None, state=None) -> Dict[str, Any]:
|
|
"""Returns the initialization arguments for this Stats object."""
|
|
if state is not None:
|
|
# Handle legacy state that uses old attribute names
|
|
is_root = state["is_root"]
|
|
is_leaf = state["is_leaf"]
|
|
return {
|
|
"is_root": is_root,
|
|
"is_leaf": is_leaf,
|
|
}
|
|
elif stats_object is not None:
|
|
return {
|
|
"is_root": stats_object.is_root,
|
|
"is_leaf": stats_object.is_leaf,
|
|
}
|
|
else:
|
|
raise ValueError("Either stats_object or state must be provided")
|
|
|
|
@abstractmethod
|
|
def __repr__(self) -> str:
|
|
...
|
|
|
|
@abstractmethod
|
|
def merge(self, incoming_stats: List["StatsBase"]) -> None:
|
|
"""Merges StatsBase objects.
|
|
|
|
Args:
|
|
incoming_stats: The list of StatsBase objects to merge.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def push(self, value: Any) -> None:
|
|
"""Pushes a value into this Stats object.
|
|
|
|
Args:
|
|
value: The value to push. Can be of any type.
|
|
GPU tensors are moved to CPU memory.
|
|
"""
|
|
assert (
|
|
self.is_leaf
|
|
), "Cannot push values to non-leaf Stats. Non-leaf Stats can only receive values via merge()."
|
|
|
|
@abstractmethod
|
|
def peek(
|
|
self, compile: bool = True, latest_merged_only: bool = False
|
|
) -> Union[Any, List[Any]]:
|
|
"""Returns the result of reducing the internal values list.
|
|
|
|
Note that this method does NOT alter the internal values list in this process.
|
|
Thus, users can call this method to get an accurate look at the reduced value(s)
|
|
given the current internal values list.
|
|
|
|
Args:
|
|
compile: If True, the result is compiled into a single value if possible.
|
|
latest_merged_only: If True, only considers the latest merged values.
|
|
This parameter only works on aggregation stats objects (is_leaf=False).
|
|
When enabled, peek() will only use the values from the most recent merge operation.
|
|
|
|
Returns:
|
|
The result of reducing the internal values list on CPU memory.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def reduce(self, compile: bool = True) -> Union[Any, "StatsBase"]:
|
|
"""Reduces the internal values.
|
|
|
|
This method should NOT be called directly by users.
|
|
It can be used as a hook to prepare the stats object for sending it to the root metrics logger and starting a new 'reduce cycle'.
|
|
|
|
The reduction logic depends on the implementation of the subclass.
|
|
Meaning that some classes may reduce to a single value, while others do not or don't even contain values.
|
|
|
|
Args:
|
|
compile: If True, the result is compiled into a single value if possible.
|
|
If False, the result is a Stats object similar to itself, but with the internal values reduced.
|
|
Returns:
|
|
The reduced value or a Stats object similar to itself, but with the internal values reduced.
|
|
"""
|