chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_failure_policy.py

@@ -0,0 +1,103 @@
+import pytest
+
+from ray.train import FailureConfig
+from ray.train.v2._internal.exceptions import WorkerGroupStartupTimeoutError
+from ray.train.v2._internal.execution.failure_handling import (
+    FailureDecision,
+    create_failure_policy,
+)
+from ray.train.v2.api.exceptions import ControllerError, WorkerGroupError
+
+
+def _controller_error(retryable):
+    return ControllerError(
+        controller_failure=WorkerGroupStartupTimeoutError(0)
+        if retryable
+        else Exception("Non-retryable error")
+    )
+
+
+def _worker_group_error_from_errors(errors):
+    return WorkerGroupError(
+        "Worker group failed",
+        dict(enumerate(errors)),
+    )
+
+
+@pytest.mark.parametrize("max_failures", [0, 1, 10])
+def test_max_failures(max_failures):
+    policy = create_failure_policy(FailureConfig(max_failures=max_failures))
+
+    for _ in range(max_failures):
+        assert (
+            policy.make_decision(
+                training_failed_error=_worker_group_error_from_errors(
+                    [RuntimeError(f"Worker {i} failed") for i in range(8)]
+                )
+            )
+            == FailureDecision.RETRY
+        )
+    assert (
+        policy.make_decision(
+            training_failed_error=_worker_group_error_from_errors(
+                [RuntimeError(f"Worker {i} failed") for i in range(8)]
+            )
+        )
+        == FailureDecision.RAISE
+    )
+
+
+@pytest.mark.parametrize("controller_failure_limit", [0, 1, 10])
+def test_max_controller_failures(controller_failure_limit):
+    policy = create_failure_policy(
+        FailureConfig(controller_failure_limit=controller_failure_limit)
+    )
+    controller_error = _controller_error(retryable=True)
+    for _ in range(controller_failure_limit):
+        assert (
+            policy.make_decision(training_failed_error=controller_error)
+            == FailureDecision.RETRY
+        )
+    assert (
+        policy.make_decision(training_failed_error=controller_error)
+        == FailureDecision.RAISE
+    )
+
+
+def test_infinite_retry():
+    policy = create_failure_policy(FailureConfig(max_failures=-1))
+    for _ in range(10):
+        assert (
+            policy.make_decision(
+                training_failed_error=WorkerGroupError(
+                    "Worker group resize failed",
+                    {0: WorkerGroupStartupTimeoutError(0)},
+                )
+            )
+            == FailureDecision.RETRY
+        )
+
+
+def test_non_retryable_error():
+    policy = create_failure_policy(FailureConfig(controller_failure_limit=10))
+    controller_error = _controller_error(retryable=False)
+    assert (
+        policy.make_decision(training_failed_error=controller_error)
+        == FailureDecision.RAISE
+    )
+
+
+def test_infinite_controller_failure_retry():
+    policy = create_failure_policy(FailureConfig(controller_failure_limit=-1))
+    controller_error = _controller_error(retryable=True)
+    for _ in range(10):
+        assert (
+            policy.make_decision(training_failed_error=controller_error)
+            == FailureDecision.RETRY
+        )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))