fix: clear stale error.pb at the start of every clustered attempt

examples/clustered/ddp_train_restart.py

@@ -1,18 +1,28 @@
 """
-DDP training that FAILS on attempt 0 and SUCCEEDS on the JobSet restart.
-
-This is a regression test for the clustered-restart fix on this branch: when a
-JobSet restarts (``JOBSET_RESTART_ATTEMPT`` > 0), ``upload_outputs`` must first
-delete the stale ``error.pb`` written by the previous failed attempt. Otherwise
-the successful retry's outputs land alongside a leftover error file and the
-execution is still reported as FAILED.
-
-Mechanics:
-    - ``ClusterFailurePolicy(max_restarts=1)`` lets the JobSet restart once.
-    - Attempt 0 (``JOBSET_RESTART_ATTEMPT`` unset / "0") raises -> writes error.pb.
-    - Attempt 1 (``JOBSET_RESTART_ATTEMPT`` == "1") runs DDP and uploads outputs.
-      With the fix, the stale error.pb is cleared and the run ends SUCCEEDED.
-      Without the fix, the run ends FAILED despite the successful retry.
+DDP training that FAILS on the first torchrun attempt and SUCCEEDS on the in-pod restart.
+
+Regression test for the stale ``error.pb`` cleanup in the clustered runtime
+(``flyte._internal.runtime.io.clear_stale_clustered_error``): whenever a clustered rank-0
+worker starts, it removes any ``error.pb`` an earlier restart left under the attempt's output
+prefix, before the task body runs. Without that cleanup a successful restart uploads
+``outputs.pb`` next to the leftover ``error.pb`` and the executor still reports the run FAILED.
+
+The stale file is produced here without any backend help:
+    - ``ClusterFailurePolicy(max_restarts=0)`` makes every attempt look terminal to the SDK's
+      terminal-attempt gate (``JOBSET_RESTART_ATTEMPT 0 >= JOBSET_MAX_RESTARTS 0``), so the first
+      failure writes ``error.pb`` right away and the worker exits 1.
+    - ``PET_MAX_RESTARTS=1`` lets torchrun restart the worker group once inside the same pod.
+      (``TorchRun(max_restarts=...)`` is not wired through to torchrun yet, so the env var is set
+      directly; torchrun reads ``PET_<FLAG>`` for every CLI flag.)
+    - Attempt 0 (``TORCHELASTIC_RESTART_COUNT`` == "0") raises on every rank.
+    - Attempt 1 (``TORCHELASTIC_RESTART_COUNT`` == "1") trains and uploads outputs. Rank-0 logs
+      "Removed stale ... error.pb" at startup.
+
+Expected: run phase SUCCEEDED. Without the cleanup: FAILED with the attempt-0 error.
+
+A JobSet-level restart (``ClusterFailurePolicy(max_restarts >= 1)``) goes through the same cleanup,
+but without free host-maintenance restarts the gate never writes a premature ``error.pb``, so that
+variant passes with or without the fix and is not a useful regression test.
 
 Run:
     uv run python examples/clustered/ddp_train_restart.py
@@ -33,14 +43,14 @@
 )
 
 # --- Knobs ---------------------------------------------------------------------------------------
-USE_GPU = True
-REPLICAS = 2  # pods (== nodes)
-NPROC_PER_NODE = 1  # processes (one per GPU) per pod  => world_size = REPLICAS * NPROC_PER_NODE
+USE_GPU = False
+REPLICAS = 1  # one pod: torchrun's in-pod restart then needs no cross-node re-rendezvous
+NPROC_PER_NODE = 2  # processes per pod  => world_size = REPLICAS * NPROC_PER_NODE
 
 _BACKEND = "nccl" if USE_GPU else "gloo"
 
 resources = (
-    flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:1")
+    flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:2")  # one GPU per process (NPROC_PER_NODE)
     if USE_GPU
     else flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi"))
 )
@@ -51,21 +61,28 @@
     resources=resources,
     replicas=REPLICAS,
     nproc_per_node=NPROC_PER_NODE,
-    runtime=TorchRun(rdzv_backend="static", max_restarts=0),
-    failure_policy=ClusterFailurePolicy(max_restarts=1),  # allow ONE JobSet restart
+    # max_restarts here is not wired through to torchrun yet; PET_MAX_RESTARTS below is what works today.
+    runtime=TorchRun(rdzv_backend="static", max_restarts=1),
+    failure_policy=ClusterFailurePolicy(max_restarts=0),  # every attempt looks terminal to the SDK gate
+    env_vars={"PET_MAX_RESTARTS": "1"},  # ONE in-pod torchrun restart (see module docstring)
 )
 
 
 @env.task
 async def train_ddp_with_restart(steps: int = 50, lr: float = 0.05) -> float:
-    """Fail on the first JobSet attempt, then train + return loss on the restart."""
-    restart_attempt = int(os.environ.get("JOBSET_RESTART_ATTEMPT", "0") or "0")
+    """Fail on the first torchrun attempt, then train + return loss on the in-pod restart."""
+    restart_attempt = int(os.environ.get("TORCHELASTIC_RESTART_COUNT", "0") or "0")
     rank = os.environ.get("RANK", "0")
-    print(f"[rank {rank}] JOBSET_RESTART_ATTEMPT={restart_attempt}", flush=True)
+    print(
+        f"[rank {rank}] TORCHELASTIC_RESTART_COUNT={restart_attempt} "
+        f"JOBSET_RESTART_ATTEMPT={flyte.ctx().restart_attempt}",
+        flush=True,
+    )
 
-    # Attempt 0 fails on every worker -> writes error.pb for the execution.
+    # Attempt 0 fails on every rank -> rank-0 writes error.pb (the SDK gate sees 0 >= 0) and every
+    # worker exits 1, so torchrun restarts the worker group in-pod.
     if restart_attempt == 0:
-        raise RuntimeError("Intentional failure on attempt 0 to force a JobSet restart")
+        raise RuntimeError("Intentional failure on torchrun attempt 0 to leave a stale error.pb behind")
 
     import torch
     import torch.distributed as dist
@@ -120,5 +137,5 @@ async def train_ddp_with_restart(steps: int = 50, lr: float = 0.05) -> float:
     run = flyte.run(train_ddp_with_restart, steps=50)
     print("Run URL:", run.url)
     run.wait()
-    # Expected WITH the fix: SUCCEEDED. Without it: FAILED (stale error.pb).
+    # Expected WITH the cleanup: SUCCEEDED. Without it: FAILED (stale error.pb from attempt 0).
     print("Final phase:", run.phase)

src/flyte/_internal/runtime/entrypoints.py

@@ -269,6 +269,13 @@ async def load_and_run_task(
     """
     sw = Stopwatch("load_and_run_task_total")
     sw.start()
