feat(cli): add --recover (rerun) and --recover-from (run) (#1237)

Adds the recovery surface to the CLI and un-gates it end to end
(ENG26-830).

```bash
# Recover from the run being rerun (fetched code, reuse its succeeded actions):
flyte rerun ul56wcvgqrb9vzhzz5l2 --recover

# Recover a fresh run with NEW local code from a prior run:
flyte run main.py main --recover-from ul56wcvgqrb9vzhzz5l2

# Force named actions to re-execute even though they succeeded in the source run:
flyte rerun ul56wcvgqrb9vzhzz5l2 --recover --force-rerun-action a3 --force-rerun-action a7
```

Both map to `with_runcontext(recover=...)`: `--recover` (bool) recovers
from the run being rerun; `--recover-from <run>` (string) recovers a
fresh `run()` from a named prior run. Recovery reuses the prior run's
succeeded actions and re-runs only what failed or changed (remote-only;
`--local` is rejected). `--force-rerun-action` (repeatable, maps to
`with_runcontext(recover_force_rerun_actions=[...])` →
`RunSpec.recover.force_rerun_actions`) forces named actions to
re-execute anyway; a listed parent re-enqueues its children (list them
too to force the whole subtree), unknown names are ignored.

Since the original draft:

- Merged main (#1236 landed there and evolved: `related_to` provenance,
`run_base_dir`, image-cache seeding).
- Reworked to the wire contract that actually merged in
flyteorg/flyte#7615: a recovery run is `RunSpec.relation = {related_to:
<run>, relation_type: RELATION_TYPE_RECOVER}`; `RunSpec.recover` carries
only optional extras (`force_rerun_actions`) and stays unset. There is
no `Recover.run_id`.
- Bumped flyteidl2 to 2.0.28 (pyproject, rs_controller pyproject +
Cargo, uv.lock), which opens the runtime gate — the
`NotImplementedError` path only triggers on older flyteidl2 builds.
- Treat `ACTION_PHASE_RECOVERED` (10) as terminal/success everywhere
phases are classified: python controller (`Action.is_terminal`, covering
both watch snapshot and live updates), rust controller
(`is_action_terminal` + crate bump — 2.0.27 dropped/panicked on unknown
phase 10), `Run.wait()`/`watch` done-checks, `flyte.models.ActionPhase`,
and TUI display. Without this, recovered children never resolve and the
run hangs re-opening `WatchForUpdates` (reproduced on devbox
2026-07-16). On resolution the child's `output_uri` points at the source
run's outputs and is consumed as-is.
- Skew-tolerant phase handling: bindings that predate
`ACTION_PHASE_RECOVERED` (e.g. a stale task-image venv) no longer kill
the run — the wire value (10) is used as a `getattr` fallback, and
phase-name rendering goes through a helper that never raises on unknown
enum values.

Recovery never reads the global cache — `overwrite_cache` /
`cache_lookup_scope` are orthogonal.

Needs the actions-service recovery logic (ENG26-828) for end-to-end
behavior.

---------

Signed-off-by: Ketan Umare <kumare3@users.noreply.github.com>
Signed-off-by: Sergey Vilgelm <sergey@union.ai>
Co-authored-by: Ketan Umare <kumare3@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Sergey Vilgelm <sergey@union.ai>

Author: 3 people

rs_controller/src/action.rs

@@ -144,6 +144,9 @@ impl Action {
                     | ActionPhase::Failed
                     | ActionPhase::Aborted
                     | ActionPhase::TimedOut
+                    // Recovered from a prior run: terminal, success-equivalent; output_uri
+                    // points at the source run's outputs (consume as-is).
+                    | ActionPhase::Recovered
             )
         } else {
             false

src/flyte/_internal/controllers/remote/_action.py

@@ -18,6 +18,11 @@
 
 ActionType = Literal["task", "trace", "condition"]
 
+# ACTION_PHASE_RECOVERED landed in flyteidl2 2.0.28; the bindings in a task image may predate
+# it. The wire value is stable, so fall back to it — never crash on an enum value the local
+# bindings don't know.
+_ACTION_PHASE_RECOVERED: int = getattr(phase_pb2, "ACTION_PHASE_RECOVERED", 10)
+
 
 # This class should be deleted following move to pyo3.
 @dataclass
@@ -64,6 +69,9 @@ def is_terminal(self) -> bool:
             phase_pb2.ACTION_PHASE_SUCCEEDED,
             phase_pb2.ACTION_PHASE_ABORTED,
             phase_pb2.ACTION_PHASE_TIMED_OUT,
+            # Recovered from a prior run: terminal and success-equivalent; output_uri points at
+            # the source run's outputs (intentional — consume as-is).
+            _ACTION_PHASE_RECOVERED,
         ]
 
     def increment_retries(self):

src/flyte/_internal/controllers/remote/_informer.py

@@ -11,6 +11,7 @@
 from flyteidl2.workflow import state_service_pb2
 
 from flyte._logging import log, logger
+from flyte._utils.helpers import action_phase_name
 
 from ._action import Action
 from ._service_protocol import ActionsService, StateService
@@ -41,13 +42,11 @@ async def observe_state(self, state: state_service_pb2.ActionUpdate) -> Action:
         """
         Add an action to the cache if it doesn't exist. This is invoked by the watch.
         """
-        logger.debug(f"Observing phase {phase_pb2.ActionPhase.Name(state.phase)} for {state.action_id.name}")
+        logger.debug(f"Observing phase {action_phase_name(state.phase)} for {state.action_id.name}")
         if state.output_uri:
             logger.debug(f"Output URI: {state.output_uri}")
         else:
-            logger.warning(
-                f"{state.action_id.name} has no output URI, in phase {phase_pb2.ActionPhase.Name(state.phase)}"
-            )
+            logger.warning(f"{state.action_id.name} has no output URI, in phase {action_phase_name(state.phase)}")
         if state.phase == phase_pb2.ACTION_PHASE_FAILED:
             logger.error(
                 f"Action {state.action_id.name} failed with error (msg):"

src/flyte/_run.py

@@ -7,7 +7,7 @@
 import sys
 import uuid
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Sequence, Tuple, Union, cast
 
 from flyte._context import Context, contextual_run, internal_ctx
 from flyte._environment import Environment
@@ -139,6 +139,25 @@ def _ambient_image_cache() -> ImageCache | None:
     return tctx.compiled_image_cache if tctx else None
 
 
+def _uri_inputs_hash(inputs_uri: str) -> str:
+    """Deterministic stand-in for ``OffloadedInputData.inputs_hash`` when the inputs blob
+    cannot be read client-side (rerun of a source run whose outputs were cleaned up).
+
+    The server computes the real hash as FNV-64a over the marshaled inputs and requires the
+    field to be non-empty (it feeds cache-key computation). Hashing the source inputs URI in
+    the same format is conservative: identical source location -> identical inputs -> same
+    key, while an accidental match with a content-derived hash is a ~2^-64 event — the same
+    collision odds the content hash itself carries. Worst case is a lost cache hit, never a
+    wrong one beyond those odds.
+    """
+    import base64
+
+    h = 0xCBF29CE484222325
+    for b in inputs_uri.encode("utf-8"):
+        h = ((h ^ b) * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF
+    return base64.urlsafe_b64encode(h.to_bytes(8, "big")).decode().rstrip("=")
+
+
 def _to_cache_lookup_scope(scope: CacheLookupScope | None = None):
     """Map the SDK cache-lookup-scope literal onto its RunSpec enum value."""
     from flyteidl2.task import run_pb2
@@ -188,6 +207,8 @@ def __init__(
         preserve_original_types: bool | None = None,
         debug: bool = False,
         recover: bool | str | None = False,
+        recover_force_rerun_actions: Sequence[str] | None = None,
+        allow_missing_source_outputs: bool = False,
         _tracker: Any = None,
         _bundle_relative_paths: tuple[str, ...] | None = None,
         _bundle_from_dir: pathlib.Path | None = None,
@@ -241,6 +262,14 @@ def __init__(
         # Carried on RunSpec.relation with RELATION_TYPE_RECOVER; remote-only; gated in
         # _apply_overrides until the flyteidl2 field + backend ship. See _resolve_recover_ref.
         self._recover = recover
+        # Escape hatch: actions that must re-execute in the recovery run even if they succeeded
+        # in the source run (RunSpec.recover.force_rerun_actions). Only valid with recover.
+        self._recover_force_rerun_actions = tuple(recover_force_rerun_actions or ())
+        if self._recover_force_rerun_actions and not self._recover:
+            raise ValueError("recover_force_rerun_actions requires recover to be set")
+        # Opt-in for rerun/recover of a source run whose outputs were cleaned up from storage:
+        # proceed with its inputs URI instead of failing. See the fallback in rerun().
+        self._allow_missing_source_outputs = allow_missing_source_outputs
 
     def _resolve_recover_ref(self, rerun_run_name: str | None) -> str | None:
         """Resolve `self._recover` to the reference run name to recover from (or None).
@@ -527,7 +556,7 @@ def _apply_overrides(self, base: Any, *, task: Any = None, relation: Tuple[Any,
             run_spec.CopyFrom(base)
             # Provenance is per-run, never inherited: a rerun of a rerun must point at its immediate
             # parent (set below from `relation`), not the grandparent captured in the prior spec.
-            for provenance_field in ("relation", "related_to"):
+            for provenance_field in ("relation", "related_to", "recover"):
                 # DESCRIPTOR internals are opaque to checkers; guard for fields absent from the current pin.
                 if provenance_field in cast(Any, run_pb2.RunSpec).DESCRIPTOR.fields_by_name:
                     run_spec.ClearField(provenance_field)
@@ -588,17 +617,35 @@ def _apply_overrides(self, base: Any, *, task: Any = None, relation: Tuple[Any,
                     cast(Any, run_spec).relation.CopyFrom(
                         _relation_pb.Relation(related_to=ref, relation_type=relation_type)
                     )
+                if kind == "recover" and self._recover_force_rerun_actions:
+                    # Escape hatch: these actions re-execute even though they succeeded in the
+                    # source run. A listed parent re-enqueues its children (list them too to force
+                    # the whole subtree); unknown names are ignored server-side.
+                    cast(Any, run_spec).recover.CopyFrom(
+                        cast(Any, run_pb2).Recover(force_rerun_actions=list(self._recover_force_rerun_actions))
+                    )
 
         return run_spec
 
     async def _submit_remote(
-        self, *, task_spec: Any, task_id: Any, proto_inputs: Any, run_spec: Any, run_id: Any, project_id: Any
+        self,
+        *,
+        task_spec: Any,
+        task_id: Any,
+        proto_inputs: Any,
+        run_spec: Any,
+        run_id: Any,
+        project_id: Any,
+        offloaded_input_data: Any = None,
     ) -> Run:
         """Upload inputs and create the run. The single network call site for remote submission.
 
         Consumes an already-built ``run_spec`` (see ``_apply_overrides``), raw proto ``inputs``
         (``flyteidl2.task.Inputs``), and a task by reference (``task_id``) or by value
-        (``task_spec``); shared by ``_run_remote`` and ``rerun``.
+        (``task_spec``); shared by ``_run_remote`` and ``rerun``. ``offloaded_input_data``
+        (``flyteidl2.common.OffloadedInputData``) references already-offloaded inputs (e.g. the
+        source run's inputs.pb on a rerun whose inputs can't be re-downloaded) and skips the
+        upload; exactly one of ``proto_inputs`` / ``offloaded_input_data`` is used.
         """
         from connectrpc.code import Code
         from connectrpc.errors import ConnectError
@@ -609,28 +656,30 @@ async def _submit_remote(
         from flyte.remote import Run
 
         try:
-            upload_req = dataproxy_service_pb2.UploadInputsRequest(inputs=proto_inputs)
-            # Pass the explicit run_base_dir so the offloaded inputs are written under the
-            # same base the CreateRun below resolves (RunSpec.run_base_dir, set in _apply_overrides).
-            # When unset the server falls back to settings/cluster default in both paths.
-            if self._run_base_dir:
-                upload_req.base_dir = self._run_base_dir
-            # Reference an already-registered task by id; otherwise upload the full spec.
-            if task_id is not None:
-                upload_req.task_id.CopyFrom(task_id)
-            else:
-                upload_req.task_spec.CopyFrom(task_spec)
-            if run_id is not None:
-                upload_req.run_id.CopyFrom(run_id)
-            else:
-                upload_req.project_id.CopyFrom(project_id)
+            if offloaded_input_data is None:
+                upload_req = dataproxy_service_pb2.UploadInputsRequest(inputs=proto_inputs)
+                # Pass the explicit run_base_dir so the offloaded inputs are written under the
+                # same base the CreateRun below resolves (RunSpec.run_base_dir, set in _apply_overrides).
+                # When unset the server falls back to settings/cluster default in both paths.
+                if self._run_base_dir:
+                    upload_req.base_dir = self._run_base_dir
+                # Reference an already-registered task by id; otherwise upload the full spec.
+                if task_id is not None:
+                    upload_req.task_id.CopyFrom(task_id)
+                else:
+                    upload_req.task_spec.CopyFrom(task_spec)
+                if run_id is not None:
+                    upload_req.run_id.CopyFrom(run_id)
+                else:
+                    upload_req.project_id.CopyFrom(project_id)
 
-            upload_resp = await get_client().dataproxy_service.upload_inputs(upload_req)
+                upload_resp = await get_client().dataproxy_service.upload_inputs(upload_req)
+                offloaded_input_data = upload_resp.offloaded_input_data
 
             create_req = run_service_pb2.CreateRunRequest(
                 run_id=run_id,
                 project_id=project_id,
-                offloaded_input_data=upload_resp.offloaded_input_data,
+                offloaded_input_data=offloaded_input_data,
                 run_spec=run_spec,
             )
             # Reference an already-registered task by id; otherwise send the full spec.
@@ -1139,6 +1188,8 @@ async def rerun(
             version = task_spec.task_template.id.version
 
         # Inputs: reuse the prior raw proto inputs, or convert new native kwargs against the interface.
+        proto_inputs = None
+        offloaded_input_data = None
         if inputs:
             if task_template is not None:
                 iface = task_template.native_interface
@@ -1149,10 +1200,67 @@ async def rerun(
             converted = await convert_from_native_to_inputs(iface, custom_context=self._custom_context, **inputs)
             proto_inputs = converted.proto_inputs
         else:
-            resp = await get_client().dataproxy_service.get_action_data(
-                request=dataproxy_service_pb2.GetActionDataRequest(action_id=action_details.pb2.id)
-            )
-            proto_inputs = resp.inputs
+            # Rerun/recover only need the source run's INPUTS. GetActionData resolves inputs AND
+            # outputs server-side concurrently and 404s wholesale when either blob has been
+            # cleaned up (retention) — and which half the error names is a race. The client has
+            # no RPC to check the inputs blob alone, so a missing-data 404 is a hard error by
+            # default; `allow_missing_source_outputs` opts into proceeding with the inputs URI
+            # (fails at runtime if the inputs turn out to be gone too).
+            from connectrpc.code import Code
+            from connectrpc.errors import ConnectError
+            from flyteidl2.common import run_pb2 as common_run_pb2
+            from flyteidl2.workflow import run_service_pb2
+
+            import flyte.errors
+
+            try:
+                resp = await get_client().dataproxy_service.get_action_data(
+                    request=dataproxy_service_pb2.GetActionDataRequest(action_id=action_details.pb2.id)
+                )
+                proto_inputs = resp.inputs
+            except ConnectError as e:
+                if e.code != Code.NOT_FOUND:
+                    raise
+                if "inputs" in str(e.message):
+                    # The inputs blob itself is gone — nothing to feed the new run; fail
+                    # fast with a clear story instead of the server's raw 404.
+                    raise flyte.errors.RuntimeUserError(
+                        "SourceRunInputsUnavailableError",
+                        f"Source run {run_name}'s inputs are no longer in storage (deleted by "
+                        f"retention/cleanup), so it cannot be rerun or recovered with its "
+                        f"original inputs. Pass new inputs explicitly instead: "
+                        f"flyte.with_runcontext(...).rerun('{run_name}', inputs={{...}}), or "
+                        f"launch fresh local code with `flyte run ... --recover-from {run_name}` "
+                        f"(inputs come from the CLI parameters).",
+                    ) from e
+                if not self._allow_missing_source_outputs:
+                    raise flyte.errors.RuntimeUserError(
+                        "SourceRunOutputsUnavailableError",
+                        f"Source run {run_name}'s outputs are no longer in storage. Rerun/recover "
+                        f"only needs its inputs, but whether those still exist cannot be verified "
+                        f"from the client. If you know the inputs are intact, retry with "
+                        f"--allow-missing-outputs "
+                        f"(with_runcontext(allow_missing_source_outputs=True)); if they were "
+                        f"deleted too, the new run would fail at runtime — pass new inputs "
+                        f"explicitly instead (rerun('{run_name}', inputs={{...}}) or "
+                        f"`flyte run ... --recover-from {run_name}`).",
+                    ) from e
+                uris = await get_client().run_service.get_action_data_u_r_is(
+                    run_service_pb2.GetActionDataURIsRequest(action_id=action_details.pb2.id)
+                )
+                if not uris.inputs_uri:
+                    raise
+                logger.warning(
+                    f"Source run {run_name} outputs are no longer in storage; proceeding with its "
+                    f"inputs at {uris.inputs_uri} (--allow-missing-outputs). If the inputs were "
+                    f"deleted too the new run will fail at runtime, and recovered actions "
+                    f"referencing deleted outputs will fail if consumed "
+                    f"(use --force-rerun-action to re-execute them)."
+                )
+                offloaded_input_data = common_run_pb2.OffloadedInputData(
+                    uri=uris.inputs_uri,
+                    inputs_hash=_uri_inputs_hash(uris.inputs_uri),
+                )
 
         run_id, project_id = self._resolve_run_target(project, domain, cfg.org)
 
@@ -1186,6 +1294,7 @@ async def rerun(
             run_spec=run_spec,
             run_id=run_id,
             project_id=project_id,
+            offloaded_input_data=offloaded_input_data,
         )
 
 
@@ -1223,6 +1332,8 @@ def with_runcontext(
     preserve_original_types: bool = False,
     debug: bool = False,
     recover: bool | str | None = False,
+    recover_force_rerun_actions: Sequence[str] | None = None,
+    allow_missing_source_outputs: bool = False,
     _tracker: Any = None,
 ) -> _Runner:
     """
@@ -1306,8 +1417,17 @@ async def example_task(x: int, y: str) -> str:
     :param recover: Recover (reuse a prior run's succeeded actions, re-running only what failed or
         changed). ``True`` recovers from the run being rerun — only valid with ``.rerun(...)``; a
         run-name string recovers from that named run and is the only form valid on ``.run(...)``.
-        Remote-only. Not yet supported by the backend (raises NotImplementedError at submit until
-        flyteidl2 RunSpec.relation ships).
+        Remote-only. Requires a backend (and flyteidl2 build) with RunSpec.relation recovery
+        support; raises NotImplementedError at submit otherwise.
+    :param recover_force_rerun_actions: Optional names of actions that must re-execute in the
+        recovery run even if they succeeded in the source run (escape hatch). A listed parent
+        action re-enqueues its children — list them too to force the whole subtree; a listed
+        condition re-pauses for a new signal. Unknown names are ignored. Only valid with
+        ``recover``.
+    :param allow_missing_source_outputs: Opt-in for ``rerun``/recover when the source run's
+        outputs were cleaned up from storage: proceed using the source inputs URI instead of
+        failing. The client cannot verify the inputs still exist — if they were deleted too,
+        the new run fails at runtime.
     :param _tracker: This is an internal only parameter used by the CLI to render the TUI.
 
     :return: runner
@@ -1358,6 +1478,8 @@ async def example_task(x: int, y: str) -> str:
         preserve_original_types=preserve_original_types,
         debug=debug,
         recover=recover,
+        recover_force_rerun_actions=recover_force_rerun_actions,
+        allow_missing_source_outputs=allow_missing_source_outputs,
         _tracker=_tracker,
     )
 

src/flyte/_utils/helpers.py

@@ -87,3 +87,19 @@ def _selector_policy():
         yield
     finally:
         asyncio.set_event_loop_policy(original_policy)  # ty: ignore[deprecated]  # kept until 3.16 drops the API
+
+
+def action_phase_name(phase: int) -> str:
+    """Human-readable name for an ActionPhase wire value.
+
+    Proto3 enums are open: the server may send values these bindings don't know
+    (e.g. ACTION_PHASE_RECOVERED from a newer flyteidl2), and ``Name()`` raises
+    ValueError on them. Never crash on display — fall back to the known name for
+    stable wire values, or a generic one otherwise.
+    """
+    from flyteidl2.common import phase_pb2
+
+    try:
+        return phase_pb2.ActionPhase.Name(phase)
+    except ValueError:
+        return "ACTION_PHASE_RECOVERED" if phase == 10 else f"ACTION_PHASE_{phase}"