chore: import upstream snapshot with attribution

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

@@ -0,0 +1,124 @@
+from contextlib import contextmanager
+from typing import Dict, Optional
+
+import ray
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    TrainContextCallback,
+    WorkerCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import TrainRunContext, get_train_context
+from ray.train.v2._internal.execution.controller.state import (
+    TrainControllerState,
+    TrainControllerStateType,
+)
+from ray.train.v2._internal.metrics.base import Metric
+from ray.train.v2._internal.metrics.controller import ControllerMetrics
+from ray.train.v2._internal.metrics.worker import WorkerMetrics
+from ray.train.v2._internal.util import time_monotonic
+
+
+class ControllerMetricsCallback(ControllerCallback, WorkerGroupCallback):
+    """Callback that records controller-specific metrics."""
+
+    def after_controller_start(self, train_run_context: TrainRunContext):
+        """Initialize metrics after controller starts."""
+        self._run_name = train_run_context.get_run_config().name
+        self._run_id = train_run_context.run_id
+        self._metrics: Dict[str, Metric] = ControllerMetrics.get_controller_metrics(
+            self._run_name, self._run_id
+        )
+        # Record initial state
+        self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
+            TrainControllerStateType.INITIALIZING
+        )
+
+    async def before_controller_shutdown(self):
+        """Shutdown metrics before controller shuts down."""
+        for metric in self._metrics.values():
+            metric.reset()
+
+    def after_controller_state_update(
+        self,
+        previous_state: TrainControllerState,
+        current_state: TrainControllerState,
+    ):
+        """Record state transitions after controller state updates."""
+        self._metrics[ControllerMetrics.CONTROLLER_STATE].record(
+            current_state._state_type
+        )
+
+    @contextmanager
+    def on_worker_group_start(self):
+        """Measure time taken to start worker group."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].record(
+            elapsed_time_s
+        )
+
+    @contextmanager
+    def on_worker_group_shutdown(self):
+        """Measure time taken to shutdown worker group."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S].record(
+            elapsed_time_s
+        )
+
+
+class WorkerMetricsCallback(WorkerCallback, TrainContextCallback):
+    """Callback that records worker-specific metrics."""
+
+    def __init__(self, train_run_context: TrainRunContext):
+        self._run_name = train_run_context.get_run_config().name
+        self._run_id = train_run_context.run_id
+        self._metrics: Optional[Dict[str, Metric]] = None
+
+    def after_init_train_context(self):
+        """Initialize metrics after train context is initialized."""
+        train_context = get_train_context()
+        core_context = ray.runtime_context.get_runtime_context()
+        world_rank = train_context.get_world_rank()
+        worker_actor_id = core_context.get_actor_id()
+        self._metrics = WorkerMetrics.get_worker_metrics(
+            self._run_name, self._run_id, world_rank, worker_actor_id
+        )
+
+    def before_worker_shutdown(self):
+        """Shutdown metrics before shutdown."""
+        if self._metrics:
+            for metric in self._metrics.values():
+                metric.reset()
+
+    @contextmanager
+    def on_report(self):
+        """
+        Context manager to measure the time taken to report a checkpoint to the storage.
+        """
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].record(elapsed_time_s)
+
+    @contextmanager
+    def on_checkpoint_sync(self):
+        """Measure time spent in the cross-rank barrier that synchronizes the
+        checkpoint directory name across all workers."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[WorkerMetrics.CHECKPOINT_SYNC_TOTAL_TIME_S].record(elapsed_time_s)
+
+    @contextmanager
+    def on_checkpoint_transfer(self):
+        """Measure time spent transferring checkpoint files to storage."""
+        start_time_s = time_monotonic()
+        yield
+        elapsed_time_s = time_monotonic() - start_time_s
+        self._metrics[WorkerMetrics.CHECKPOINT_TRANSFER_TOTAL_TIME_S].record(
+            elapsed_time_s
+        )