From f43abfb547b991590980a7aa4069c46955c4a00f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Sun, 31 May 2026 09:36:48 +0000
Subject: [PATCH] fix(ci): bound JET pipeline polling with a watchdog to
 prevent indefinite hangs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The release driver blocked on `pipeline.wait(max_wait_time=7d, retries_on_error=3)`.
jetclient's `_wait_for_state` only re-checks `max_wait_time` between polls and
`retries_on_error` only fires on a *raised* error, so when a status poll's HTTP
GET wedges on a silently-dropped TCP connection (no socket read timeout), the
wait never returns. A weekly release job hung ~28h after its downstream had
already finished and had to be cancelled by hand.

Replace the single blocking wait with `wait_for_pipeline_completion`, which
drives the status poll directly and guards each individual poll with a SIGALRM
watchdog: a hung connection is interrupted and retried, while a 12h wall-clock
deadline bounds the total wait. On deadline the downstream is cancelled so the
existing iteration-retry logic kicks in.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 .../python_scripts/launch_jet_workload.py     | 87 ++++++++++++++++++-
 1 file changed, 85 insertions(+), 2 deletions(-)

diff --git a/tests/test_utils/python_scripts/launch_jet_workload.py b/tests/test_utils/python_scripts/launch_jet_workload.py
index d5629bd432d..5365ebdc118 100644
--- a/tests/test_utils/python_scripts/launch_jet_workload.py
+++ b/tests/test_utils/python_scripts/launch_jet_workload.py
@@ -82,6 +82,79 @@ def sigterm_handler(_signo, _stack_frame):
     signal.signal(signal.SIGTERM, sigterm_handler)
 
 
+TERMINAL_PIPELINE_STATUSES = frozenset(
+    {
+        PipelineStatus.SUCCESS,
+        PipelineStatus.FAILED,
+        PipelineStatus.CANCELED,
+        PipelineStatus.SUBMISSION_FAILED,
+    }
+)
+
+
+class StatusPollTimeout(Exception):
+    """Raised by the watchdog when a single pipeline status poll exceeds its budget."""
+
+
+def wait_for_pipeline_completion(
+    pipeline: jetclient.JETPipeline,
+    max_wait_time: int = 60 * 60 * 12,
+    interval: int = 60,
+    poll_timeout: int = 60 * 3,
+) -> PipelineStatus:
+    """Block until a downstream JET pipeline reaches a terminal status.
+
+    jetclient's own ``pipeline.wait`` polls the GitLab API with no socket read
+    timeout and only re-checks ``max_wait_time`` between polls, so a silently
+    dropped TCP connection wedges ``recv`` forever and the wait never returns
+    (observed: a release job hung ~28h after its downstream had already
+    finished). This drives the status poll directly and guards each individual
+    poll with a SIGALRM watchdog, so a hung connection is interrupted and
+    retried, while a wall-clock deadline bounds the total wait.
+
+    Args:
+        pipeline: The submitted downstream JET pipeline to watch.
+        max_wait_time: Wall-clock budget in seconds before giving up.
+        interval: Seconds to sleep between status polls.
+        poll_timeout: Per-poll watchdog budget in seconds.
+
+    Returns:
+        The terminal status reached by the pipeline.
+
+    Raises:
+        jetclient.facades.objects.util.WaitTimeExceeded: If the pipeline does
+            not reach a terminal status within ``max_wait_time`` seconds.
+    """
+
+    def raise_poll_timeout(_signo, _stack_frame):
+        raise StatusPollTimeout
+
+    deadline = time.monotonic() + max_wait_time
+    previous_handler = signal.signal(signal.SIGALRM, raise_poll_timeout)
+    try:
+        while time.monotonic() < deadline:
+            signal.alarm(poll_timeout)
+            try:
+                status = pipeline.get_status()
+            except (StatusPollTimeout, jetclient.clients.gitlab.GitlabAPIError) as e:
+                logger.warning("Pipeline status poll failed (%s); retrying", type(e).__name__)
+                status = None
+            finally:
+                signal.alarm(0)
+
+            if status in TERMINAL_PIPELINE_STATUSES:
+                return status
+
+            time.sleep(interval)
+
+        raise jetclient.facades.objects.util.WaitTimeExceeded(
+            f"Pipeline {pipeline.jet_id} did not reach a terminal status "
+            f"within {max_wait_time} seconds"
+        )
+    finally:
+        signal.signal(signal.SIGALRM, previous_handler)
+
+
 def launch_and_wait_for_completion(
     test_case: str,
     environment: str,
@@ -187,9 +260,19 @@ def launch_and_wait_for_completion(
         pipeline.jet_id,
     )
 
-    pipeline.wait(max_wait_time=60 * 60 * 24 * 7, interval=60 * 1, retries_on_error=3)
+    try:
+        status = wait_for_pipeline_completion(pipeline)
+        logger.info(f"Pipeline terminated; status: {status}")
+    except jetclient.facades.objects.util.WaitTimeExceeded:
+        logger.error(
+            "Pipeline %s exceeded the wall-clock budget; cancelling so the iteration retries.",
+            pipeline.jet_id,
+        )
+        try:
+            pipeline.cancel()
+        except jetclient.clients.gitlab.GitlabAPIError:
+            logger.exception("Failed to cancel pipeline %s", pipeline.jet_id)
 
-    logger.info(f"Pipeline terminated; status: {pipeline.get_status()}")
     return pipeline