[SPARK-52843][PYTHON][TESTS] Support retry timeout tests

### What changes were proposed in this pull request?
Support retry timeout tests

### Why are the changes needed?
We have a batch of flaky (ml, streaming, etc) tests which are prone to timeout, existing `eventually` with `timeout` cannot handle this case

### Does this PR introduce _any_ user-facing change?
No, test-only

### How was this patch tested?
added ut

### Was this patch authored or co-authored using generative AI tooling?
no

Closes #51537 from zhengruifeng/retry_timeout_tests.

Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>

Author: zhengruifeng

python/pyspark/testing/utils.py

@@ -132,15 +132,15 @@ def write_int(i):
     return struct.pack("!i", i)
 
 
-def timeout(seconds):
+def timeout(timeout):
     def decorator(func):
         def handler(signum, frame):
-            raise TimeoutError(f"Function {func.__name__} timed out after {seconds} seconds")
+            raise TimeoutError(f"Function {func.__name__} timed out after {timeout} seconds")
 
         def wrapper(*args, **kwargs):
             signal.alarm(0)
             signal.signal(signal.SIGALRM, handler)
-            signal.alarm(seconds)
+            signal.alarm(timeout)
             try:
                 result = func(*args, **kwargs)
             finally:
@@ -155,6 +155,7 @@ def wrapper(*args, **kwargs):
 def eventually(
     timeout=30.0,
     catch_assertions=False,
+    catch_timeout=False,
 ):
     """
     Wait a given amount of time for a condition to pass, else fail with an error.
@@ -176,9 +177,14 @@ def eventually(
         If False (default), do not catch AssertionErrors.
         If True, catch AssertionErrors; continue, but save
         error to throw upon timeout.
+    catch_timeout : bool
+        If False (default), do not catch TimeoutError.
+        If True, catch TimeoutError; continue, but save
+        error to throw upon timeout.
     """
     assert timeout > 0
     assert isinstance(catch_assertions, bool)
+    assert isinstance(catch_timeout, bool)
 
     def decorator(condition: Callable) -> Callable:
         assert isinstance(condition, Callable)
@@ -191,21 +197,26 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
             while time() - start_time < timeout:
                 numTries += 1
 
-                if catch_assertions:
-                    try:
-                        lastValue = condition(*args, **kwargs)
-                    except AssertionError as e:
-                        lastValue = e
-                else:
+                try:
                     lastValue = condition(*args, **kwargs)
+                except AssertionError as e:
+                    if catch_assertions:
+                        lastValue = e
+                    else:
+                        raise e
+                except TimeoutError as e:
+                    if catch_timeout:
+                        lastValue = e
+                    else:
+                        raise e
 
                 if lastValue is True or lastValue is None:
                     return
 
                 print(f"\nAttempt #{numTries} failed!\n{lastValue}")
                 sleep(0.01)
 
-            if isinstance(lastValue, AssertionError):
+            if isinstance(lastValue, (AssertionError, TimeoutError)):
                 raise lastValue
             else:
                 raise AssertionError(

python/pyspark/tests/test_util.py

@@ -148,6 +148,17 @@ def test_parse_memory(self):
         with self.assertRaisesRegex(ValueError, "invalid format"):
             _parse_memory("2gs")
 
+    @eventually(timeout=180, catch_timeout=True)
+    @timeout(timeout=1)
+    def test_retry_timeout_test(self):
+        import random
+
+        if random.random() < 0.5:
+            print("hanging for 1 hour")
+            time.sleep(3600)  # Simulate a long-running operation
+        else:
+            print("succeeding immediately")
+
 
 class HandleWorkerExceptionTests(unittest.TestCase):
     exception_bytes = b"ValueError: test_message"