[train] add TrainControllerState metrics

[matthewdeng]

Summary

1. Refactors out the metrics handling logic into an internal metrics module.
2. Implements a new metric for tracking TrainController state over time, similar to Ray Core's Task/Actor state metrics.
3. Add a Grafana dashboard panel that shows the states.

Refactoring

All the metrics handling logic is now abstracted away in an internal metrics module.

python/ray/train/v2/_internal/metrics
├── __init__.py
├── base.py
├── controller.py
└── worker.py

As a result, the ControllerMetricsCallback and WorkerMetricsCallback can now be thin layers that map the callback events to calls to MetricsTracker.update.

class Metric:
    def start(self):
        ...
    
    def record(self, tags: Dict[str, str], value: T):
        ...
    
    def get_value(self, tags: Dict[str, str]) -> T:
        ...
    
    def reset(self):
        ...
    
    def shutdown(self):
        ...

TrainControllerState

Added a new metric to track TrainControllerState.

Its key is defined ("ray_train_run_name", "ray_train_controller_state"), and this should be taken into consideration when defining the visualization for this metric.

Example

Ran this locally, with a small modification to print out calls to ray.util.metrics.Gauge.set().

Repro script

import time
from ray.train.torch import TorchTrainer

def train_func():
    time.sleep(10)

trainer = TorchTrainer(train_func)
trainer.fit()

Logs

(TrainController pid=39720) SET train_controller_state 1 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'INITIALIZING'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'INITIALIZING'}
(TrainController pid=39720) SET train_controller_state 1 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'SCHEDULING'}
(TrainController pid=39720) Started training worker group of size 1: 
(TrainController pid=39720) - (ip=127.0.0.1, pid=39728) world_rank=0, local_rank=0, node_rank=0
(RayTrainWorker pid=39728) Setting up process group for: env:// [rank=0, world_size=1]
(TrainController pid=39720) SET train_worker_group_start_total_time_s 1.9262119578197598 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'SCHEDULING'}
(TrainController pid=39720) SET train_controller_state 1 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'RUNNING'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'RUNNING'}
(TrainController pid=39720) SET train_controller_state 1 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36', 'ray_train_controller_state': 'FINISHED'}
(TrainController pid=39720) SET train_worker_group_shutdown_total_time_s 0.004093542229384184 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_worker_group_start_total_time_s 0.0 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_worker_group_shutdown_total_time_s 0.0 {'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_controller_state': 'INITIALIZING', 'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_controller_state': 'SCHEDULING', 'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_controller_state': 'RUNNING', 'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}
(TrainController pid=39720) SET train_controller_state 0 {'ray_train_controller_state': 'FINISHED', 'ray_train_run_name': 'ray_train_run-2025-05-07_17-56-36'}

Dashboard

Added a new dashboard panel that shows the state.

Note that the dashboard has 15 second increments, so states that are shorter than this period may not show up.

Repro script:

import time

from ray.train import RunConfig
from ray.train.backend import Backend, BackendConfig
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer


class SlowBackend(Backend):
    def on_start(self, worker_group, backend_config):
        time.sleep(100)

    def on_shutdown(self, worker_group, backend_config):
        time.sleep(100)


class SlowBackendConfig(BackendConfig):
    @property
    def backend_cls(self):
        return SlowBackend


def train_func():
    time.sleep(100)


run_config = RunConfig(name="slow-backend-run")


trainer = DataParallelTrainer(
    train_func, backend_config=SlowBackendConfig(), run_config=run_config
)
trainer.fit()

[justinvyu]

The "additional tags" argument seems like a confusing implementation detail of the state update metric. Wondering if a pattern like this is simpler:

class EnumMetric(Metric):
    def __init__(self, enum_cls):
        self._lock = nullcontext()

    def create_gauge(self):
        self._gauge = Gauge(...)
    
    def set_lock(self, lock):
        self._lock = lock

    def update(self, new_state: Enum):
        with self.lock:
            self._value = new_state

    def push(self):
        with self._lock:
            for option in Enum.options:
                if self._value == option:
                    self._gauge.set(1, {"ray_train_controller_state": self._value})
                else:
                    self._gauge.set(0, {"ray_train_controller_state": self._value})


controller_state_metric = EnumMetric(ControllerState, ...)
tracker = MetricsTracker([controller_state_metric])  # This creates and sets the lock on all metrics passed in
controller_state_metric.update(curr_controller_state)

[justinvyu]

this description doesn't match with what's happening below.

Implicitly adding values for cumulative metrics when you call "update" is a bit misleading.

What about having keeping an "accumulation_fn" in the Metric dataclass (ex: lambda accumulated_val, curr: curr and lambda accumulated_val, curr: accumulated_val + curr). Then, use the metric's accumulation function to update the underlying value.

[matthewdeng]

Ooh this is a great callout. I was overindexing a bit on this particular new metric that I added.

[justinvyu]

what's the purpose of these tag keys?

[matthewdeng]

This is meant to validate the tag keys that need to be specified when logging this metric. There's actually validation that happens at the lower level if the tags aren't passed in, but I can add a quick validation at the update/record layer as well so it fast-fails!

[justinvyu]

Is gauge.set() a very cheap operation? Was there a point to doing the background thread loop in the first place?

[justinvyu]

should we set it to default?

[justinvyu]

or maybe just remove "Default" from the base class since it's not used there

[matthewdeng]

yes good call, I will remove since Metric is now abstract!

[matthewdeng]

Oh actually I will keep it and convert to use default because I think it's nice to guarantee this in the logic of get_value

[justinvyu]

🆒