chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,3 @@
|
||||
from .controller import TrainController
|
||||
|
||||
__all__ = ["TrainController"]
|
||||
@@ -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
|
||||
)
|
||||
@@ -0,0 +1,183 @@
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
import ray
|
||||
from ray.train.v2._internal.state.util import is_actor_alive
|
||||
from ray.util.placement_group import PlacementGroup, remove_placement_group
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PlacementGroupCleaner:
|
||||
"""Detached helper that ensures PG cleanup if Ray Train Controller exits ungracefully.
|
||||
|
||||
This actor should be created with lifetime='detached' to avoid being
|
||||
fate-shared with the Train controller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
controller_actor_id: str,
|
||||
check_interval_s: float,
|
||||
get_actor_timeout_s: float,
|
||||
stop_timeout: Optional[float],
|
||||
):
|
||||
self._controller_actor_id = controller_actor_id
|
||||
self._check_interval_s = check_interval_s
|
||||
self._get_actor_timeout_s = get_actor_timeout_s
|
||||
self._stop_timeout = stop_timeout
|
||||
self._pg_queue: queue.Queue = queue.Queue()
|
||||
self._stop_event = threading.Event()
|
||||
self._monitor_thread: Optional[threading.Thread] = None
|
||||
self._exiting: bool = False
|
||||
|
||||
def register_placement_group(self, placement_group: PlacementGroup):
|
||||
logger.debug(
|
||||
"PlacementGroupCleaner registered placement group %s for controller %s",
|
||||
placement_group.id,
|
||||
self._controller_actor_id,
|
||||
)
|
||||
# Send placement group update to the monitor thread via queue
|
||||
self._pg_queue.put(placement_group)
|
||||
|
||||
def start_monitoring(self):
|
||||
"""Start monitoring the controller and placement group."""
|
||||
if self._monitor_thread is not None and self._monitor_thread.is_alive():
|
||||
# Thread already running, just return True
|
||||
logger.debug("Monitor thread already running")
|
||||
return True
|
||||
|
||||
self._monitor_thread = threading.Thread(
|
||||
target=self._monitor_loop,
|
||||
name="PlacementGroupCleanerMonitor",
|
||||
daemon=True,
|
||||
)
|
||||
self._monitor_thread.start()
|
||||
logger.debug("PlacementGroupCleaner started monitoring in background thread")
|
||||
return True
|
||||
|
||||
def _monitor_loop(self):
|
||||
"""Monitor controller; remove PG when controller is gone.
|
||||
|
||||
This runs continuously until controller dies or stop() is called.
|
||||
Uses a queue to receive placement group updates.
|
||||
"""
|
||||
curr_placement_group: Optional[PlacementGroup] = None
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
# Check for new placement group updates from queue
|
||||
try:
|
||||
pg = self._pg_queue.get(timeout=self._check_interval_s)
|
||||
curr_placement_group = pg
|
||||
logger.debug(f"Updated current placement group to {pg.id}")
|
||||
except queue.Empty:
|
||||
pass # continue to monitor current placement group
|
||||
|
||||
# Check if controller is still alive
|
||||
try:
|
||||
alive = is_actor_alive(
|
||||
actor_id=self._controller_actor_id,
|
||||
timeout=self._get_actor_timeout_s,
|
||||
)
|
||||
except ray.util.state.exception.RayStateApiException:
|
||||
logger.warning(
|
||||
"Failed to query Ray Train Controller actor state. "
|
||||
"State API may be temporarily unavailable. Continuing to monitor."
|
||||
)
|
||||
continue
|
||||
|
||||
# Cleanup if controller is dead
|
||||
if not alive:
|
||||
# Drain any queued placement groups
|
||||
while True:
|
||||
try:
|
||||
pg = self._pg_queue.get_nowait()
|
||||
curr_placement_group = pg
|
||||
except queue.Empty:
|
||||
break
|
||||
self._cleanup_placement_group(curr_placement_group)
|
||||
break
|
||||
|
||||
# Exit the actor after cleanup since controller is dead
|
||||
self._exit()
|
||||
self._monitor_thread = None
|
||||
|
||||
def _cleanup_placement_group(self, placement_group: Optional[PlacementGroup]):
|
||||
"""Clean up the current placement group if it hasn't been removed."""
|
||||
if placement_group is None:
|
||||
logger.debug("No placement group registered; skipping cleanup.")
|
||||
return
|
||||
|
||||
if self._is_placement_group_removed(placement_group):
|
||||
logger.debug(
|
||||
"Controller actor died but placement group already removed; "
|
||||
"skipping cleanup."
|
||||
)
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
f"Detected that the Ray Train controller actor ({self._controller_actor_id}) is dead. "
|
||||
f"Cleaning up placement group = [{placement_group.id}] created by this run."
|
||||
)
|
||||
try:
|
||||
remove_placement_group(placement_group)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to clean up placement group: {e}")
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
f"Placement group = [{placement_group.id}] cleaned up successfully"
|
||||
)
|
||||
|
||||
def _stop_monitor_thread(self):
|
||||
"""Stop the monitor thread and wait for it to exit.
|
||||
|
||||
Returns:
|
||||
bool: True if the thread was stopped, False if there was no active thread.
|
||||
"""
|
||||
if self._monitor_thread is None or not self._monitor_thread.is_alive():
|
||||
return False
|
||||
|
||||
# Signal stop and wait for thread to exit
|
||||
self._stop_event.set()
|
||||
self._monitor_thread.join(timeout=self._stop_timeout)
|
||||
if self._monitor_thread.is_alive():
|
||||
logger.warning(
|
||||
"Monitor thread did not exit within %.2f seconds", self._stop_timeout
|
||||
)
|
||||
return False
|
||||
|
||||
self._monitor_thread = None
|
||||
return True
|
||||
|
||||
def stop(self):
|
||||
"""Request the cleaner to stop monitoring and exit."""
|
||||
self._stop_monitor_thread()
|
||||
self._exit()
|
||||
|
||||
def _is_placement_group_removed(self, placement_group: PlacementGroup) -> bool:
|
||||
"""Check if a placement group has been removed."""
|
||||
try:
|
||||
table = ray.util.placement_group_table(placement_group)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to query placement group table: {e}. "
|
||||
"Assuming placement group is not removed."
|
||||
)
|
||||
return False
|
||||
if "state" not in table:
|
||||
return True
|
||||
return table["state"] == "REMOVED"
|
||||
|
||||
def _exit(self):
|
||||
"""Exit the actor."""
|
||||
if self._exiting:
|
||||
return
|
||||
self._exiting = True
|
||||
try:
|
||||
ray.actor.exit_actor()
|
||||
except Exception as e:
|
||||
# If exit fails for any reason, just log it.
|
||||
logger.warning(f"Failed to exit actor: {e}")
|
||||
@@ -0,0 +1,156 @@
|
||||
from enum import Enum
|
||||
|
||||
from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
|
||||
ScalingDecision,
|
||||
)
|
||||
from ray.train.v2.api.exceptions import TrainingFailedError
|
||||
|
||||
|
||||
class TrainControllerStateType(Enum):
|
||||
"""Enum representing different states of the train controller.
|
||||
|
||||
States:
|
||||
INITIALIZING: The train controller is starting up. This is always the initial
|
||||
state of the controller.
|
||||
SCHEDULING: The train controller is in the process of scheduling a new worker
|
||||
group.
|
||||
RESCHEDULING: The train controller is in the process of rescheduling the worker
|
||||
group.
|
||||
RUNNING: The train controller is actively running training tasks.
|
||||
RESTARTING: The train controller is in the process of recovering from an error.
|
||||
RESIZING: The train controller is in the process of resizing a running worker
|
||||
group.
|
||||
SHUTTING_DOWN: The train controller has already shut down the worker group and
|
||||
and is in the process of shutting itself down.
|
||||
ERRORED: A terminal state indicating that training has encountered an error and
|
||||
cannot continue.
|
||||
FINISHED: A terminal state indicating that training has completed.
|
||||
ABORTED: A terminal state indicating that training has been aborted.
|
||||
|
||||
Args:
|
||||
state_name: The name of the state.
|
||||
is_terminal: Whether this is a terminal state that should not be further processed.
|
||||
needs_new_run_attempt: Whether this state requires starting a new run attempt, where
|
||||
a run attempt is a logical unit that encompasses both scheduling workers and
|
||||
executing training on those workers.
|
||||
"""
|
||||
|
||||
INITIALIZING = ("INITIALIZING", False, True)
|
||||
SCHEDULING = ("SCHEDULING", False, False)
|
||||
RESCHEDULING = ("RESCHEDULING", False, False)
|
||||
RUNNING = ("RUNNING", False, False)
|
||||
RESTARTING = ("RESTARTING", False, True)
|
||||
RESIZING = ("RESIZING", False, True)
|
||||
SHUTTING_DOWN = ("SHUTTING_DOWN", False, False)
|
||||
ERRORED = ("ERRORED", True, False)
|
||||
FINISHED = ("FINISHED", True, False)
|
||||
ABORTED = ("ABORTED", True, False)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
state_name: str,
|
||||
is_terminal: bool,
|
||||
needs_new_run_attempt: bool,
|
||||
):
|
||||
self.state_name = state_name
|
||||
self.is_terminal = is_terminal
|
||||
self.needs_new_run_attempt = needs_new_run_attempt
|
||||
|
||||
|
||||
class TrainControllerState:
|
||||
"""Base class for all train controller states.
|
||||
|
||||
Methods:
|
||||
get_type() -> TrainControllerStateType: Returns the type of the state.
|
||||
is_terminal() -> bool: Returns whether the state is terminal.
|
||||
needs_new_run_attempt() -> bool: Returns whether a new run attempt is needed.
|
||||
"""
|
||||
|
||||
def __init__(self, state_type: TrainControllerStateType):
|
||||
self._state_type = state_type
|
||||
|
||||
def __repr__(self) -> str:
|
||||
attrs = {
|
||||
"type": self._state_type.name,
|
||||
"is_terminal": self._state_type.is_terminal,
|
||||
"needs_new_run_attempt": self._state_type.needs_new_run_attempt,
|
||||
**{k: v for k, v in vars(self).items() if not k.startswith("_")},
|
||||
}
|
||||
attrs_str = "\n ".join(f"{k}={v}" for k, v in attrs.items())
|
||||
return f"{self.__class__.__name__}(\n {attrs_str}\n)"
|
||||
|
||||
def is_terminal(self) -> bool:
|
||||
return self._state_type.is_terminal
|
||||
|
||||
def needs_new_run_attempt(self) -> bool:
|
||||
return self._state_type.needs_new_run_attempt
|
||||
|
||||
|
||||
class InitializingState(TrainControllerState):
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.INITIALIZING)
|
||||
|
||||
|
||||
class SchedulingState(TrainControllerState):
|
||||
def __init__(self, scaling_decision: ScalingDecision):
|
||||
super().__init__(state_type=TrainControllerStateType.SCHEDULING)
|
||||
self.scaling_decision = scaling_decision
|
||||
|
||||
|
||||
class ReschedulingState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.RESCHEDULING)
|
||||
self.training_failed_error = training_failed_error
|
||||
|
||||
|
||||
class RunningState(TrainControllerState):
|
||||
# TODO: Split into multiple more granular states, or add more fields.
|
||||
# For example, we may want to indicate if any health checks failed.
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.RUNNING)
|
||||
|
||||
|
||||
class RestartingState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.RESTARTING)
|
||||
self.training_failed_error = training_failed_error
|
||||
|
||||
|
||||
class ResizingState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
scaling_decision: ScalingDecision,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.RESIZING)
|
||||
self.scaling_decision = scaling_decision
|
||||
|
||||
|
||||
class ShuttingDownState(TrainControllerState):
|
||||
def __init__(self, next_state: "TrainControllerState"):
|
||||
super().__init__(state_type=TrainControllerStateType.SHUTTING_DOWN)
|
||||
self.next_state = next_state
|
||||
|
||||
|
||||
class ErroredState(TrainControllerState):
|
||||
def __init__(
|
||||
self,
|
||||
training_failed_error: TrainingFailedError,
|
||||
):
|
||||
super().__init__(state_type=TrainControllerStateType.ERRORED)
|
||||
self.training_failed_error = training_failed_error
|
||||
|
||||
|
||||
class FinishedState(TrainControllerState):
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.FINISHED)
|
||||
|
||||
|
||||
class AbortedState(TrainControllerState):
|
||||
def __init__(self):
|
||||
super().__init__(state_type=TrainControllerStateType.ABORTED)
|
||||
Reference in New Issue
Block a user