from unittest.mock import create_autospec import pytest import ray from ray.train.backend import Backend, BackendConfig from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR from ray.train.v2._internal.exceptions import ( WorkerGroupStartupFailedError, WorkerGroupStartupTimeoutError, ) 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 ( AbortedState, ErroredState, FinishedState, InitializingState, ReschedulingState, ResizingState, RestartingState, RunningState, SchedulingState, ShuttingDownState, TrainControllerState, ) from ray.train.v2._internal.execution.failure_handling import FailureDecision from ray.train.v2._internal.execution.scaling_policy import ( NoopDecision, ResizeDecision, ) from ray.train.v2._internal.execution.worker_group import ( WorkerGroupPollStatus, WorkerStatus, ) 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) # Make polling interval 0 to speed up tests 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() @pytest.mark.asyncio async def test_resize(): scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) train_run_context = create_dummy_run_context() controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=MockFailurePolicy(failure_config=None), ) decisions = [ NoopDecision(), ResizeDecision(num_workers=2, resources_per_worker={}), NoopDecision(), NoopDecision(), ResizeDecision(num_workers=10, resources_per_worker={}), NoopDecision(), ResizeDecision(num_workers=10, resources_per_worker={}), ResizeDecision(num_workers=20, resources_per_worker={}), NoopDecision(), ResizeDecision(num_workers=5, resources_per_worker={}), ] assert isinstance(controller.get_state(), InitializingState) assert controller.get_worker_group() is None # Noop decision should be ignored scaling_policy.queue_recovery_decision(NoopDecision()) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), InitializingState) assert controller.get_worker_group() is None # Start with 1 worker scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=1, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) assert controller.get_worker_group() is None await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group = controller.get_worker_group() assert worker_group is not None assert worker_group.has_started() num_workers = len(worker_group.get_workers()) assert num_workers == 1 for decision in decisions: prev_num_workers = num_workers prev_worker_group = worker_group scaling_policy.queue_monitor_decision(decision) if isinstance(decision, NoopDecision): await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group = controller.get_worker_group() assert worker_group is not None assert worker_group is prev_worker_group assert worker_group.has_started() num_workers = len(worker_group.get_workers()) assert num_workers == prev_num_workers else: # TODO: refactor common "run and check" sequences like this. await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), ResizingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group = controller.get_worker_group() assert worker_group is not None assert worker_group is not prev_worker_group assert worker_group.has_started() num_workers = len(worker_group.get_workers()) assert num_workers == decision.num_workers @pytest.mark.asyncio async def test_failure_handling(): scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) failure_policy = MockFailurePolicy(failure_config=None) train_run_context = create_dummy_run_context() controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=failure_policy, ) assert isinstance(controller.get_state(), InitializingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=2, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group_before_failure = controller.get_worker_group() controller.get_worker_group().error_worker(1) failure_policy.queue_decision(FailureDecision.RETRY) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RestartingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=4, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) # After failure recovery, worker group should be a new instance (full restart). assert controller.get_worker_group() is not worker_group_before_failure DummyWorkerGroup.set_poll_failure(RuntimeError("Simulated poll failure")) failure_policy.queue_decision(FailureDecision.RAISE) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), ShuttingDownState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), ErroredState) @pytest.mark.parametrize( "error_type", [WorkerGroupStartupFailedError, WorkerGroupStartupTimeoutError(2)] ) @pytest.mark.asyncio async def test_worker_group_start_failure(error_type): """Check that controller can gracefully handle worker group start failures.""" scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) failure_policy = MockFailurePolicy(failure_config=None) train_run_context = create_dummy_run_context() controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=failure_policy, ) DummyWorkerGroup.set_start_failure(error_type) assert isinstance(controller.get_state(), InitializingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=2, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) # Worker group will fail to start, but controller should not raise # and should go into RESCHEDULING state. failure_policy.queue_decision(FailureDecision.RETRY) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), ReschedulingState) # Let the worker group start successfully the 2nd time. DummyWorkerGroup.set_start_failure(None) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=2, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) @pytest.mark.asyncio async def test_poll_frequency(monkeypatch): monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "1") async def sleep_mock(t): sleep_calls.append(t) sleep_calls = [] monkeypatch.setattr("asyncio.sleep", sleep_mock) train_run_context = create_dummy_run_context() scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=None, ) # Mock worker group to avoid actual polling controller._worker_group = create_autospec(DummyWorkerGroup, instance=True) controller._worker_group.poll_status.return_value = WorkerGroupPollStatus( worker_statuses={} ) num_polls = 5 for _ in range(num_polls): await controller._poll_workers() # No sleep calls for the first poll assert len(sleep_calls) == num_polls - 1 @pytest.mark.asyncio async def test_controller_callback(monkeypatch): """Check that all controller callback hooks are called.""" 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 self.before_abort_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 def before_controller_abort(self): self.before_abort_called = True callback = AssertCallback() scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) failure_policy = MockFailurePolicy(failure_config=None) train_run_context = create_dummy_run_context() controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=failure_policy, callbacks=[callback], ) assert callback.start_called mock_exit_actor = create_autospec(ray.actor.exit_actor) monkeypatch.setattr("ray.actor.exit_actor", mock_exit_actor) await controller.abort() assert callback.before_abort_called assert isinstance(callback.latest_state_update[1], AbortedState) # Reset the state to InitializingState to test the control loop controller._set_state(InitializingState()) 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 @pytest.mark.asyncio async def test_controller_abort(monkeypatch): mock_exit_actor = create_autospec(ray.actor.exit_actor) monkeypatch.setattr("ray.actor.exit_actor", mock_exit_actor) scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) failure_policy = MockFailurePolicy(failure_config=None) train_run_context = create_dummy_run_context() controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=failure_policy, ) await controller.abort() assert mock_exit_actor.call_count == 1 assert isinstance(controller.get_state(), AbortedState) @pytest.mark.asyncio async def test_shutdown_failure_on_finished_path(): """Shutdown failure on the finished path transitions to ErroredState.""" def failing_shutdown(): raise RuntimeError("Simulated shutdown failure") 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, ) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=2, resources_per_worker={}) ) await controller._run_control_loop_iteration() # Init -> Scheduling await controller._run_control_loop_iteration() # Scheduling -> Running 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) controller.get_worker_group().shutdown = failing_shutdown await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), ErroredState) assert isinstance(controller.get_state().training_failed_error, ControllerError) class _MockReplicaGroupBackend(Backend): has_replica_groups = True class _MockReplicaGroupBackendConfig(BackendConfig): @property def backend_cls(self): return _MockReplicaGroupBackend @pytest.mark.asyncio async def test_resize_and_fail_with_replica_groups(): """Test partial replica group replacement vs full restart with has_replica_groups. Four scenarios: 1) Same size + no poll_status → regular full restart path 2) Same size + poll_status with errors → partial replacement path 3) Different size + poll_status → regular full restart path 4) Same size + all replica groups failing → regular full restart path """ scaling_policy = MockScalingPolicy(scaling_config=ScalingConfig()) failure_policy = MockFailurePolicy(failure_config=None) train_run_context = create_dummy_run_context( backend_config=_MockReplicaGroupBackendConfig(), ) controller = TrainController( train_fn_ref=DummyObjectRefWrapper(lambda: None), train_run_context=train_run_context, scaling_policy=scaling_policy, failure_policy=failure_policy, ) # Start with 4 workers. assert isinstance(controller.get_state(), InitializingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=4, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) initial_worker_group = controller.get_worker_group() assert initial_worker_group is not None assert initial_worker_group.has_started() assert len(initial_worker_group.get_workers()) == 4 # --- Case 1: same size, no poll_status → regular full restart path --- controller.get_worker_group().error_worker(1) failure_policy.queue_decision(FailureDecision.RETRY) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RestartingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=4, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group_after_case1 = controller.get_worker_group() assert worker_group_after_case1 is not initial_worker_group assert worker_group_after_case1.has_started() assert len(worker_group_after_case1.get_workers()) == 4 # --- Case 2: same size, failure poll_status → partial replacement path --- poll_status = WorkerGroupPollStatus( worker_statuses={ 0: WorkerStatus(running=True), 1: WorkerStatus(running=False, error=RuntimeError("Worker 1 failed")), 2: WorkerStatus(running=True), 3: WorkerStatus(running=True), }, worker_rank_to_replica_group_rank={0: 0, 1: 0, 2: 1, 3: 1}, ) controller.get_worker_group().get_latest_poll_status = lambda: poll_status controller.get_worker_group().error_worker(1) failure_policy.queue_decision(FailureDecision.RETRY) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RestartingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=4, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group_after_case2 = controller.get_worker_group() assert worker_group_after_case2 is worker_group_after_case1 assert worker_group_after_case2.has_started() assert worker_group_after_case2._replaced_replica_groups == [0] # Clear the error so the next poll is clean. worker_group_after_case2.clear_worker() # --- Case 3: different size, failure poll_status → regular full restart path --- controller.get_worker_group().error_worker(2) failure_policy.queue_decision(FailureDecision.RETRY) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RestartingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=6, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group_after_case3 = controller.get_worker_group() assert worker_group_after_case3 is not worker_group_after_case2 assert worker_group_after_case3.has_started() assert len(worker_group_after_case3.get_workers()) == 6 # --- Case 4: same size, all replica groups failing → regular full restart path --- all_failing_poll_status = WorkerGroupPollStatus( worker_statuses={ 0: WorkerStatus(running=False, error=RuntimeError("Worker 0 failed")), 1: WorkerStatus(running=False, error=RuntimeError("Worker 1 failed")), 2: WorkerStatus(running=False, error=RuntimeError("Worker 2 failed")), 3: WorkerStatus(running=False, error=RuntimeError("Worker 3 failed")), 4: WorkerStatus(running=False, error=RuntimeError("Worker 4 failed")), 5: WorkerStatus(running=False, error=RuntimeError("Worker 5 failed")), }, worker_rank_to_replica_group_rank={0: 0, 1: 0, 2: 0, 3: 1, 4: 1, 5: 1}, ) controller.get_worker_group().get_latest_poll_status = ( lambda: all_failing_poll_status ) controller.get_worker_group().error_worker(0) failure_policy.queue_decision(FailureDecision.RETRY) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RestartingState) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=6, resources_per_worker={}) ) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), SchedulingState) await controller._run_control_loop_iteration() assert isinstance(controller.get_state(), RunningState) worker_group_after_case4 = controller.get_worker_group() assert worker_group_after_case4 is not worker_group_after_case3 assert worker_group_after_case4.has_started() assert len(worker_group_after_case4.get_workers()) == 6 @pytest.mark.asyncio async def test_shutdown_failure_on_errored_path(): """Shutdown failure on the errored path preserves the original training error.""" def failing_shutdown(): raise RuntimeError("Simulated shutdown failure") 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, ) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=2, resources_per_worker={}) ) await controller._run_control_loop_iteration() # Init -> Scheduling await controller._run_control_loop_iteration() # Scheduling -> Running 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 controller.get_worker_group().shutdown = failing_shutdown 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_shutdown_and_callback_both_fail_on_finished_path(): """When both worker group shutdown and shutdown callback fail on the finished path, the shutdown error takes precedence (callback error is logged).""" def failing_shutdown(): raise RuntimeError("Simulated shutdown failure") class FailingShutdownHookCallback(ControllerCallback): async def before_controller_shutdown(self): raise ValueError("Intentional error in shutdown 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=[FailingShutdownHookCallback()], ) scaling_policy.queue_recovery_decision( ResizeDecision(num_workers=2, resources_per_worker={}) ) await controller._run_control_loop_iteration() # Init -> Scheduling await controller._run_control_loop_iteration() # Scheduling -> Running 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) controller.get_worker_group().shutdown = failing_shutdown await controller._run_control_loop_iteration() # Shutdown error takes precedence over callback error. assert isinstance(controller.get_state(), ErroredState) assert isinstance(controller.get_state().training_failed_error, ControllerError) assert ( "shutdown" in str(controller.get_state().training_failed_error.controller_failure).lower() ) @pytest.mark.asyncio async def test_abort_resilient_to_callback_failure(monkeypatch): """abort() completes even when a callback raises.""" class FailingAbortCallback(ControllerCallback): def before_controller_abort(self): raise ValueError("Intentional error in abort callback") mock_exit_actor = create_autospec(ray.actor.exit_actor) monkeypatch.setattr("ray.actor.exit_actor", mock_exit_actor) 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=[FailingAbortCallback()], ) await controller.abort() assert mock_exit_actor.call_count == 1 assert isinstance(controller.get_state(), AbortedState) if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", "-x", __file__]))