wehub-resource-sync
674 lines
26 KiB
Python
674 lines
26 KiB
Python
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__]))
|