chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,865 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Callable, List, Optional
|
||||
|
||||
import pandas as pd
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray.exceptions import AsyncioActorExit
|
||||
from ray.train.v2._internal.constants import (
|
||||
DEFAULT_ENABLE_CONTROLLER_LOGGING,
|
||||
DEFAULT_ENABLE_PREEMPTION_WATCHER,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL_S,
|
||||
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR,
|
||||
HEALTH_CHECK_INTERVAL_S_ENV_VAR,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback import (
|
||||
ControllerCallback,
|
||||
ReportCallback,
|
||||
TrainContextCallback,
|
||||
WorkerCallback,
|
||||
WorkerGroupCallback,
|
||||
)
|
||||
from ray.train.v2._internal.execution.callback_manager import CallbackManager
|
||||
from ray.train.v2._internal.execution.checkpoint.checkpoint_manager import (
|
||||
CheckpointManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.checkpoint.report_handler import (
|
||||
ReportCallbackHandler,
|
||||
)
|
||||
from ray.train.v2._internal.execution.checkpoint.validation_manager import (
|
||||
ValidationManager,
|
||||
)
|
||||
from ray.train.v2._internal.execution.context import TrainRunContext
|
||||
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,
|
||||
FailurePolicy,
|
||||
)
|
||||
from ray.train.v2._internal.execution.scaling_policy import (
|
||||
NoopDecision,
|
||||
ResizeDecision,
|
||||
ScalingPolicy,
|
||||
)
|
||||
from ray.train.v2._internal.execution.worker_group import (
|
||||
WorkerGroup,
|
||||
WorkerGroupContext,
|
||||
WorkerGroupPollStatus,
|
||||
)
|
||||
from ray.train.v2._internal.logging import LoggingManager
|
||||
from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
|
||||
from ray.train.v2.api.callback import RayTrainCallback
|
||||
from ray.train.v2.api.exceptions import (
|
||||
ControllerError,
|
||||
TrainingFailedError,
|
||||
)
|
||||
from ray.train.v2.api.report_config import CheckpointConsistencyMode
|
||||
from ray.train.v2.api.result import Result
|
||||
from ray.train.v2.api.validation_config import ValidationConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint
|
||||
|
||||
from ray.util.tpu import get_tpu_num_slices_for_workers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainControllerLoopIterationResult:
|
||||
"""The result of a single iteration of the control loop."""
|
||||
|
||||
run_attempt_id: str
|
||||
previous_state: TrainControllerState
|
||||
next_state: TrainControllerState
|
||||
training_failed_error: Optional[TrainingFailedError] = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"TrainControllerLoopIterationResult(\n"
|
||||
f" run_attempt_id={self.run_attempt_id},\n"
|
||||
f" previous_state={self.previous_state._state_type.state_name},\n"
|
||||
f" next_state={self.next_state._state_type.state_name}\n"
|
||||
f" training_failed_error={self.training_failed_error}\n"
|
||||
f")"
|
||||
)
|
||||
|
||||
|
||||
class TrainController:
|
||||
"""Manages the execution of a distributed training job.
|
||||
|
||||
Responsibilities include:
|
||||
* Triggering the training function to run on the worker group.
|
||||
* Monitoring the status of the worker group.
|
||||
* Handling scaling decisions by restarting the worker group.
|
||||
* Handling failure decisions by restarting the worker group or terminating training.
|
||||
* Running callback logic on different hooks in the control loop.
|
||||
"""
|
||||
|
||||
worker_group_cls = WorkerGroup
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_fn_ref: ObjectRefWrapper[Callable[[], None]],
|
||||
train_run_context: TrainRunContext,
|
||||
scaling_policy: ScalingPolicy,
|
||||
failure_policy: FailurePolicy,
|
||||
callbacks: Optional[List[RayTrainCallback]] = None,
|
||||
validation_config: Optional[ValidationConfig] = None,
|
||||
):
|
||||
self._train_run_context = train_run_context
|
||||
if ray_constants.env_bool(
|
||||
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
|
||||
DEFAULT_ENABLE_CONTROLLER_LOGGING,
|
||||
):
|
||||
LoggingManager.configure_controller_logger(self._train_run_context)
|
||||
self._train_fn_ref = train_fn_ref
|
||||
self._scaling_policy = scaling_policy
|
||||
self._failure_policy = failure_policy
|
||||
self._run_config = self._train_run_context.run_config
|
||||
self._callbacks = callbacks or []
|
||||
self._storage_context = self._train_run_context.run_config.storage_context
|
||||
|
||||
self._checkpoint_manager = CheckpointManager(
|
||||
checkpoint_config=self._run_config.checkpoint_config,
|
||||
storage_context=self._storage_context,
|
||||
)
|
||||
if validation_config:
|
||||
validation_manager = ValidationManager(
|
||||
checkpoint_manager=self._checkpoint_manager,
|
||||
validation_config=validation_config,
|
||||
)
|
||||
else:
|
||||
validation_manager = None
|
||||
report_handler = ReportCallbackHandler(
|
||||
report_callbacks=(
|
||||
[self._checkpoint_manager]
|
||||
+ ([validation_manager] if validation_manager else [])
|
||||
+ [c for c in self._callbacks if isinstance(c, ReportCallback)]
|
||||
)
|
||||
)
|
||||
|
||||
# Group callbacks by the hooks they're subscribed to.
|
||||
self._controller_callbacks = (
|
||||
[
|
||||
self._scaling_policy,
|
||||
]
|
||||
+ ([validation_manager] if validation_manager else [])
|
||||
+ [c for c in self._callbacks if isinstance(c, ControllerCallback)]
|
||||
)
|
||||
self._controller_callback_manager = CallbackManager(self._controller_callbacks)
|
||||
# Group callbacks that will be propagated to the worker group,
|
||||
# train worker and the train context.
|
||||
self._worker_group_callbacks_to_propagate = (
|
||||
[report_handler]
|
||||
+ ([validation_manager] if validation_manager else [])
|
||||
+ [
|
||||
c
|
||||
for c in self._callbacks
|
||||
if isinstance(
|
||||
c, (WorkerGroupCallback, WorkerCallback, TrainContextCallback)
|
||||
)
|
||||
]
|
||||
+ [self._checkpoint_manager]
|
||||
)
|
||||
|
||||
self._health_check_interval_s = float(
|
||||
os.getenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, DEFAULT_HEALTH_CHECK_INTERVAL_S)
|
||||
)
|
||||
|
||||
self._manages_replica_groups = (
|
||||
train_run_context.backend_config.backend_cls.has_replica_groups
|
||||
if train_run_context.backend_config
|
||||
else False
|
||||
)
|
||||
|
||||
# Register the preemption-observability callback when not in TorchFT
|
||||
# mode (replica groups handle peer loss via their own quorum).
|
||||
enable_preemption_watcher = ray_constants.env_bool(
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR,
|
||||
DEFAULT_ENABLE_PREEMPTION_WATCHER,
|
||||
)
|
||||
if self._manages_replica_groups:
|
||||
if enable_preemption_watcher and ray_constants.env_set_by_user(
|
||||
ENABLE_PREEMPTION_WATCHER_ENV_VAR
|
||||
):
|
||||
logger.info(
|
||||
"The preemption watcher is not compatible with replica "
|
||||
"groups (e.g. TorchFT), which handle peer loss via their "
|
||||
"own quorum; skipping it."
|
||||
)
|
||||
elif enable_preemption_watcher:
|
||||
from ray.train.v2._internal.callbacks.preemption_callback import (
|
||||
PreemptionCallback,
|
||||
)
|
||||
|
||||
self._worker_group_callbacks_to_propagate.append(PreemptionCallback())
|
||||
|
||||
self._worker_group: Optional[WorkerGroup] = None
|
||||
self._state = InitializingState()
|
||||
self._return_value: Optional[Any] = None
|
||||
|
||||
# TODO: These can be attributes of a RunAttempt?
|
||||
self._latest_poll_time = float("-inf")
|
||||
|
||||
# Generate an initial run attempt ID so that `_run_controller_hook`
|
||||
# can reference it if a callback fails during `_start`.
|
||||
self._generate_run_attempt_id()
|
||||
self._start()
|
||||
|
||||
def _run_controller_hook(
|
||||
self,
|
||||
hook_name: str,
|
||||
*args,
|
||||
invoke_failure_decision_callbacks: bool = True,
|
||||
**context,
|
||||
) -> Optional["TrainControllerLoopIterationResult"]:
|
||||
"""Invoke a named controller hook and catch any exceptions.
|
||||
|
||||
This method invokes all callbacks registered for the given controller hook.
|
||||
If a callback raises an error, the error is routed through the failure policy
|
||||
and may produce a ``TrainControllerLoopIterationResult``, indicating that the
|
||||
current controller step should exit early with this failure result.
|
||||
|
||||
Args:
|
||||
hook_name: The controller hook name to invoke.
|
||||
*args: Positional arguments to pass to the hook.
|
||||
invoke_failure_decision_callbacks: Whether to invoke failure-decision hooks
|
||||
when handling a callback failure.
|
||||
**context: Keyword arguments to pass to the hook.
|
||||
|
||||
Returns:
|
||||
failure_result: A``TrainControllerLoopIterationResult`` if the hook execution results
|
||||
in an early exit from the controller loop to raise the callback error,
|
||||
or ``None`` if hook execution completes successfully.
|
||||
"""
|
||||
try:
|
||||
self._controller_callback_manager.invoke(hook_name, *args, **context)
|
||||
except ControllerError as error:
|
||||
failure_decision = self._failure_policy.make_decision(
|
||||
training_failed_error=error,
|
||||
)
|
||||
# Avoid re-entering controller callback hooks while handling a callback failure.
|
||||
return self._execute_failure_decision(
|
||||
failure_decision,
|
||||
training_failed_error=error,
|
||||
invoke_failure_decision_callbacks=invoke_failure_decision_callbacks,
|
||||
)
|
||||
return None
|
||||
|
||||
def _execute_resize_decision(
|
||||
self, decision: ResizeDecision
|
||||
) -> TrainControllerLoopIterationResult:
|
||||
"""Executes resize decisions.
|
||||
|
||||
Errors from worker group shutdown, callbacks, or worker group startup
|
||||
are allowed to propagate to the catch-all in ``run()``.
|
||||
"""
|
||||
|
||||
failure_result = self._run_controller_hook(
|
||||
"before_controller_execute_resize_decision", decision
|
||||
)
|
||||
if failure_result:
|
||||
return failure_result
|
||||
current_num_workers = (
|
||||
len(self._worker_group.get_workers()) if self._worker_group else 0
|
||||
)
|
||||
poll_status = (
|
||||
self._worker_group.get_latest_poll_status() if self._worker_group else None
|
||||
)
|
||||
failing_rgs = (
|
||||
poll_status.failing_replica_group_indices if poll_status else set()
|
||||
)
|
||||
all_rgs = poll_status.all_replica_group_indices if poll_status else set()
|
||||
if (
|
||||
self._manages_replica_groups
|
||||
and bool(failing_rgs)
|
||||
and failing_rgs != all_rgs
|
||||
and self._worker_group
|
||||
# TODO: relax this after integrating replica groups with elastic training.
|
||||
and decision.num_workers == current_num_workers
|
||||
):
|
||||
# Torchft: replace only failing replica groups.
|
||||
self._replace_bad_workers(poll_status)
|
||||
else:
|
||||
# Standard: full restart.
|
||||
if self._worker_group:
|
||||
self._shutdown_worker_group()
|
||||
self._start_worker_group(
|
||||
num_workers=decision.num_workers,
|
||||
resources_per_worker=decision.resources_per_worker,
|
||||
)
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=self._state,
|
||||
next_state=RunningState(),
|
||||
)
|
||||
|
||||
def _replace_bad_workers(self, poll_status: WorkerGroupPollStatus):
|
||||
"""Replace failing replica groups in the worker group.
|
||||
|
||||
Args:
|
||||
poll_status: The poll status containing error information.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
failing_rg_indices = poll_status.failing_replica_group_indices
|
||||
|
||||
if not failing_rg_indices:
|
||||
logger.warning("No failing replica groups found in poll status.")
|
||||
return
|
||||
|
||||
logger.info(f"Replacing failing replica groups: {failing_rg_indices}")
|
||||
|
||||
for rg_index in failing_rg_indices:
|
||||
# TODO: parallelize this.
|
||||
# TODO: also ensure that if earlier replacements succeed and later replacements fail,
|
||||
# we don't redo the earlier replacements.
|
||||
# See https://github.com/ray-project/ray/pull/61475#discussion_r3055217289
|
||||
self._worker_group.replace_replica_group(rg_index)
|
||||
|
||||
def _get_retry_state(
|
||||
self,
|
||||
controller_state: TrainControllerState,
|
||||
training_failed_error: TrainingFailedError,
|
||||
) -> TrainControllerState:
|
||||
if isinstance(controller_state, RunningState):
|
||||
return RestartingState(training_failed_error=training_failed_error)
|
||||
elif isinstance(controller_state, SchedulingState):
|
||||
return ReschedulingState(training_failed_error=training_failed_error)
|
||||
else:
|
||||
# Cannot retry from this state (e.g. InitializingState,
|
||||
# ShuttingDownState); force shutdown with error.
|
||||
logger.warning(
|
||||
"Cannot retry from state %s; forcing shutdown.",
|
||||
type(controller_state).__name__,
|
||||
)
|
||||
return ShuttingDownState(
|
||||
next_state=ErroredState(training_failed_error=training_failed_error)
|
||||
)
|
||||
|
||||
def _execute_failure_decision(
|
||||
self,
|
||||
failure_decision: FailureDecision,
|
||||
training_failed_error: TrainingFailedError,
|
||||
invoke_failure_decision_callbacks: bool = True,
|
||||
) -> TrainControllerLoopIterationResult:
|
||||
"""Executes failure handling decisions for a scheduling or poll error."""
|
||||
|
||||
controller_state = self.get_state()
|
||||
|
||||
if invoke_failure_decision_callbacks:
|
||||
failure_result = self._run_controller_hook(
|
||||
"before_controller_execute_failure_decision",
|
||||
failure_decision,
|
||||
invoke_failure_decision_callbacks=False,
|
||||
)
|
||||
if failure_result:
|
||||
return failure_result
|
||||
|
||||
# TODO: What should we do here?
|
||||
# This currently never happens because there must be errors.
|
||||
if failure_decision == FailureDecision.NOOP:
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=controller_state,
|
||||
training_failed_error=training_failed_error,
|
||||
)
|
||||
|
||||
if failure_decision == FailureDecision.RETRY:
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=self._get_retry_state(
|
||||
controller_state, training_failed_error
|
||||
),
|
||||
)
|
||||
elif failure_decision == FailureDecision.RAISE:
|
||||
next_state = ShuttingDownState(
|
||||
next_state=ErroredState(
|
||||
training_failed_error=training_failed_error,
|
||||
),
|
||||
)
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=next_state,
|
||||
training_failed_error=training_failed_error,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unexpected failure decision: {failure_decision}")
|
||||
|
||||
async def _poll_workers(self) -> WorkerGroupPollStatus:
|
||||
# Ensure that the time between polls is at least HEALTH_CHECK_INTERVAL_S.
|
||||
time_since_last_poll = time_monotonic() - self._latest_poll_time
|
||||
if time_since_last_poll < self._health_check_interval_s:
|
||||
remaining_time = max(
|
||||
self._health_check_interval_s - time_since_last_poll, 0
|
||||
)
|
||||
await asyncio.sleep(remaining_time)
|
||||
if self.get_state().is_terminal():
|
||||
logger.debug(
|
||||
f"Controller is unexpectedly in terminal state {self.get_state()} after "
|
||||
"sleeping and before polling workers. Exiting actor."
|
||||
)
|
||||
ray.actor.exit_actor()
|
||||
|
||||
status = self._worker_group.poll_status(timeout=self._health_check_interval_s)
|
||||
self._latest_poll_time = time_monotonic()
|
||||
return status
|
||||
|
||||
def _start_worker_group(self, num_workers: int, resources_per_worker: dict) -> None:
|
||||
"""Start the worker group and launch the train function.
|
||||
|
||||
Args:
|
||||
num_workers: The number of workers to start.
|
||||
resources_per_worker: The resources per worker to start.
|
||||
|
||||
Raises:
|
||||
Exception: If the worker group failed to start.
|
||||
"""
|
||||
placement_strategy = self._scaling_policy.scaling_config.placement_strategy
|
||||
scaling_config = self._train_run_context.scaling_config
|
||||
|
||||
# Check for `label_selector` to influence WorkerGroup scheduling.
|
||||
label_selector = scaling_config._label_selector_per_worker(num_workers)
|
||||
for callback in self._controller_callbacks:
|
||||
selector = callback.on_controller_start_worker_group(
|
||||
scaling_config=scaling_config, num_workers=num_workers
|
||||
)
|
||||
if selector:
|
||||
if label_selector:
|
||||
logger.warning(
|
||||
f"Overriding `ScalingConfig.label_selector` {label_selector} "
|
||||
f"with label_selector returned by user-specified callback {selector}"
|
||||
)
|
||||
label_selector = [selector.copy() for _ in range(num_workers)]
|
||||
|
||||
# Calculate num_slices for the worker group if using TPU.
|
||||
num_slices = 1
|
||||
if scaling_config.use_tpu:
|
||||
num_slices = get_tpu_num_slices_for_workers(
|
||||
topology=scaling_config.topology,
|
||||
accelerator_type=scaling_config.accelerator_type,
|
||||
num_workers=num_workers,
|
||||
resources_per_worker=resources_per_worker,
|
||||
)
|
||||
|
||||
worker_group_context = WorkerGroupContext(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
train_fn_ref=self._train_fn_ref,
|
||||
num_workers=num_workers,
|
||||
resources_per_worker=resources_per_worker,
|
||||
placement_strategy=placement_strategy,
|
||||
label_selector=label_selector,
|
||||
num_slices=num_slices,
|
||||
)
|
||||
self._worker_group = self.worker_group_cls.create(
|
||||
train_run_context=self._train_run_context,
|
||||
worker_group_context=worker_group_context,
|
||||
callbacks=self._worker_group_callbacks_to_propagate,
|
||||
)
|
||||
|
||||
def _start(self):
|
||||
failure_result = self._run_controller_hook(
|
||||
"after_controller_start", self._train_run_context
|
||||
)
|
||||
if failure_result:
|
||||
self._set_state(failure_result.next_state)
|
||||
|
||||
async def _shutdown(self) -> "TrainControllerLoopIterationResult":
|
||||
"""Execute shutdown and return the final state transition.
|
||||
|
||||
Shutdown errors are never retried. If an error occurs during shutdown:
|
||||
- If we're already shutting down after a training error
|
||||
(next_state is ErroredState), the original error is preserved.
|
||||
- Otherwise the shutdown error becomes the training failure.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
assert isinstance(controller_state, ShuttingDownState)
|
||||
|
||||
shutdown_error = None
|
||||
|
||||
# TODO: move to __del__ after https://github.com/ray-project/ray/issues/53169
|
||||
if self._worker_group:
|
||||
try:
|
||||
self._shutdown_worker_group()
|
||||
except Exception as e:
|
||||
logger.exception("Error shutting down worker group.")
|
||||
shutdown_error = ControllerError(e)
|
||||
|
||||
try:
|
||||
await self._controller_callback_manager.async_invoke(
|
||||
"before_controller_shutdown"
|
||||
)
|
||||
except ControllerError as e:
|
||||
if shutdown_error:
|
||||
logger.warning(
|
||||
"An additional error occurred in the before_controller_shutdown "
|
||||
"callback after a worker group shutdown error. "
|
||||
"This error is being ignored to preserve the original "
|
||||
"shutdown error. Error: %s",
|
||||
e,
|
||||
)
|
||||
else:
|
||||
shutdown_error = e
|
||||
|
||||
if shutdown_error:
|
||||
if isinstance(controller_state.next_state, ErroredState):
|
||||
logger.warning(
|
||||
"Another error occurred during shutdown after a training error. "
|
||||
"This error is being ignored to preserve the original "
|
||||
"training error. Error: %s",
|
||||
shutdown_error,
|
||||
)
|
||||
else:
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=ErroredState(training_failed_error=shutdown_error),
|
||||
training_failed_error=shutdown_error,
|
||||
)
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=controller_state.next_state,
|
||||
)
|
||||
|
||||
def _shutdown_worker_group(self):
|
||||
"""Shutdown the worker group and set the worker group to None."""
|
||||
self._worker_group.shutdown()
|
||||
self._worker_group = None
|
||||
|
||||
def get_worker_group(self) -> Optional[WorkerGroup]:
|
||||
return self._worker_group
|
||||
|
||||
def get_state(self) -> TrainControllerState:
|
||||
return self._state
|
||||
|
||||
def _set_state(self, state: TrainControllerState):
|
||||
previous_state = self._state
|
||||
self._state = state
|
||||
|
||||
failure_result = self._run_controller_hook(
|
||||
"after_controller_state_update", previous_state, state
|
||||
)
|
||||
if failure_result:
|
||||
# If we're transitioning into a terminal state, or if we're already in the shutdown path to an errored terminal state
|
||||
# (ShuttingDownState -> ErroredState), preserve the original failure as the
|
||||
# surfaced error. A failure in a state-update callback should not overwrite
|
||||
# the underlying root-cause error.
|
||||
if state.is_terminal() or (
|
||||
isinstance(state, ShuttingDownState)
|
||||
and isinstance(state.next_state, ErroredState)
|
||||
):
|
||||
logger.warning(
|
||||
"A callback failed during a terminal state transition. "
|
||||
"This failure is being ignored to preserve the original "
|
||||
"training result. Error: %s",
|
||||
failure_result.training_failed_error,
|
||||
)
|
||||
return
|
||||
|
||||
# NOTE: We intentionally do *not* re-invoke `after_controller_state_update`
|
||||
# for this transition to avoid re-entering callback hooks while handling
|
||||
# a callback failure.
|
||||
self._state = failure_result.next_state
|
||||
|
||||
def _make_and_handle_scaling_decision_for_non_running_worker_group(
|
||||
self,
|
||||
controller_state: TrainControllerState,
|
||||
) -> TrainControllerLoopIterationResult:
|
||||
"""Make a scaling decision for a non-running worker group and return the appropriate next state.
|
||||
|
||||
This method should be called when entering a state that requires a scaling decision
|
||||
for a non-running worker group.
|
||||
|
||||
This method handles the complete flow of:
|
||||
1. Shutting down the non-running worker group if it still exists.
|
||||
2. Getting a scaling decision for a non-running worker group
|
||||
3. Determining the next state based on the decision type
|
||||
4. Creating and returning the iteration result
|
||||
|
||||
Args:
|
||||
controller_state: The current controller state
|
||||
|
||||
Returns:
|
||||
TrainControllerLoopIterationResult with the appropriate next state
|
||||
"""
|
||||
scaling_decision = (
|
||||
self._scaling_policy.make_decision_for_non_running_worker_group()
|
||||
)
|
||||
|
||||
if isinstance(scaling_decision, NoopDecision):
|
||||
next_state = controller_state
|
||||
elif isinstance(scaling_decision, ResizeDecision):
|
||||
next_state = SchedulingState(scaling_decision)
|
||||
else:
|
||||
raise ValueError(f"Unexpected scaling decision: {scaling_decision}")
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=next_state,
|
||||
)
|
||||
|
||||
async def _step(self) -> TrainControllerLoopIterationResult:
|
||||
"""Run a single iteration of the control loop.
|
||||
|
||||
Returns:
|
||||
The result of the iteration.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
|
||||
if isinstance(
|
||||
controller_state, (InitializingState, RestartingState, ReschedulingState)
|
||||
):
|
||||
return self._make_and_handle_scaling_decision_for_non_running_worker_group(
|
||||
controller_state
|
||||
)
|
||||
elif isinstance(controller_state, SchedulingState):
|
||||
assert isinstance(controller_state.scaling_decision, ResizeDecision)
|
||||
return self._execute_resize_decision(controller_state.scaling_decision)
|
||||
elif isinstance(controller_state, RunningState):
|
||||
worker_group_status: WorkerGroupPollStatus = await self._poll_workers()
|
||||
|
||||
if worker_group_status.finished and not worker_group_status.errors:
|
||||
self._return_value = worker_group_status.worker_statuses[0].return_value
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=ShuttingDownState(
|
||||
next_state=FinishedState(),
|
||||
),
|
||||
)
|
||||
if worker_group_status.errors:
|
||||
worker_group_error = worker_group_status.get_worker_group_error()
|
||||
failure_decision = self._failure_policy.make_decision(
|
||||
training_failed_error=worker_group_error,
|
||||
)
|
||||
return self._execute_failure_decision(
|
||||
failure_decision, training_failed_error=worker_group_error
|
||||
)
|
||||
|
||||
scaling_decision = (
|
||||
self._scaling_policy.make_decision_for_running_worker_group(
|
||||
worker_group_state=self.get_worker_group().get_worker_group_state(),
|
||||
worker_group_status=worker_group_status,
|
||||
)
|
||||
)
|
||||
|
||||
if isinstance(scaling_decision, NoopDecision):
|
||||
next_state = RunningState()
|
||||
elif isinstance(scaling_decision, ResizeDecision):
|
||||
next_state = ResizingState(
|
||||
scaling_decision=scaling_decision,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unexpected scaling decision: {scaling_decision}")
|
||||
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=next_state,
|
||||
)
|
||||
elif isinstance(controller_state, ResizingState):
|
||||
return TrainControllerLoopIterationResult(
|
||||
run_attempt_id=self._get_run_attempt_id(),
|
||||
previous_state=controller_state,
|
||||
next_state=SchedulingState(
|
||||
scaling_decision=controller_state.scaling_decision
|
||||
),
|
||||
)
|
||||
elif isinstance(controller_state, ShuttingDownState):
|
||||
return await self._shutdown()
|
||||
else:
|
||||
raise ValueError(f"Unexpected controller state: {controller_state}")
|
||||
|
||||
def _generate_run_attempt_id(self):
|
||||
self._run_attempt_id = uuid.uuid4().hex
|
||||
return self._run_attempt_id
|
||||
|
||||
def _get_run_attempt_id(self):
|
||||
return self._run_attempt_id
|
||||
|
||||
async def _run_control_loop_iteration(self):
|
||||
"""Run a single iteration of the control loop.
|
||||
|
||||
Steps:
|
||||
1. Poll the worker group for status.
|
||||
2. If the worker group is initializing or recovering from an error,
|
||||
make a scaling decision and execute it.
|
||||
3. If the worker group has finished, set the controller state to FINISHED.
|
||||
4. If the worker group has errors, make a failure decision and execute it.
|
||||
5. Otherwise, the worker group is running healthily.
|
||||
Query the scaling policy for a scaling decision and execute it.
|
||||
|
||||
Errors raised by ``_step`` are caught and routed through the failure
|
||||
policy (retry / raise). If the failure policy itself fails, the
|
||||
controller is forced into ``ErroredState`` as a last resort.
|
||||
|
||||
``AsyncioActorExit`` is always re-raised so that the actor can shut
|
||||
down cleanly.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
assert not controller_state.is_terminal()
|
||||
|
||||
if controller_state.needs_new_run_attempt():
|
||||
self._generate_run_attempt_id()
|
||||
|
||||
try:
|
||||
result = await self._step()
|
||||
except AsyncioActorExit:
|
||||
raise
|
||||
except Exception as e:
|
||||
# Preserve the original error type if it is already a
|
||||
# TrainingFailedError (e.g. WorkerGroupError); otherwise
|
||||
# wrap it in a ControllerError.
|
||||
if isinstance(e, TrainingFailedError):
|
||||
training_error = e
|
||||
else:
|
||||
# Log the full traceback only for unexpected errors.
|
||||
logger.exception("Error in control loop iteration: %s", e)
|
||||
training_error = ControllerError(e)
|
||||
try:
|
||||
failure_decision = self._failure_policy.make_decision(
|
||||
training_failed_error=training_error,
|
||||
)
|
||||
result = self._execute_failure_decision(
|
||||
failure_decision,
|
||||
training_failed_error=training_error,
|
||||
)
|
||||
except Exception:
|
||||
# Last resort: force into errored state, bypassing callbacks.
|
||||
logger.exception(
|
||||
"Failed to execute failure decision, forcing error state."
|
||||
)
|
||||
self._state = ErroredState(training_failed_error=training_error)
|
||||
return
|
||||
|
||||
self._set_state(result.next_state)
|
||||
|
||||
async def run(self):
|
||||
"""Run the main control loop. Exits when training is finished or errored."""
|
||||
while not self.get_state().is_terminal():
|
||||
await self._run_control_loop_iteration()
|
||||
|
||||
# Call after_controller_finish with the final result.
|
||||
result = self._build_result()
|
||||
failure_result = self._run_controller_hook(
|
||||
"after_controller_finish", result, invoke_failure_decision_callbacks=False
|
||||
)
|
||||
# Since we are already in a terminal state, a callback failure should
|
||||
# not overwrite the training outcome — log and preserve the result.
|
||||
if failure_result:
|
||||
logger.warning(
|
||||
"A callback failed after training finished. "
|
||||
"This failure is being ignored to preserve the original "
|
||||
"training result. Error: %s",
|
||||
failure_result.training_failed_error,
|
||||
)
|
||||
|
||||
async def abort(self):
|
||||
"""Trigger callback abort hooks and terminate the controller process."""
|
||||
# Do not abort run if it's already finished.
|
||||
if self.get_state().is_terminal():
|
||||
return
|
||||
|
||||
self._controller_callback_manager.invoke_best_effort("before_controller_abort")
|
||||
|
||||
# Intentionally abort worker group before setting train run state because
|
||||
# we only reconcile the states of live train runs.
|
||||
try:
|
||||
if self._worker_group:
|
||||
self._worker_group.abort()
|
||||
self._set_state(AbortedState())
|
||||
except Exception as e:
|
||||
logger.exception("Error aborting worker group: %s", e)
|
||||
|
||||
ray.actor.exit_actor()
|
||||
|
||||
def _build_result(self) -> Result:
|
||||
storage = self._checkpoint_manager._storage_context
|
||||
|
||||
latest_checkpoint_result = self._checkpoint_manager.latest_checkpoint_result
|
||||
latest_metrics = (
|
||||
latest_checkpoint_result.metrics if latest_checkpoint_result else None
|
||||
)
|
||||
latest_checkpoint = (
|
||||
latest_checkpoint_result.checkpoint if latest_checkpoint_result else None
|
||||
)
|
||||
best_checkpoints = [
|
||||
(r.checkpoint, r.metrics)
|
||||
for r in self._checkpoint_manager.best_checkpoint_results
|
||||
]
|
||||
|
||||
# Provide the history of metrics attached to checkpoints as a dataframe.
|
||||
metrics_dataframe = None
|
||||
if best_checkpoints:
|
||||
metrics_dataframe = pd.DataFrame([m for _, m in best_checkpoints])
|
||||
|
||||
return Result(
|
||||
metrics=latest_metrics,
|
||||
checkpoint=latest_checkpoint,
|
||||
error=self.get_training_failed_error(),
|
||||
path=storage.experiment_fs_path,
|
||||
best_checkpoints=best_checkpoints,
|
||||
metrics_dataframe=metrics_dataframe,
|
||||
_storage_filesystem=storage.storage_filesystem,
|
||||
return_value=self._return_value,
|
||||
)
|
||||
|
||||
def get_result(self) -> Result:
|
||||
"""Get the final training result from the TrainController."""
|
||||
|
||||
controller_state = self.get_state()
|
||||
if not controller_state.is_terminal():
|
||||
raise ValueError(
|
||||
f"Cannot get result when controller is in state {controller_state}"
|
||||
)
|
||||
return self._build_result()
|
||||
|
||||
def get_training_failed_error(self) -> Optional[TrainingFailedError]:
|
||||
"""Get the training failed error from the controller state.
|
||||
|
||||
Returns:
|
||||
The training failed error if the controller is in an errored state,
|
||||
None otherwise.
|
||||
"""
|
||||
controller_state = self.get_state()
|
||||
|
||||
if isinstance(controller_state, ErroredState):
|
||||
return controller_state.training_failed_error
|
||||
|
||||
return None
|
||||
|
||||
async def get_all_reported_checkpoints(
|
||||
self,
|
||||
current_report_index: int,
|
||||
consistency_mode: CheckpointConsistencyMode = CheckpointConsistencyMode.VALIDATED,
|
||||
timeout_s: Optional[float] = None,
|
||||
) -> List["ReportedCheckpoint"]:
|
||||
return await self._checkpoint_manager.get_all_reported_checkpoints(
|
||||
current_report_index, consistency_mode, timeout_s
|
||||
)
|
||||
Reference in New Issue
Block a user