merge main and resolve

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

Author: wild-endeavor

examples/advanced/hash_root_action.py

@@ -0,0 +1,87 @@
+"""
+Content-based caching for the *root* task of a run.
+
+`hash_flyte_dataframe.py` shows content-based caching between tasks: a driver produces a
+DataFrame and calls a cached consumer twice, and the second call hits. This example covers the
+other entrypoint — passing a locally-built DataFrame straight into `flyte.run(...)`, so the
+cached task *is* the root action of the run.
+
+The two paths compute the cache key differently. A sub-action's key is computed by the
+controller in-process, which substitutes `Literal.hash` for the literal's contents. A root
+action's key is derived from the offloaded inputs, which reference the upload URI — and every
+`flyte.run` uploads the DataFrame to a fresh URI. So without a content hash the second run
+misses even though the bytes are identical.
+
+Passing `hash_method=` to `DataFrame.from_local_sync` makes both runs agree: the key follows
+the content, not where it happened to land in blob storage.
+
+Run it twice; `check_cache_hit` below asserts the second run returns the first run's value.
+"""
+
+import pandas as pd
+
+import flyte
+from flyte import Cache
+from flyte.io import DataFrame, HashFunction
+
+img = flyte.Image.from_debian_base(name="flyte-root-hash").with_pip_packages("pandas", "pyarrow")
+
+env = flyte.TaskEnvironment(
+    "flyte_root_action_hash",
+    image=img,
+    resources=flyte.Resources(cpu="1", memory="2Gi"),
+)
+
+SAMPLE_DATA = {"id": [1, 2, 3, 4, 5], "value": [100, 200, 300, 400, 500]}
+
+
+def hash_pandas_dataframe(df: pd.DataFrame) -> str:
+    """Content-based hash: the same rows always produce the same digest."""
+    return str(pd.util.hash_pandas_object(df).sum())
+
+
+@env.task(cache=Cache(behavior="override", version_override="v1"))
+async def main(df: DataFrame) -> str:
+    """Cached root task.
+
+    The random number is the cache probe: it is regenerated on every real execution, so two
+    runs returning the same string can only mean the second one was served from the cache.
+    """
+    import random
+
+    pdf = await df.open(pd.DataFrame).all()
+    return f"rows={len(pdf)}, total={pdf['value'].sum()}, random={random.randint(1, 1000000)}"
+
+
+def build_input() -> DataFrame:
+    """The DataFrame to submit, tagged with a content-based hash.
+
+    Without `hash_method` the cache key would follow the (per-run, always new) upload URI and
+    the second run would miss.
+    """
+    return DataFrame.from_local_sync(
+        pd.DataFrame(SAMPLE_DATA),
+        hash_method=HashFunction.from_fn(hash_pandas_dataframe),
+    )
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+
+    # Two independent submissions of the same content. Each uploads to its own URI.
+    run1 = flyte.run(main, df=build_input())
+    print(f"Run 1: {run1.url}")
+    run1.wait()
+    result1 = run1.outputs()[0]
+
+    run2 = flyte.run(main, df=build_input())
+    print(f"Run 2: {run2.url}")
+    run2.wait()
+    result2 = run2.outputs()[0]
+
+    print(f"\nRun 1: {result1}")
+    print(f"Run 2: {result2}")
+    if result1 == result2:
+        print("\n✓ Cache hit — the new upload URI did not change the cache key.")
+    else:
+        print("\n✗ Cache miss — the root action's key still tracks the upload URI.")

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)

examples/integration_tests.py

@@ -393,6 +393,35 @@ async def test_advanced_local_tasks(flyte_client):
     await _run_and_wait(parallel_main_no_io, "test_advanced_local_tasks", q="hello")
 
 
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_advanced_hash_root_action(flyte_client):
+    """Content-based caching must work when the cached task is the run's root action.
+
+    A sub-action's cache key is computed by the controller, which substitutes `Literal.hash`
+    for the literal's contents; a root action's key is derived from the offloaded inputs, which
+    reference the upload URI. Each `flyte.run` uploads to a fresh URI, so this only passes if
+    the content hash — not the URI — drives the key.
+    """
+    from examples.advanced.hash_root_action import build_input, main
+
+    results = []
+    for i in (1, 2):
+        # Rebuilt each time, so run 2 uploads identical bytes to a brand-new URI.
+        run = await flyte.with_runcontext(log_level=logging.DEBUG).run.aio(main, df=build_input())
+        print(f"\n[test_advanced_hash_root_action] Run {i}: {run.url}")
+        run.wait()
+        detail = await run.action.details()
+        if detail.error_info:
+            raise RuntimeError(f"Run {i} failed with error: {detail.error_info.message}")
+        results.append(run.outputs()[0])
+
+    # The task embeds a fresh random number on every real execution, so identical outputs
+    # mean run 2 was served from run 1's cache entry.
+    assert results[0] == results[1], f"root action cache miss across a new upload URI: {results[0]!r} != {results[1]!r}"
+    print("  Cache hit across a new upload URI\n")
+
+
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_advanced_multi_loops(flyte_client):