wehub-resource-sync
428 lines
16 KiB
Python
428 lines
16 KiB
Python
"""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__]))
|