fix(kubeflow): unique TrainJob name + forward all ranks (deduped)

- TrainJob name is now <basename>-<uuid6> (RFC-1123, <=33 chars) via the new
  train_job_basename field, decoupled from the experiment name. The uuid makes
  every launch unique, so concurrent/retried jobs never collide on the API
  server (the descriptive experiment name is intentionally non-unique).
- fetch_logs now forwards every rank to stdout, de-duplicated: torchrun runs
  the same entrypoint on all ranks so startup/config/NCCL lines are identical;
  we strip the per-rank [pod/...]/[defaultN] markers and forward each distinct
  message once. This stops dropping the per-step loss line and wandb URL, which
  Megatron emits from a single layout-dependent rank (neither rank 0 nor last).
  The full per-rank stream still goes to log-allranks_0.out untouched.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

nemo_run/core/execution/kubeflow.py

@@ -14,12 +14,12 @@
 # limitations under the License.
 
 import getpass
-import json
 import logging
 import os
 import re
 import subprocess
 import time
+import uuid
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Iterable, Optional
@@ -110,6 +110,10 @@ class KubeflowExecutor(Executor):
     # the PVC sync.  Use this to include local scripts/files that are not
     # generated by the packager (e.g. a hand-written training script).
     workdir_local_path: Optional[str] = None
+    # Human-readable base for the generated TrainJob name. The k8s name becomes
+    # ``<basename>-<uuid6>`` (RFC-1123 safe, ≤33 chars); the uuid keeps every
+    # launch unique. Falls back to the launch ``name`` when unset.
+    train_job_basename: Optional[str] = None
 
     def __post_init__(self):
         if not _KUBERNETES_AVAILABLE:
@@ -281,6 +285,26 @@ def get_job_body(self, name: str, command: list[str]) -> dict:
 
     # ── Submit / status / cancel / logs ──────────────────────────────────────
 
