chore: import upstream snapshot with attribution

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

@@ -0,0 +1,427 @@
+"""Tests for controller callback behaviour.
+
+Verifies that controller callback hooks are invoked at the correct lifecycle
+points, and that callback exceptions are properly caught, routed through the
+failure decision pipeline, and do not crash the control loop.
+"""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+import ray
+from ray.exceptions import AsyncioActorExit
+from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
+from ray.train.v2._internal.execution.callback import ControllerCallback
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.controller import TrainController
+from ray.train.v2._internal.execution.controller.state import (
+    ErroredState,
+    FinishedState,
+    InitializingState,
+    RunningState,
+    SchedulingState,
+    ShuttingDownState,
+    TrainControllerState,
+)
+from ray.train.v2._internal.execution.failure_handling import FailureDecision
+from ray.train.v2._internal.execution.scaling_policy import ResizeDecision
+from ray.train.v2.api.config import ScalingConfig
+from ray.train.v2.api.exceptions import ControllerError
+from ray.train.v2.tests.util import (
+    DummyObjectRefWrapper,
+    DummyWorkerGroup,
+    MockFailurePolicy,
+    MockScalingPolicy,
+    create_dummy_run_context,
+)
+
+pytestmark = pytest.mark.usefixtures("mock_runtime_context")
+
+
+@pytest.fixture(autouse=True)
+def patch_worker_group(monkeypatch):
+    monkeypatch.setattr(TrainController, "worker_group_cls", DummyWorkerGroup)
+    monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0")
+    yield
+    DummyWorkerGroup.set_poll_failure(None)
+    DummyWorkerGroup.set_start_failure(None)
+
+
+@pytest.fixture(autouse=True)
+def ray_start():
+    ray.init()
+    yield
+    ray.shutdown()
+
+
+async def _create_controller_and_drive_to_running(callbacks=None, num_workers=2):
+    """Create a controller and drive it to RunningState."""
+    scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig())
+    failure_policy = MockFailurePolicy(failure_config=None)
+    controller = TrainController(
+        train_fn_ref=DummyObjectRefWrapper(lambda: None),
+        train_run_context=create_dummy_run_context(),
+        scaling_policy=scaling_policy,
+        failure_policy=failure_policy,
+        callbacks=callbacks or [],
+    )
+    scaling_policy.queue_recovery_decision(
+        ResizeDecision(num_workers=num_workers, resources_per_worker={})
+    )
+    await controller._run_control_loop_iteration()  # Init -> Scheduling
+    await controller._run_control_loop_iteration()  # Scheduling -> Running
+    assert isinstance(controller.get_state(), RunningState)
+    return controller, scaling_policy, failure_policy
+
+
+# ---------------------------------------------------------------------------
+# Happy-path: verify hooks are invoked
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_controller_callback_hooks_are_invoked():
+    """Check that all controller callback hooks are called at the right time."""
+
+    class AssertCallback(ControllerCallback):
+        def __init__(self):
+            self.start_called = False
+            self.latest_state_update = None
+            self.failure_decision_called = False
+            self.resize_decision_called = False
+            self.shutdown_called = False
+
+        def after_controller_start(self, train_run_context: TrainRunContext):
+            self.start_called = True
+
+        def after_controller_state_update(
+            self,
+            previous_state: TrainControllerState,
+            current_state: TrainControllerState,
+        ):
+            self.latest_state_update = (previous_state, current_state)
+
+        def before_controller_execute_failure_decision(
+            self,
+            failure_decision: FailureDecision,
+        ):
+            self.failure_decision_called = True
+
+        def before_controller_execute_resize_decision(
+            self,
+            resize_decision: ResizeDecision,
+        ):
+            self.resize_decision_called = True
+
+        async def before_controller_shutdown(self):
+            self.shutdown_called = True
+
+    callback = AssertCallback()
+
+    scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig())
+    failure_policy = MockFailurePolicy(failure_config=None)
+
+    controller = TrainController(
+        train_fn_ref=DummyObjectRefWrapper(lambda: None),
+        train_run_context=create_dummy_run_context(),
+        scaling_policy=scaling_policy,
+        failure_policy=failure_policy,
+        callbacks=[callback],
+    )
+
+    assert callback.start_called
+
+    scaling_policy.queue_recovery_decision(
+        ResizeDecision(num_workers=2, resources_per_worker={})
+    )
+
+    await controller._run_control_loop_iteration()
+    assert not callback.resize_decision_called
+    assert isinstance(callback.latest_state_update[0], InitializingState)
+    assert isinstance(callback.latest_state_update[1], SchedulingState)
+
+    await controller._run_control_loop_iteration()
+    assert callback.resize_decision_called
+    assert isinstance(callback.latest_state_update[0], SchedulingState)
+    assert isinstance(callback.latest_state_update[1], RunningState)
+
+    controller.get_worker_group().error_worker(1)
+    failure_policy.queue_decision(FailureDecision.RAISE)
+
+    assert not callback.failure_decision_called
+    await controller._run_control_loop_iteration()
+    assert callback.failure_decision_called
+    assert isinstance(callback.latest_state_update[0], RunningState)
+    assert isinstance(callback.latest_state_update[1], ShuttingDownState)
+
+    await controller._run_control_loop_iteration()
+    assert isinstance(callback.latest_state_update[0], ShuttingDownState)
+    assert isinstance(callback.latest_state_update[1], ErroredState)
+    assert callback.shutdown_called
+
+
+# ---------------------------------------------------------------------------
+# Error handling: callback failures during lifecycle hooks
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_callback_error_during_state_update():
+    """A callback error in after_controller_state_update should not crash the
+    control loop; it should transition into shutdown → errored. Callback
+    failures during terminal state transitions must not loop indefinitely."""
+
+    class FailingStateUpdateCallback(ControllerCallback):
+        def after_controller_state_update(
+            self,
+            previous_state: TrainControllerState,
+            current_state: TrainControllerState,
+        ):
+            raise ValueError("Intentional error in state update callback")
+
+    scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig())
+    failure_policy = MockFailurePolicy(failure_config=None)
+
+    controller = TrainController(
+        train_fn_ref=DummyObjectRefWrapper(lambda: None),
+        train_run_context=create_dummy_run_context(),
+        scaling_policy=scaling_policy,
+        failure_policy=failure_policy,
+        callbacks=[FailingStateUpdateCallback()],
+    )
+
+    failure_policy.queue_decision(FailureDecision.RAISE)
+    scaling_policy.queue_recovery_decision(
+        ResizeDecision(num_workers=2, resources_per_worker={})
+    )
+
+    await controller._run_control_loop_iteration()
+    assert isinstance(controller.get_state(), ShuttingDownState)
+    assert isinstance(controller.get_state().next_state, ErroredState)
+    assert isinstance(
+        controller.get_state().next_state.training_failed_error, ControllerError
+    )
+
+    await controller._run_control_loop_iteration()
+    assert isinstance(controller.get_state(), ErroredState)
+    assert isinstance(controller.get_state().training_failed_error, ControllerError)
+
+
+@pytest.mark.asyncio
+async def test_callback_error_during_start():
+    """A callback error in after_controller_start routes through the failure
+    policy and transitions into ShuttingDownState(ErroredState)."""
+
+    class FailingStartCallback(ControllerCallback):
+        def after_controller_start(self, train_run_context):
+            raise ValueError("Intentional error in start callback")
+
+    scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig())
+    failure_policy = MockFailurePolicy(failure_config=None)
+    failure_policy.queue_decision(FailureDecision.RAISE)
+
+    controller = TrainController(
+        train_fn_ref=DummyObjectRefWrapper(lambda: None),
+        train_run_context=create_dummy_run_context(),
+        scaling_policy=scaling_policy,
+        failure_policy=failure_policy,
+        callbacks=[FailingStartCallback()],
+    )
+
+    # _start() failed during __init__.
+    assert isinstance(controller.get_state(), ShuttingDownState)
+    assert isinstance(controller.get_state().next_state, ErroredState)
+    assert isinstance(
+        controller.get_state().next_state.training_failed_error, ControllerError
+    )
+
+    # Complete shutdown → terminal ErroredState.
+    await controller._run_control_loop_iteration()
+    assert isinstance(controller.get_state(), ErroredState)
+    assert isinstance(controller.get_state().training_failed_error, ControllerError)
+
+
+@pytest.mark.asyncio
+async def test_callback_error_during_finish_is_suppressed():
+    """A callback error in after_controller_finish is caught and suppressed.
+    The controller's terminal state must be preserved."""
+
+    class FailingFinishCallback(ControllerCallback):
+        def after_controller_finish(self, result):
+            raise ValueError("Intentional error in finish callback")
+
+    controller, _, _ = await _create_controller_and_drive_to_running(
+        callbacks=[FailingFinishCallback()]
+    )
+
+    # Finish all workers → ShuttingDown(Finished) → Finished.
+    for i in range(2):
+        controller.get_worker_group().finish_worker(i)
+    await controller._run_control_loop_iteration()  # Running -> ShuttingDown(Finished)
+    await controller._run_control_loop_iteration()  # ShuttingDown -> Finished
+    assert isinstance(controller.get_state(), FinishedState)
+
+    # Simulate what run() does after the loop: invoke after_controller_finish.
+    result = controller._build_result()
+    failure_result = controller._run_controller_hook("after_controller_finish", result)
+
+    # Callback error produces a failure result, but state stays FinishedState.
+    assert failure_result is not None
+    assert isinstance(controller.get_state(), FinishedState)
+
+
+@pytest.mark.asyncio
+async def test_callback_error_during_shutdown_hook_on_finished_path():
+    """A before_controller_shutdown callback failure on the finished path
+    transitions to ErroredState with a ControllerError."""
+
+    class FailingShutdownHookCallback(ControllerCallback):
+        async def before_controller_shutdown(self):
+            raise ValueError("Intentional error in shutdown callback")
+
+    controller, _, _ = await _create_controller_and_drive_to_running(
+        callbacks=[FailingShutdownHookCallback()]
+    )
+    for i in range(2):
+        controller.get_worker_group().finish_worker(i)
+    await controller._run_control_loop_iteration()  # Running -> ShuttingDown(Finished)
+    assert isinstance(controller.get_state().next_state, FinishedState)
+
+    await controller._run_control_loop_iteration()
+    assert isinstance(controller.get_state(), ErroredState)
+    assert isinstance(controller.get_state().training_failed_error, ControllerError)
+
+
+@pytest.mark.asyncio
+async def test_callback_error_during_shutdown_hook_on_errored_path():
+    """A before_controller_shutdown callback failure on the errored path
+    preserves the original training error."""
+
+    class FailingShutdownHookCallback(ControllerCallback):
+        async def before_controller_shutdown(self):
+            raise ValueError("Intentional error in shutdown callback")
+
+    controller, _, failure_policy = await _create_controller_and_drive_to_running(
+        callbacks=[FailingShutdownHookCallback()]
+    )
+    controller.get_worker_group().error_worker(0)
+    failure_policy.queue_decision(FailureDecision.RAISE)
+    await controller._run_control_loop_iteration()  # Running -> ShuttingDown(Errored)
+    original_error = controller.get_state().next_state.training_failed_error
+
+    await controller._run_control_loop_iteration()
+    assert isinstance(controller.get_state(), ErroredState)
+    assert controller.get_state().training_failed_error is original_error
+
+
+@pytest.mark.asyncio
+async def test_async_shutdown_callback_is_awaited():
+    """Verify that an async before_controller_shutdown callback is actually
+    awaited (not silently dropped)."""
+
+    class AsyncShutdownCallback(ControllerCallback):
+        def __init__(self):
+            self.shutdown_called = False
+
+        async def before_controller_shutdown(self):
+            self.shutdown_called = True
+
+    callback = AsyncShutdownCallback()
+    controller, _, _ = await _create_controller_and_drive_to_running(
+        callbacks=[callback]
+    )
+    for i in range(2):
+        controller.get_worker_group().finish_worker(i)
+    await controller._run_control_loop_iteration()  # Running -> ShuttingDown(Finished)
+    await controller._run_control_loop_iteration()  # ShuttingDown -> Finished
+    assert isinstance(controller.get_state(), FinishedState)
+    assert callback.shutdown_called
+
+
+@pytest.mark.asyncio
+async def test_async_shutdown_callback_error_is_caught():
+    """An async before_controller_shutdown callback that raises is caught
+    and transitions to ErroredState on the finished path."""
+
+    class FailingAsyncShutdownCallback(ControllerCallback):
+        async def before_controller_shutdown(self):
+            raise ValueError("Async shutdown callback failure")
+
+    controller, _, _ = await _create_controller_and_drive_to_running(
+        callbacks=[FailingAsyncShutdownCallback()]
+    )
+    for i in range(2):
+        controller.get_worker_group().finish_worker(i)
+    await controller._run_control_loop_iteration()  # Running -> ShuttingDown(Finished)
+    assert isinstance(controller.get_state().next_state, FinishedState)
+
+    await controller._run_control_loop_iteration()
+    assert isinstance(controller.get_state(), ErroredState)
+    assert isinstance(controller.get_state().training_failed_error, ControllerError)
+
+
+@pytest.mark.asyncio
+async def test_top_level_safety_net_catches_unhandled_error():
+    """The safety net in _run_control_loop_iteration catches exceptions that
+    escape _step and forces the controller into ErroredState."""
+
+    class CrashingStateUpdateCallback(ControllerCallback):
+        """Raises a non-ControllerError that bypasses _set_state's handling."""
+
+        def __init__(self):
+            self._crash_count = 0
+
+        def after_controller_state_update(
+            self,
+            previous_state: TrainControllerState,
+            current_state: TrainControllerState,
+        ):
+            # Only crash on the first few transitions to test the safety net,
+            # then allow terminal transitions to proceed.
+            if isinstance(current_state, (ErroredState, ShuttingDownState)):
+                return
+            self._crash_count += 1
+            if self._crash_count <= 1:
+                raise RuntimeError("Unexpected crash in state update")
+
+    scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig())
+    failure_policy = MockFailurePolicy(failure_config=None)
+    failure_policy.queue_decision(FailureDecision.RAISE)
+    failure_policy.queue_decision(FailureDecision.RAISE)
+
+    controller = TrainController(
+        train_fn_ref=DummyObjectRefWrapper(lambda: None),
+        train_run_context=create_dummy_run_context(),
+        scaling_policy=scaling_policy,
+        failure_policy=failure_policy,
+        callbacks=[CrashingStateUpdateCallback()],
+    )
+
+    scaling_policy.queue_recovery_decision(
+        ResizeDecision(num_workers=2, resources_per_worker={})
+    )
+
+    # The run() method should not raise — the safety net catches it.
+    await controller.run()
+    assert controller.get_state().is_terminal()
+
+
+@pytest.mark.asyncio
+async def test_asyncio_actor_exit_is_reraised():
+    """AsyncioActorExit raised by _step must propagate through
+    _run_control_loop_iteration without being caught by the safety net."""
+
+    controller, _, _ = await _create_controller_and_drive_to_running()
+
+    with patch.object(
+        controller, "_step", new_callable=AsyncMock, side_effect=AsyncioActorExit()
+    ):
+        with pytest.raises(AsyncioActorExit):
+            await controller._run_control_loop_iteration()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))