wehub-resource-sync
67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
from typing import Dict, Union
|
|
|
|
from ray.train.v2._internal.execution.controller.state import TrainControllerStateType
|
|
from ray.train.v2._internal.metrics.base import (
|
|
RUN_ID_TAG_KEY,
|
|
RUN_NAME_TAG_KEY,
|
|
EnumMetric,
|
|
TimeMetric,
|
|
)
|
|
|
|
|
|
class ControllerMetrics:
|
|
"""Factory for creating controller-specific metrics.
|
|
|
|
This class defines all metrics used to track the state and performance of the
|
|
training controller. Each metric is defined with its name, type, default value,
|
|
description, and required tags.
|
|
"""
|
|
|
|
# ===== Metric Names =====
|
|
CONTROLLER_STATE = "train_controller_state"
|
|
WORKER_GROUP_START_TOTAL_TIME_S = "train_worker_group_start_total_time_s"
|
|
WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S = "train_worker_group_shutdown_total_time_s"
|
|
|
|
# ===== Tag Keys =====
|
|
CONTROLLER_STATE_TAG_KEY = "ray_train_controller_state"
|
|
|
|
@classmethod
|
|
def _create_time_metric(
|
|
cls, name: str, description: str, base_tags: Dict[str, str]
|
|
) -> TimeMetric:
|
|
return TimeMetric(
|
|
name=name,
|
|
description=description,
|
|
base_tags=base_tags,
|
|
)
|
|
|
|
@classmethod
|
|
def _create_controller_state_metric(
|
|
cls, base_tags: Dict[str, str]
|
|
) -> EnumMetric[TrainControllerStateType]:
|
|
return EnumMetric[TrainControllerStateType](
|
|
name=cls.CONTROLLER_STATE,
|
|
description="Current state of the Ray Train controller",
|
|
base_tags=base_tags,
|
|
enum_tag_key=cls.CONTROLLER_STATE_TAG_KEY,
|
|
)
|
|
|
|
@classmethod
|
|
def get_controller_metrics(
|
|
cls, run_name: str, run_id: str
|
|
) -> Dict[str, Union[TimeMetric, EnumMetric[TrainControllerStateType]]]:
|
|
base_tags = {RUN_NAME_TAG_KEY: run_name, RUN_ID_TAG_KEY: run_id}
|
|
return {
|
|
cls.WORKER_GROUP_START_TOTAL_TIME_S: cls._create_time_metric(
|
|
cls.WORKER_GROUP_START_TOTAL_TIME_S,
|
|
"Total time taken to start the worker group",
|
|
base_tags,
|
|
),
|
|
cls.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S: cls._create_time_metric(
|
|
cls.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S,
|
|
"Total time taken to shutdown the worker group",
|
|
base_tags,
|
|
),
|
|
cls.CONTROLLER_STATE: cls._create_controller_state_metric(base_tags),
|
|
}
|