+    if output_path:
+        # Clustered rank-0 only; a no-op for regular tasks. Must precede every upload_error site (the load
+        # failure below, taskrunner.extract_download_run_upload, runtime._run_and_stop) and the task body.
+        # rusty.run_task bypasses this function but never hosts clustered workers (clustered.py execs a0).
+        from .io import clear_stale_clustered_error
+
+        await clear_stale_clustered_error(output_path)
     try:
         task = await _download_and_load_task(code_bundle, resolver, resolver_args)
     except Exception as e:

src/flyte/_internal/runtime/io.py

@@ -8,6 +8,7 @@
 
 from flyteidl2.core import execution_pb2
 from flyteidl2.task import common_pb2
+from fsspec.asyn import AsyncFileSystem
 
 import flyte.storage as storage
 from flyte._logging import logger
@@ -40,6 +41,9 @@ def _is_nonzero_rank_clustered_worker() -> bool:
 
 
 def _get_clustered_restart_attempt() -> int | None:
+    """JOBSET_RESTART_ATTEMPT mirrors JobSet `Status.Restarts`, which also counts free host-maintenance
+    restarts, so it can run ahead of the charged budget (`RestartsCountTowardsMax`, never exposed to pods).
+    """
     raw_attempt = os.environ.get("JOBSET_RESTART_ATTEMPT")
     if raw_attempt is None:
         return None
@@ -62,13 +66,18 @@ def _get_clustered_max_restarts() -> int | None:
 
 
 def _is_terminal_clustered_attempt() -> bool:
-    """Whether a failure in this attempt should write error.pb.
+    """Best-effort guess whether a failure in this attempt should write error.pb.
 
     For a clustered/jobset task the JobSet restarts the whole pod set up to `max_restarts` times
