From aadd2bb3d442e3d35275675d12329a2e0fae3be1 Mon Sep 17 00:00:00 2001
From: Nikolaus Schuetz <nikolauspschuetz@gmail.com>
Date: Thu, 9 Jul 2026 00:09:03 -0700
Subject: [PATCH] test: Retry CliRunner apply subprocess on timeout to reduce
 flakiness
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CliRunner.run() (used for 'feast apply' in unit tests) spawns a fresh child
interpreter that must import Feast (pandas/pyarrow/...) from cold. On loaded
CI runners — observed on unit-test-python (3.11, macos-14) — this can exceed
the 120s timeout, failing the test with:

    AssertionError: Command timed out after 120s: ['apply']

Unlike the materializing paths guarded by run_with_output(), commands routed
through run() do not materialize and so cannot trigger the Dask atexit hang the
timeout was added for. A timeout here is therefore almost always transient cold
-start slowness, not a genuine deadlock, so retrying once (with warm OS/page
caches) clears it while a real repeated stall still fails after the attempts
are exhausted.

Builds on the CliRunner subprocess stabilization in #6560.

Signed-off-by: Nikolaus Schuetz <nikolauspschuetz@gmail.com>
---
 sdk/python/tests/utils/cli_repo_creator.py | 49 +++++++++++++++-------
 1 file changed, 33 insertions(+), 16 deletions(-)

diff --git a/sdk/python/tests/utils/cli_repo_creator.py b/sdk/python/tests/utils/cli_repo_creator.py
index 03c0976e880..48491610174 100644
--- a/sdk/python/tests/utils/cli_repo_creator.py
+++ b/sdk/python/tests/utils/cli_repo_creator.py
@@ -63,26 +63,43 @@ def _subprocess_env(self) -> dict[str, str]:
         env["PYTHONPATH"] = os.pathsep.join(parent_paths)
         return env
 
-    def run(self, args: List[str], cwd: Path) -> subprocess.CompletedProcess:
+    def run(
+        self, args: List[str], cwd: Path, attempts: int = 2
+    ) -> subprocess.CompletedProcess:
         # Apply a conservative timeout to prevent CI hangs from Dask atexit-handler
         # stalls or other subprocess blockages.
         timeout = 120
 
-        try:
-            return subprocess.run(
-                [sys.executable, cli.__file__] + args,
-                cwd=cwd,
-                capture_output=True,
-                timeout=timeout,
-                env=self._subprocess_env(),
-            )
-        except subprocess.TimeoutExpired:
-            return subprocess.CompletedProcess(
-                args=[sys.executable, cli.__file__] + args,
-                returncode=-1,
-                stdout=b"",
-                stderr=f"Command timed out after {timeout}s: {args}".encode(),
-            )
+        # Commands routed through here (e.g. `apply`) do not materialize, so they
+        # cannot trigger the Dask atexit hang described in the module docstring. A
+        # timeout is therefore almost always transient slowness — a cold child
+        # interpreter importing Feast (pandas/pyarrow/...) on a loaded CI runner,
+        # which has been observed to exceed the timeout on the slower macOS
+        # runners. Retry once so a single slow cold start does not flake the test:
+        # the retry runs with warm OS/page caches and typically finishes quickly,
+        # while a genuine repeated stall still fails once attempts are exhausted.
+        full_args = [sys.executable, cli.__file__] + args
+        result = None
+        for attempt in range(1, attempts + 1):
+            try:
+                return subprocess.run(
+                    full_args,
+                    cwd=cwd,
+                    capture_output=True,
+                    timeout=timeout,
+                    env=self._subprocess_env(),
+                )
+            except subprocess.TimeoutExpired:
+                result = subprocess.CompletedProcess(
+                    args=full_args,
+                    returncode=-1,
+                    stdout=b"",
+                    stderr=(
+                        f"Command timed out after {timeout}s "
+                        f"(attempt {attempt}/{attempts}): {args}"
+                    ).encode(),
+                )
+        return result
 
     def run_with_output(self, args: List[str], cwd: Path) -> Tuple[int, bytes]:
         is_teardown = "teardown" in args