fix(kubeflow): stream only rank 0 + last rank, write all ranks to disk

KubeflowExecutor.fetch_logs followed every replica and forwarded all ranks to
the caller, so at scale the aggregate output overran CI/runner job-log size
limits (a 16-node x 8-GPU run exceeded GitLab's 128MB cap). Now it still tails
every rank (kubectl logs -l <jobset> --prefix --max-log-requests num_nodes) and
writes the complete multi-rank output to <job_dir>/log-allranks_0.out, but
forwards only global rank 0 (node 0, [default0]) and the last global rank
(node num_nodes-1, [default nproc_per_node-1]) to stdout. Downstream log
validation that globs log*.out still sees every rank via the on-disk file.

Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

nemo_run/core/execution/kubeflow.py

@@ -16,6 +16,7 @@
 import getpass
 import logging
 import os
+import re
 import subprocess
 import time
 from dataclasses import dataclass, field
@@ -331,6 +332,13 @@ def fetch_logs(
         until pods are running (up to 10 minutes).  Otherwise it returns the last
         *lines* lines from a single ``kubectl logs`` call.
         """
+        # Tail every rank, but forward only global rank 0 and the last global
+        # rank to the caller (stdout / CI job log). Streaming all ranks at scale
+        # overruns CI/runner job-log size limits, yet the full multi-rank output
+        # is still written to <job_dir>/log-allranks_0.out so downstream log
+        # validation (which globs log*.out) sees every rank. --prefix tags each
+        # line with its pod name so we can recover the node (completion index)
+        # and pair it with torchrun's [defaultN] local-rank marker.
         label_selector = f"jobset.sigs.k8s.io/jobset-name={job_name}"
         cmd = [
             "kubectl",
@@ -339,11 +347,29 @@ def fetch_logs(
             label_selector,
             "-n",
             self.namespace,
+            "--prefix",
             "--tail",
             str(lines),
             "--max-log-requests",
             str(self.num_nodes),
         ]
+        last_node = max(self.num_nodes - 1, 0)
+        last_local = max(self.nproc_per_node() - 1, 0)
+        node_re = re.compile(r"node-0-(\d+)-")
+        local_re = re.compile(r"\[default(\d+)\]")
+
+        def _forward_to_stdout(log_line: str) -> bool:
+            """True for global rank 0 (node 0, local 0) and the last global rank."""
+            node_match = node_re.search(log_line)
+            local_match = local_re.search(log_line)
+            if not node_match or not local_match:
+                return False
+            node, local = int(node_match.group(1)), int(local_match.group(1))
+            return (node == 0 and local == 0) or (node == last_node and local == last_local)
+
+        all_ranks_path = os.path.join(self.job_dir, "log-allranks_0.out")
+        os.makedirs(self.job_dir, exist_ok=True)
+
         if stream:
             cmd.append("-f")
             # Retry kubectl logs -f until the job reaches a terminal state.
@@ -354,16 +380,21 @@ def fetch_logs(
                 )
                 lines_yielded = 0
                 try:
-                    for line in iter(proc.stdout.readline, ""):
-                        if line:
-                            lines_yielded += 1
-                            yield line
-                        if proc.poll() is not None:
-                            for remaining in proc.stdout:
-                                if remaining:
+                    with open(all_ranks_path, "a") as all_ranks_file:
+                        for line in iter(proc.stdout.readline, ""):
+                            if line:
+                                all_ranks_file.write(line)
+                                if _forward_to_stdout(line):
                                     lines_yielded += 1
-                                    yield remaining
-                            break
+                                    yield line
+                            if proc.poll() is not None:
+                                for remaining in proc.stdout:
+                                    if remaining:
+                                        all_ranks_file.write(remaining)
+                                        if _forward_to_stdout(remaining):
+                                            lines_yielded += 1
+                                            yield remaining
+                                break
                 except Exception as e:
                     logger.warning("Error streaming logs: %s; retrying", e)
                 finally:
@@ -381,7 +412,11 @@ def fetch_logs(
                 time.sleep(5)
         else:
             result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
-            yield from result.stdout.splitlines()
+            with open(all_ranks_path, "a") as all_ranks_file:
+                for line in result.stdout.splitlines():
+                    all_ranks_file.write(line + "\n")
+                    if _forward_to_stdout(line):
+                        yield line
 
     def cancel(
         self,