fix: make fetch_logs streaming resilient to flaky kubectl exits

Replace the brittle `lines_yielded > 0` and 10-minute deadline heuristics
with `status()`-based termination: the retry loop now runs until the job
reaches SUCCEEDED or FAILED, handling slow container pulls, mid-stream
crashes, and transient network failures correctly.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

nemo_run/core/execution/kubeflow.py

@@ -342,10 +342,9 @@ def fetch_logs(
         ]
         if stream:
             cmd.append("-f")
-            # Pods may not be running yet when the log thread starts.  Retry
-            # kubectl logs -f until we get output (or 10 minutes pass).
-            deadline = time.time() + 600
-            while time.time() < deadline:
+            # Retry kubectl logs -f until the job reaches a terminal state.
+            # This handles both pods not yet running and transient mid-stream failures.
+            while True:
                 proc = subprocess.Popen(
                     cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, bufsize=1
                 )
@@ -362,14 +361,20 @@ def fetch_logs(
                                     yield remaining
                             break
                 except Exception as e:
-                    logger.error("Error streaming logs: %s", e)
-                    break
+                    logger.warning("Error streaming logs: %s; retrying", e)
                 finally:
                     proc.terminate()
                     proc.wait(timeout=2)
-                if lines_yielded > 0:
-                    break  # kubectl exited after producing output — job done
-                time.sleep(5)  # no pods running yet, retry
+                state = self.status(job_name)
+                if state in (KubeflowJobState.SUCCEEDED, KubeflowJobState.FAILED):
+                    break  # job reached a terminal state, stop streaming
+                logger.warning(
+                    "kubectl logs exited (rc=%d, lines=%d, state=%s); retrying",
+                    proc.returncode,
+                    lines_yielded,
+                    state,
+                )
+                time.sleep(5)
         else:
             result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
             yield from result.stdout.splitlines()

test/core/execution/test_kubeflow.py

@@ -588,45 +588,51 @@ def test_launch_wait_exits_on_failed(self, executor, mock_k8s_clients):
             _, state = executor.launch("test-job", ["echo"], wait=True, timeout=30)
         assert state == KubeflowJobState.FAILED
 
-    # ── fetch_logs streaming: retry when no lines yielded ────────────────────
+    # ── fetch_logs streaming: retry until terminal state ─────────────────────
 
-    def test_fetch_logs_stream_retries_when_no_output_then_succeeds(
+    def test_fetch_logs_stream_retries_until_terminal_state(
         self, executor, mock_k8s_clients
     ):
-        """First Popen yields nothing; second yields a line — loop exits after output."""
+        """First Popen yields nothing and job is RUNNING; second yields a line and job is
+        SUCCEEDED — loop exits on terminal status."""
         import io
 
         empty_proc = MagicMock()
         empty_proc.stdout = io.StringIO("")
         empty_proc.poll.return_value = None
+        empty_proc.returncode = 1
 
         output_proc = MagicMock()
         output_proc.stdout = io.StringIO("some output\n")
         output_proc.poll.return_value = None
-
-        procs = [empty_proc, output_proc]
+        output_proc.returncode = 0
 
         with (
-            patch("subprocess.Popen", side_effect=procs),
+            patch("subprocess.Popen", side_effect=[empty_proc, output_proc]),
             patch("time.sleep"),
+            patch.object(
+                executor,
+                "status",
+                side_effect=[KubeflowJobState.RUNNING, KubeflowJobState.SUCCEEDED],
+            ),
         ):
             lines = list(executor.fetch_logs("my-job", stream=True))
 
         assert "some output\n" in lines
 
     def test_fetch_logs_stream_handles_exception(self, executor, mock_k8s_clients):
-        """Exception inside the readline loop is caught; generator terminates cleanly."""
+        """Exception inside the readline loop is caught; loop exits when job is terminal."""
 
         mock_proc = MagicMock()
-
         mock_proc.stdout.readline.side_effect = OSError("read error")
         mock_proc.poll.return_value = None
+        mock_proc.returncode = 1
 
         with (
             patch("subprocess.Popen", return_value=mock_proc),
             patch("time.sleep"),
+            patch.object(executor, "status", return_value=KubeflowJobState.FAILED),
         ):
-            # Should not raise; returns empty (error path)
             lines = list(executor.fetch_logs("my-job", stream=True))
 
         assert lines == []