-    within a single Flyte attempt. We only write error.pb on the terminal attempt (budget exhausted)
-    so transient restarts don't leave a stale error that a later successful restart would have to
-    delete. Returns True (write) for non-clustered tasks, and as a safe fallback whenever the budget
-    is unknown — errors must never be silently hidden.
+    within a single Flyte attempt, so error.pb is only written once the budget looks exhausted, which
+    avoids needless writes on transient restarts. The guess can be EARLY but never late:
+    JOBSET_RESTART_ATTEMPT mirrors JobSet `Status.Restarts`, which also counts free host-maintenance
+    restarts (`restart_on_host_maintenance`), while the budget is charged from
+    `RestartsCountTowardsMax`, which pods cannot see. A premature error.pb is harmless because
+    `clear_stale_clustered_error` removes it at the start of the next attempt — that cleanup, not this
+    gate, is what keeps a later successful attempt from being reported as failed. Returns True (write)
+    for non-clustered tasks, and as a safe fallback whenever the budget is unknown — errors must never
+    be silently hidden.
     """
     attempt = _get_clustered_restart_attempt()
     if attempt is None:
@@ -79,6 +88,50 @@ def _is_terminal_clustered_attempt() -> bool:
     return attempt >= max_restarts
 
 
+async def _delete_path(path: str) -> None:
+    """Delete one object. Raises FileNotFoundError when it is already gone (GCS/Azure/local; S3 returns OK)."""
+    fs = storage.get_underlying_filesystem(path=path)
+    if isinstance(fs, AsyncFileSystem):
+        # The sync rm_file() of an AsyncFileSystem re-enters the running loop and fails from inside a
+        # coroutine; obstore's FsspecStore (s3/gs/abfs) implements the async _rm_file directly.
+        await fs._rm_file(path)  # pylint: disable=W0212
+        return
+    fs.rm_file(path)
+
+
+async def clear_stale_clustered_error(output_path: str) -> None:
+    """Remove an error.pb left under this attempt's output prefix by an earlier restart of the pod set.
+
+    Runs at startup of every clustered rank-0 worker (JobSet whole-set restarts and in-pod torchrun
+    restarts both re-exec `a0`). Deliberately unconditional rather than keyed on a restart counter:
+    JOBSET_RESTART_ATTEMPT mirrors JobSet `Status.Restarts`, which also counts free host-maintenance
+    restarts, so `_is_terminal_clustered_attempt` can write error.pb on a non-terminal attempt and a
+    later successful attempt would then be reported as failed (the executor reads error.pb before
+    outputs.pb). Restart attempts are strictly sequential, so there is no concurrent writer, and only
+    rank-0 ever writes or deletes the file. Never raises: a failed delete just leaves today's behavior.
+    """
+    if not _is_clustered_worker() or _is_nonzero_rank_clustered_worker():
+        return
+    error_uri = error_path(output_path)
+    restart_ctx = (
+        f"JOBSET_RESTART_ATTEMPT={os.environ.get('JOBSET_RESTART_ATTEMPT')}, "
+        f"TORCHELASTIC_RESTART_COUNT={os.environ.get('TORCHELASTIC_RESTART_COUNT')}"
+    )
+    try:
+        if not await storage.exists(error_uri):
+            return
+        await _delete_path(error_uri)
+    except FileNotFoundError:
+        logger.debug(f"Stale {error_uri} disappeared before it could be deleted ({restart_ctx})")
+        return
+    except Exception as e:
+        logger.warning(f"Could not remove stale {error_uri} ({restart_ctx}): {e}")
+        return
+    # Warning, not info: the default pod log level is WARNING, and this line is the operator's evidence
+    # that the terminal-attempt gate fired early. It fires at most once per restarted attempt.
+    logger.warning(f"Removed stale {error_uri} left by an earlier restart of this task ({restart_ctx})")
+
+
 def pkl_path(base_path: str, pkl_name: str) -> str:
     return storage.join(base_path, f"{pkl_name}{_PKL_EXT}")
 
@@ -148,11 +201,15 @@ async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str, re
     # so they don't race to clobber error.pb.
     if _is_nonzero_rank_clustered_worker():
         return error_uri
-    # For a clustered task, only write error.pb once the JobSet has exhausted its restart budget.
-    # Transient restarts recover on their own, so writing on every attempt would leave a stale
-    # error that a later successful restart would have to delete.
+    # For a clustered task, only write error.pb once the JobSet looks to have exhausted its restart
+    # budget. Transient restarts recover on their own; if this guess is early (free host-maintenance
+    # restarts inflate JOBSET_RESTART_ATTEMPT), clear_stale_clustered_error removes the file at the
+    # start of the next attempt.
     if not _is_terminal_clustered_attempt():
-        logger.info(f"Skipping error.pb on transient JobSet restart (budget remaining): {error_uri}")
+        logger.info(
+            f"Skipping error.pb on transient JobSet restart (budget remaining, "
+            f"attempt={_get_clustered_restart_attempt()} max_restarts={_get_clustered_max_restarts()}): {error_uri}"
+        )
         return error_uri
     error_document = execution_pb2.ErrorDocument(
         error=execution_pb2.ContainerError(

src/flyte/clustered/_environment.py

@@ -36,7 +36,8 @@ class ClusterFailurePolicy:
         max_restarts: Number of times the entire JobSet may be restarted before Flyte
             surfaces a RetryableFailure.
         restart_on_host_maintenance: When True, node evictions (DisruptionTarget condition)
-            trigger a free restart that does not consume the max_restarts budget.
+            trigger a free restart that does not consume the max_restarts budget. Free restarts
+            still increment `flyte.ctx().restart_attempt`.
     """
 
     max_restarts: int = 0

src/flyte/models.py

@@ -384,6 +384,10 @@ def master_port(self) -> Optional[int]:
 
     @property
     def restart_attempt(self) -> Optional[int]:
+        """How many times the JobSet has restarted the whole pod set within this Flyte attempt; None
+        outside clustered tasks. Free host-maintenance restarts count too, so treat this as a
+        "has the set restarted" counter, not as a position within `ClusterFailurePolicy.max_restarts`.
+        """
         v = os.environ.get("JOBSET_RESTART_ATTEMPT")
         return int(v) if v is not None else None