+    def _trainjob_name(self, fallback: str) -> str:
+        """RFC-1123 base name ``<basename>-<uuid6>`` (≤33 chars), generated once.
+
+        Shared by the TrainJob and its data-mover pod (created in ``package()``,
+        before ``launch()``) so both are valid, unique per launch — the uuid
+        avoids API-server collisions — and consistent. The basename is
+        ``train_job_basename`` (e.g. the model recipe) or the caller's name,
+        sanitized to lowercase alphanumerics + dashes; capped at 33 chars to
+        stay under the 63-char label limit with room for the ``-data-mover``
+        suffix.
+        """
+        cached = getattr(self, "_k8s_job_name", None)
+        if cached is not None:
+            return cached
+        base = re.sub(r"[^a-z0-9-]+", "-", (self.train_job_basename or fallback or "job").lower()).strip("-")
+        uid = uuid.uuid4().hex[:6]
+        base = base[: 33 - len(uid) - 1].strip("-") or "job"
+        self._k8s_job_name = f"{base}-{uid}"
+        return self._k8s_job_name
+
     def launch(
         self,
         name: str,
@@ -295,7 +319,7 @@ def launch(
         observed ``RUNNING``, ``SUCCEEDED``, or ``FAILED`` state when *wait* is ``True``.
         Raises ``RuntimeError`` if the job already exists or *timeout* expires.
         """
-        name = name.replace("_", "-").replace(".", "-").lower()
+        name = self._trainjob_name(name)
         job_body = self.get_job_body(name, cmd)
         try:
             self._custom_objects_api.create_namespaced_custom_object(
@@ -413,21 +437,17 @@ def fetch_logs(
         *lines* lines from a single ``kubectl logs`` call.
         """
         # Tail every rank to <job_dir>/log-allranks_0.out (downstream log
-        # validation globs log*.out and needs 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 job-log limits.
-        #
-        # Identifying the last global rank requires the authoritative node rank,
-        # NOT the pod name. Kubeflow Trainer binds torchrun's PET_NODE_RANK to
-        # the indexed-Job completion index, stamped on each pod as the
-        # `batch.kubernetes.io/job-completion-index` label. So:
-        #     global_rank = job_completion_index * nproc_per_node + local_rank
-        # `--prefix` tags each line with `[pod/<pod>/<container>]`; we map that
-        # pod name → completion index (refreshed on every (re)connect, since a
-        # gang restart spawns new pod names) and pair it with torchrun's
-        # `[defaultN]` local-rank marker. `--tail=-1` replays each pod's full
-        # history on (re)attach so mid-run lines are never dropped (the previous
-        # `--tail <lines>` snapshot missed the last rank's per-step lines).
+        # validation globs log*.out and needs every rank). Forward all ranks to
+        # the caller (stdout / CI job log) too, but de-duplicated: torchrun runs
+        # the same entrypoint on every rank, so the bulk of the volume (startup,
+        # config dump, NCCL init) is byte-identical across ranks. We forward each
+        # distinct message once — which is also why the rank-specific loss line
+        # (emitted by a single, parallelism-layout-dependent rank that is usually
+        # neither rank 0 nor the last rank) and genuine per-rank errors are no
+        # longer dropped. `--prefix` tags each line with `[pod/<pod>/<container>]`
+        # and torchrun adds `[defaultN]`; both are stripped to form the dedup key.
+        # `--tail=-1` replays each pod's full history on (re)attach so mid-run
+        # lines are never dropped.
         label_selector = f"jobset.sigs.k8s.io/jobset-name={job_name}"
         cmd = [
             "kubectl",
@@ -442,61 +462,40 @@ def fetch_logs(
             "--max-log-requests",
             str(self.num_nodes),
         ]
-        nproc = self.nproc_per_node()
-        last_rank = max(self.num_nodes * nproc - 1, 0)
-        pod_re = re.compile(r"pod/([^/]+)/")
-        local_re = re.compile(r"\[default(\d+)\]")
-
-        def _pod_index_map() -> dict[str, int]:
-            """Map pod name → job-completion-index (== torchrun node rank)."""
-            try:
-                out = subprocess.run(
-                    [
-                        "kubectl",
-                        "get",
-                        "pods",
-                        "-n",
-                        self.namespace,
-                        "-l",
-                        label_selector,
-                        "-o",
-                        "json",
-                    ],
-                    capture_output=True,
-                    text=True,
-                    timeout=timeout,
-                )
-                items = json.loads(out.stdout).get("items", [])
-            except Exception as e:
-                logger.warning("Could not list pods for %s: %s", job_name, e)
-                return {}
-            mapping: dict[str, int] = {}
-            for item in items:
-                meta = item.get("metadata", {})
-                name = meta.get("name")
-                idx = (meta.get("labels", {}) or {}).get("batch.kubernetes.io/job-completion-index")
-                if name is not None and idx is not None:
-                    mapping[name] = int(idx)
-            return mapping
-
-        def _forward_to_stdout(log_line: str, pod_index: dict[str, int]) -> bool:
-            """True for the first and last *global rank* only.
-
-            Kubeflow Trainer sets torchrun's PET_NODE_RANK from the JobSet
-            completion-index label (static), so the global rank is
-            ``node_rank * nproc_per_node + local_rank`` where node_rank is the
-            pod's completion index and local_rank is torchrun's ``[defaultN]``
-            marker. We forward global rank 0 and ``world_size - 1`` only.
-            """
-            pod_match = pod_re.search(log_line)
-            local_match = local_re.search(log_line)
-            if not pod_match or not local_match:
-                return False
-            node = pod_index.get(pod_match.group(1))
-            if node is None:
+        # Collapse the near-simultaneous cross-rank burst with a *sliding time
+        # window* (cf. ClusterShell `clush -b`, which gathers identical output
+        # across nodes into one line). torchrun runs the same entrypoint on
+        # every rank, so startup/config/NCCL lines arrive as a burst of
+        # byte-identical copies; we strip the per-rank `[pod/<pod>/<container>]`
+        # and `[defaultN]` markers to form a dedup key and suppress a key only
+        # if an identical line was already forwarded within `dedup_window_s`.
+        # Unlike a global set this is bounded in both memory and time: a line
+        # that legitimately recurs later (e.g. a periodic "saving checkpoint")
+        # is forwarded again once the window passes, and a continuously
+        # repeating line is rate-limited to once per window rather than
+        # suppressed for the whole run. Lines whose body differs across ranks
+        # (per-step loss, `[rankN]` errors) keep distinct keys and are never
+        # collapsed. The full per-rank stream still goes to log-allranks_0.out.
+        rank_marker_re = re.compile(r"\[pod/[^\]]+\]\s*|\[default\d+\]:?\s*")
+        dedup_window_s = 60.0
+        last_forwarded: dict[str, float] = {}
+
+        def _should_forward(log_line: str) -> bool:
+            key = rank_marker_re.sub("", log_line).strip()
+            if not key:
+                return True
+            now = time.monotonic()
+            prev = last_forwarded.get(key)
+            if prev is not None and now - prev < dedup_window_s:
                 return False
-            global_rank = node * nproc + int(local_match.group(1))
-            return global_rank == 0 or global_rank == last_rank
+            last_forwarded[key] = now
+            # Bound memory: once the map is large, drop keys older than the
+            # window (they can no longer suppress anything).
+            if len(last_forwarded) > 20000:
+                stale = now - dedup_window_s
+                for k in [k for k, t in last_forwarded.items() if t < stale]:
+                    del last_forwarded[k]
+            return True
 
         all_ranks_path = os.path.join(self.job_dir, "log-allranks_0.out")
         os.makedirs(self.job_dir, exist_ok=True)
@@ -506,7 +505,6 @@ def _forward_to_stdout(log_line: str, pod_index: dict[str, int]) -> bool:
             # 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:
-                pod_index = _pod_index_map()
                 proc = subprocess.Popen(
                     cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, bufsize=1
                 )
@@ -516,14 +514,14 @@ def _forward_to_stdout(log_line: str, pod_index: dict[str, int]) -> bool:
                         for line in iter(proc.stdout.readline, ""):
                             if line:
                                 all_ranks_file.write(line)
-                                if _forward_to_stdout(line, pod_index):
+                                if _should_forward(line):
                                     lines_yielded += 1
                                     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, pod_index):
+                                        if _should_forward(remaining):
                                             lines_yielded += 1
                                             yield remaining
                                 break
@@ -543,12 +541,11 @@ def _forward_to_stdout(log_line: str, pod_index: dict[str, int]) -> bool:
                 )
                 time.sleep(5)
         else:
-            pod_index = _pod_index_map()
             result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
             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, pod_index):
+                    if _should_forward(line):
                         yield line
 
     def cancel(
@@ -614,7 +611,7 @@ def cancel(
     # ── Workdir sync helpers ──────────────────────────────────────────────────
 
     def _data_mover_pod_name(self, job_name: str) -> str:
-        return f"{job_name}-data-mover"
+        return f"{self._trainjob_name(job_name)}-data-mover"
 
     def _start_data_mover_pod(self, pod_name: str, timeout: int = 120) -> None:
         """Spin up a throw-away Alpine pod that mounts workdir_pvc and blocks until Running.