chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_metrics.py

@@ -0,0 +1,363 @@
+import asyncio
+import enum
+from unittest.mock import MagicMock
+
+import pytest
+
+import ray
+from ray.train.v2._internal.callbacks.metrics import (
+    ControllerMetricsCallback,
+    WorkerMetricsCallback,
+)
+from ray.train.v2._internal.execution.controller.state import (
+    TrainControllerState,
+    TrainControllerStateType,
+)
+from ray.train.v2._internal.metrics.base import EnumMetric, TimeMetric
+from ray.train.v2._internal.metrics.controller import ControllerMetrics
+from ray.train.v2._internal.metrics.worker import WorkerMetrics
+from ray.train.v2.api.config import RunConfig
+from ray.train.v2.tests.util import create_dummy_run_context
+
+
+class MockGauge:
+    """Mock class for ray.util.metrics.Gauge."""
+
+    def __init__(self, name: str, description: str, tag_keys: tuple = ()):
+        self._values: dict[set[str], float] = {}
+
+    def set(self, value: float, tags: dict):
+        self._values[frozenset(tags.items())] = value
+
+
+@pytest.fixture
+def mock_gauge(monkeypatch):
+    """Fixture that replaces ray.util.metrics.Gauge with MockGauge."""
+    monkeypatch.setattr(ray.train.v2._internal.metrics.base, "Gauge", MockGauge)
+    return MockGauge
+
+
+def mock_time_monotonic(monkeypatch, time_values: list[float]):
+    time_index = 0
+
+    def mock_time():
+        nonlocal time_index
+        value = time_values[time_index]
+        time_index = time_index + 1
+        return value
+
+    monkeypatch.setattr(
+        ray.train.v2._internal.callbacks.metrics, "time_monotonic", mock_time
+    )
+
+
+def mock_start_end_time(monkeypatch, time_values: list[tuple[float, float]]):
+    """Mock the time_monotonic function to return the start and end times.
+
+    This assumes that time_monotonic is called in the order of the start and end times.
+    """
+    all_times = []
+    for start, end in time_values:
+        all_times.append(start)
+        all_times.append(end)
+    mock_time_monotonic(monkeypatch, all_times)
+
+
+def test_time_metric(monkeypatch, mock_gauge):
+    base_tags = {"run_name": "test_run"}
+    metric = TimeMetric(
+        name="test_time",
+        description="Test time metric",
+        base_tags=base_tags,
+    )
+
+    # Test recording values
+    metric.record(1.0)
+    assert metric.get_value() == 1.0
+
+    # Test updating metric
+    metric.record(2.0)
+    assert metric.get_value() == 3.0
+
+    # Test reset
+    metric.reset()
+    assert metric.get_value() == 0.0
+
+
+def test_enum_metric(monkeypatch, mock_gauge):
+    class TestEnum(enum.Enum):
+        A = "A"
+        B = "B"
+        C = "C"
+
+    base_tags = {"run_name": "test_run"}
+    metric = EnumMetric[TestEnum](
+        name="test_enum",
+        description="Test enum metric",
+        base_tags=base_tags,
+        enum_tag_key="state",
+    )
+
+    # Test recording values
+    metric.record(TestEnum.A)
+    assert metric.get_value(TestEnum.A) == 1
+    assert metric.get_value(TestEnum.B) == 0
+    assert metric.get_value(TestEnum.C) == 0
+
+    metric.record(TestEnum.B)
+    assert metric.get_value(TestEnum.A) == 0
+    assert metric.get_value(TestEnum.B) == 1
+    assert metric.get_value(TestEnum.C) == 0
+
+    metric.record(TestEnum.C)
+    assert metric.get_value(TestEnum.A) == 0
+    assert metric.get_value(TestEnum.B) == 0
+    assert metric.get_value(TestEnum.C) == 1
+
+    # Test reset
+    metric.reset()
+    assert metric.get_value(TestEnum.A) == 0
+    assert metric.get_value(TestEnum.B) == 0
+    assert metric.get_value(TestEnum.C) == 0
+
+
+def test_worker_metrics_callback(monkeypatch, mock_gauge):
+    t1 = 0.0
+    t2 = 1.0
+    t3 = 10.0
+    t4 = 12.0
+    mock_start_end_time(monkeypatch, [(t1, t2), (t3, t4)])
+
+    mock_train_context = MagicMock()
+    mock_train_context.get_world_rank.return_value = 1
+    mock_train_context.train_run_context = create_dummy_run_context()
+    monkeypatch.setattr(
+        ray.train.v2._internal.callbacks.metrics,
+        "get_train_context",
+        lambda: mock_train_context,
+    )
+
+    callback = WorkerMetricsCallback(train_run_context=create_dummy_run_context())
+    callback.after_init_train_context()
+
+    # Check if the gauges is updated with the correct metrics
+    with callback.on_report():
+        pass
+    assert (
+        callback._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].get_value()
+        == t2 - t1
+    )
+
+    # Check if the gauges is updated with the correct metrics
+    with callback.on_report():
+        pass
+    assert callback._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].get_value() == (
+        t2 - t1
+    ) + (t4 - t3)
+
+    callback.before_worker_shutdown()
+    assert (
+        callback._metrics[WorkerMetrics.REPORT_TOTAL_BLOCKED_TIME_S].get_value() == 0.0
+    )
+
+
+@pytest.mark.parametrize(
+    "context_manager_name,metric_name",
+    [
+        ("on_checkpoint_sync", WorkerMetrics.CHECKPOINT_SYNC_TOTAL_TIME_S),
+        ("on_checkpoint_transfer", WorkerMetrics.CHECKPOINT_TRANSFER_TOTAL_TIME_S),
+    ],
+)
+def test_worker_checkpoint_metrics_callback(
+    monkeypatch, mock_gauge, context_manager_name, metric_name
+):
+    t1 = 0.0
+    t2 = 1.0
+    t3 = 10.0
+    t4 = 12.0
+    mock_start_end_time(monkeypatch, [(t1, t2), (t3, t4)])
+
+    mock_train_context = MagicMock()
+    mock_train_context.get_world_rank.return_value = 1
+    mock_train_context.train_run_context = create_dummy_run_context()
+    monkeypatch.setattr(
+        ray.train.v2._internal.callbacks.metrics,
+        "get_train_context",
+        lambda: mock_train_context,
+    )
+
+    callback = WorkerMetricsCallback(train_run_context=create_dummy_run_context())
+    callback.after_init_train_context()
+
+    context_manager = getattr(callback, context_manager_name)
+
+    with context_manager():
+        pass
+    assert callback._metrics[metric_name].get_value() == t2 - t1
+
+    with context_manager():
+        pass
+    assert callback._metrics[metric_name].get_value() == (t2 - t1) + (t4 - t3)
+
+    callback.before_worker_shutdown()
+    assert callback._metrics[metric_name].get_value() == 0.0
+
+
+def test_worker_metrics_callback_shutdown_without_init(mock_gauge):
+    """before_worker_shutdown should not crash when _metrics is None."""
+    callback = WorkerMetricsCallback(train_run_context=create_dummy_run_context())
+    # _metrics is None — after_init_train_context was never called
+    callback.before_worker_shutdown()  # should not raise
+
+
+def test_controller_metrics_callback(monkeypatch, mock_gauge):
+    t1 = 0.0
+    t2 = 1.0
+    t3 = 10.0
+    t4 = 12.0
+    mock_start_end_time(monkeypatch, [(t1, t2), (t3, t4)])
+
+    mock_train_context = MagicMock()
+    mock_train_context.get_run_config.return_value = RunConfig(name="test_run_name")
+    monkeypatch.setattr(
+        ray.train.v2._internal.execution.context,
+        "get_train_context",
+        lambda: mock_train_context,
+    )
+
+    callback = ControllerMetricsCallback()
+    callback.after_controller_start(train_run_context=create_dummy_run_context())
+
+    # Check if the gauges is updated with the correct metrics
+    with callback.on_worker_group_start():
+        pass
+    assert (
+        callback._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].get_value()
+        == t2 - t1
+    )
+    assert (
+        callback._metrics[
+            ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S
+        ].get_value()
+        == 0.0
+    )
+
+    # Check if the gauges is updated with the correct metrics
+    with callback.on_worker_group_shutdown():
+        pass
+    assert (
+        callback._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].get_value()
+        == t2 - t1
+    )
+    assert (
+        callback._metrics[
+            ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S
+        ].get_value()
+        == t4 - t3
+    )
+
+    asyncio.run(callback.before_controller_shutdown())
+    assert (
+        callback._metrics[ControllerMetrics.WORKER_GROUP_START_TOTAL_TIME_S].get_value()
+        == 0.0
+    )
+    assert (
+        callback._metrics[
+            ControllerMetrics.WORKER_GROUP_SHUTDOWN_TOTAL_TIME_S
+        ].get_value()
+        == 0.0
+    )
+
+
+def test_controller_state_metrics(monkeypatch, mock_gauge):
+    """Test controller state transition metrics."""
+    mock_train_context = MagicMock()
+    mock_train_context.get_run_config.return_value = RunConfig(name="test_run_name")
+    monkeypatch.setattr(
+        ray.train.v2._internal.execution.context,
+        "get_train_context",
+        lambda: mock_train_context,
+    )
+
+    callback = ControllerMetricsCallback()
+    callback.after_controller_start(train_run_context=create_dummy_run_context())
+
+    # Test initial state
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.INITIALIZING
+        )
+        == 1
+    )
+
+    # Test state transition
+
+    previous_state = TrainControllerState(TrainControllerStateType.INITIALIZING)
+    current_state = TrainControllerState(TrainControllerStateType.RUNNING)
+    callback.after_controller_state_update(previous_state, current_state)
+
+    # Verify state counts
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.INITIALIZING
+        )
+        == 0
+    )
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.RUNNING
+        )
+        == 1
+    )
+
+    # Test another state transition
+    previous_state = TrainControllerState(TrainControllerStateType.RUNNING)
+    current_state = TrainControllerState(TrainControllerStateType.FINISHED)
+    callback.after_controller_state_update(previous_state, current_state)
+
+    # Verify updated state counts
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.INITIALIZING
+        )
+        == 0
+    )
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.RUNNING
+        )
+        == 0
+    )
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.FINISHED
+        )
+        == 1
+    )
+
+    asyncio.run(callback.before_controller_shutdown())
+
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.INITIALIZING
+        )
+        == 0
+    )
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.RUNNING
+        )
+        == 0
+    )
+    assert (
+        callback._metrics[ControllerMetrics.CONTROLLER_STATE].get_value(
+            TrainControllerStateType.FINISHED
+        )
+        == 0
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))