chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/execution/failure_handling/__init__.py

@@ -0,0 +1,16 @@
+# isort: off
+from .failure_policy import FailureDecision, FailurePolicy
+from .default import DefaultFailurePolicy
+from .factory import create_failure_policy
+
+# isort: on
+
+__all__ = [
+    "DefaultFailurePolicy",
+    "FailureDecision",
+    "FailurePolicy",
+    "create_failure_policy",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/v2/_internal/execution/failure_handling/default.py

@@ -0,0 +1,100 @@
+import logging
+
+from .failure_policy import FailureDecision, FailurePolicy
+from ray.train.v2._internal.exceptions import (
+    WorkerGroupStartupFailedError,
+    WorkerGroupStartupTimeoutError,
+)
+from ray.train.v2.api.config import FailureConfig
+from ray.train.v2.api.exceptions import (
+    ControllerError,
+    TrainingFailedError,
+    WorkerGroupError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+RETRYABLE_CONTROLLER_ERRORS = (
+    WorkerGroupStartupFailedError,
+    WorkerGroupStartupTimeoutError,
+)
+
+
+class DefaultFailurePolicy(FailurePolicy):
+    def __init__(self, failure_config: FailureConfig):
+        super().__init__(failure_config)
+        self._worker_group_failures = 0
+        self._controller_failures = 0
+
+    def _log_decision(
+        self,
+        decision: FailureDecision,
+        training_failed_error: TrainingFailedError,
+        error_count: int,
+        retry_limit: int,
+    ):
+        if isinstance(training_failed_error, ControllerError):
+            error_source = "controller"
+        elif isinstance(training_failed_error, WorkerGroupError):
+            error_source = "worker group"
+        else:
+            raise ValueError(f"Unknown error type: {type(training_failed_error)}")
+
+        logger.info(
+            f"[FailurePolicy] {decision.value}\n"
+            f"  Source: {error_source}\n"
+            f"  Error count: {error_count} (max allowed: {retry_limit})\n"
+            f"Error: {training_failed_error}",
+            exc_info=(
+                type(training_failed_error),
+                training_failed_error,
+                training_failed_error.__traceback__,
+            ),
+        )
+
+    def _is_retryable_error(self, training_failed_error: TrainingFailedError) -> bool:
+        if isinstance(training_failed_error, WorkerGroupError):
+            return True
+        elif isinstance(training_failed_error, ControllerError):
+            return isinstance(
+                training_failed_error.controller_failure, RETRYABLE_CONTROLLER_ERRORS
+            )
+        return False
+
+    def make_decision(
+        self,
+        training_failed_error: TrainingFailedError,
+    ) -> FailureDecision:
+
+        if not self._is_retryable_error(training_failed_error):
+            decision = FailureDecision.RAISE
+            error_count = 1
+            retry_limit = 0
+        else:
+            if isinstance(training_failed_error, ControllerError):
+                self._controller_failures += 1
+                error_count = self._controller_failures
+                retry_limit = (
+                    self.failure_config.controller_failure_limit
+                    if self.failure_config.controller_failure_limit != -1
+                    else float("inf")
+                )
+            elif isinstance(training_failed_error, WorkerGroupError):
+                self._worker_group_failures += 1
+                error_count = self._worker_group_failures
+                retry_limit = (
+                    self.failure_config.max_failures
+                    if self.failure_config.max_failures != -1
+                    else float("inf")
+                )
+            else:
+                raise ValueError(f"Unknown error type: {type(training_failed_error)}")
+
+            if error_count > retry_limit:
+                decision = FailureDecision.RAISE
+            else:
+                decision = FailureDecision.RETRY
+
+        self._log_decision(decision, training_failed_error, error_count, retry_limit)
+        return decision

python/ray/train/v2/_internal/execution/failure_handling/factory.py

@@ -0,0 +1,13 @@
+from ray.train import FailureConfig
+from ray.train.v2._internal.execution.failure_handling import (
+    DefaultFailurePolicy,
+    FailurePolicy,
+)
+
+
+def create_failure_policy(failure_config: FailureConfig) -> FailurePolicy:
+    """Create a failure policy from the given failure config.
+
+    Defaults to the `DefaultFailurePolicy` implementation.
+    """
+    return DefaultFailurePolicy(failure_config=failure_config)

python/ray/train/v2/_internal/execution/failure_handling/failure_policy.py

@@ -0,0 +1,29 @@
+import abc
+from enum import Enum
+
+from ray.train.v2.api.config import FailureConfig
+from ray.train.v2.api.exceptions import TrainingFailedError
+
+
+class FailureDecision(Enum):
+    RETRY = "RETRY"
+    RAISE = "RAISE"
+    NOOP = "NOOP"
+
+
+class FailurePolicy(abc.ABC):
+    """A policy that determines how to handle user and system failures.
+    FailurePolicy will handle the controller failure and worker errors during training.
+
+    This can be used to implement fault tolerance and error recovery.
+    """
+
+    def __init__(self, failure_config: FailureConfig):
+        self.failure_config = failure_config
+
+    @abc.abstractmethod
+    def make_decision(
+        self,
+        training_failed_error: TrainingFailedError,
+    ) -> FailureDecision:
+        raise NotImplementedError