Skip to content

Allow opt-in recovery after task output upload - #3483

Open
1fanwang wants to merge 1 commit into
flyteorg:masterfrom
1fanwang:1fanwang/recover-post-output-shutdown-crashes
Open

1fanwang wants to merge 1 commit into
flyteorg:masterfrom
1fanwang:1fanwang/recover-post-output-shutdown-crashes

Conversation

@1fanwang

@1fanwang 1fanwang commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

Tracking issue

There is no linked public issue.

Why are the changes needed?

The fast-execution launcher reports failure when a native shutdown signal arrives after task results have been uploaded. A task that returns 42 and uploads its result still leaves the launcher at exit 245 when its exit hook receives SIGSEGV.

What changes were proposed in this pull request?

Add a recovery mode that is off by default and is switched on through the task environment. When it is on, the launcher restores the SDK's completed status only when a fresh completion record identifies its direct child and the upload path has finished. Task hooks run before upload, so a plugin cannot safely record this boundary.

Recovery excludes termination signals, forwarded cancellation, and earlier failures. It preserves error documents and the SDK's explicit failure status. Exit hooks still run, but a crash can interrupt remaining cleanup.

How was this patch tested?

Testing Done

The reproducer runs each commit's entrypoint source against the same installed SDK and dependencies, and it switches recovery on by setting FLYTE_RECOVER_SHUTDOWN_CRASH in the task environment. It launches the fast-execute and execute CLIs, unpacks a local code distribution, uploads a real protobuf result, and injects a signal in an exit hook. This is controlled fault injection, not a reproduction of a Ray native defect. No remote Flyte deployment or Ray cluster was exercised.

  • Local code review completed
# Scenario Command Result
1 Baseline entrypoint and fixed entrypoint with recovery disabled .venv/bin/python /tmp/flytekit-shutdown-repro-9d81b06f.py Both exit 245 after SIGSEGV despite the uploaded value 42.
2 Recovery enabled after SIGSEGV, SIGABRT, or SIGBUS Same invocation The launcher exits 0 and the decoded result is 42.
3 Crash before upload or failed upload Same invocation The launcher exits 245 and neither result nor error artifact exists.
4 SIGTERM after upload Same invocation The launcher retains exit 241 even though the result exists.

Setup process

Python 3.12.13 on macOS arm64, from the repository root with this PR checked out:

uv venv --python python3 .venv
SETUPTOOLS_SCM_PRETEND_VERSION=1.999.0dev0 uv pip install --python .venv/bin/python --quiet -r dev-requirements.in

The reproducer exports the original entrypoint from c3735a6 and the fixed entrypoint from the PR commit. Save the source below as /tmp/flytekit-shutdown-repro-9d81b06f.py.

.venv/bin/python /tmp/flytekit-shutdown-repro-9d81b06f.py
baseline-enabled: entrypoint=c3735a6a0adad174c09f1cfff69fe99090ca5b8c recovery=True exit=245
EXIT_HOOK signal=SIGSEGV
outputs.pb: o0=42; error.pb: absent
fixed-default: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=False exit=245
EXIT_HOOK signal=SIGSEGV
outputs.pb: o0=42; error.pb: absent
fixed-segv: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=True exit=0
EXIT_HOOK signal=SIGSEGV
{"asctime": "2026-11-04 22:30:05,734", "name": "flytekit", "levelname": "WARNING", "message": "Task subprocess crashed with signal 11 after output upload; restoring completed SDK exit status 0."}
outputs.pb: o0=42; error.pb: absent
fixed-abrt: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=True exit=0
EXIT_HOOK signal=SIGABRT
{"asctime": "2026-11-04 22:30:06,854", "name": "flytekit", "levelname": "WARNING", "message": "Task subprocess crashed with signal 6 after output upload; restoring completed SDK exit status 0."}
outputs.pb: o0=42; error.pb: absent
fixed-bus: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=True exit=0
EXIT_HOOK signal=SIGBUS
{"asctime": "2026-11-04 22:30:07,999", "name": "flytekit", "levelname": "WARNING", "message": "Task subprocess crashed with signal 10 after output upload; restoring completed SDK exit status 0."}
outputs.pb: o0=42; error.pb: absent
fixed-early: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=True exit=245
outputs.pb: absent; error.pb: absent
fixed-termination: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=True exit=241
EXIT_HOOK signal=SIGTERM
outputs.pb: o0=42; error.pb: absent
fixed-upload-failure: entrypoint=25c96e2cd3ba7e6997681db0fb9c83b528e985e5 recovery=True exit=245
EXIT_HOOK signal=SIGSEGV
outputs.pb: absent; error.pb: absent
Full launcher logs: /tmp/flyte-shutdown-4kk_xso9
Reproducer source: flytekit-shutdown-repro-9d81b06f.py
import os
import signal
import subprocess
import sys
import tarfile
import tempfile
from pathlib import Path

from flyteidl.core.literals_pb2 import LiteralMap

from flytekit.core.constants import ERROR_FILE_NAME, FLYTE_RECOVER_SHUTDOWN_CRASH, OUTPUT_FILE_NAME
from flytekit.loggers import LOGGING_FMT_ENV_VAR, LOGGING_RICH_FMT_ENV_VAR

BASELINE: str = "c3735a6a0adad174c09f1cfff69fe99090ca5b8c"
FIXED: str = "25c96e2cd3ba7e6997681db0fb9c83b528e985e5"
TASK_SOURCE: str = """\
import atexit
import os
import resource
import signal

from flytekit import task

resource.setrlimit(resource.RLIMIT_CORE, (0, 0))

def shutdown() -> None:
    crash_signal = getattr(signal, os.environ["SHUTDOWN_SIGNAL"])
    print(f"EXIT_HOOK signal={crash_signal.name}", flush=True)
    os.kill(os.getpid(), crash_signal)

@task
def value() -> int:
    atexit.register(shutdown)
    if os.environ["SHUTDOWN_MODE"] == "early":
        os.kill(os.getpid(), signal.SIGSEGV)
    return 42
"""

root: Path = Path(tempfile.mkdtemp(prefix="flyte-shutdown-", dir="/tmp"))
repo: Path = Path.cwd()
task_module: Path = root / "shutdown_task.py"
task_module.write_text(TASK_SOURCE)
distribution: Path = root / "code.tar.gz"
with tarfile.open(name=distribution, mode="w:gz") as archive:
    archive.add(name=task_module, arcname=task_module.name)
entrypoints: dict[str, Path] = {}
for revision in (BASELINE, FIXED):
    entrypoint: Path = root / f"entrypoint-{revision}.py"
    entrypoint.write_bytes(
        subprocess.check_output(args=["git", "show", f"{revision}:flytekit/bin/entrypoint.py"], cwd=repo)
    )
    entrypoints[revision] = entrypoint

cases: tuple[tuple[str, str, bool, str, signal.Signals, int, bool], ...] = (
    ("baseline-enabled", BASELINE, True, "shutdown", signal.SIGSEGV, (-signal.SIGSEGV) % 256, True),
    ("fixed-default", FIXED, False, "shutdown", signal.SIGSEGV, (-signal.SIGSEGV) % 256, True),
    ("fixed-segv", FIXED, True, "shutdown", signal.SIGSEGV, 0, True),
    ("fixed-abrt", FIXED, True, "shutdown", signal.SIGABRT, 0, True),
    ("fixed-bus", FIXED, True, "shutdown", signal.SIGBUS, 0, True),
    ("fixed-early", FIXED, True, "early", signal.SIGSEGV, (-signal.SIGSEGV) % 256, False),
    ("fixed-termination", FIXED, True, "shutdown", signal.SIGTERM, (-signal.SIGTERM) % 256, True),
    ("fixed-upload-failure", FIXED, True, "upload_failure", signal.SIGSEGV, (-signal.SIGSEGV) % 256, False),
)
for label, revision, enabled, mode, crash_signal, expected_exit, expect_output in cases:
    case: Path = root / label
    case.mkdir()
    inputs: Path = case / "inputs.pb"
    inputs.write_bytes(LiteralMap().SerializeToString())
    destination: Path = case / "code"
    destination.mkdir()
    output_prefix: Path = case / "outputs"
    if mode == "upload_failure":
        output_prefix.write_text("This file prevents the output-directory upload.")
    command: list[str] = [
        sys.executable, str(entrypoints[revision]), "pyflyte-fast-execute",
        "--additional-distribution", str(distribution), "--dest-dir", str(destination),
        "--", sys.executable, str(entrypoints[revision]), "pyflyte-execute",
        "--inputs", str(inputs), "--output-prefix", str(output_prefix),
        "--raw-output-data-prefix", str(case / "raw"),
        "--resolver", "flytekit.core.python_auto_container.default_task_resolver",
        "--", "task-module", "shutdown_task", "task-name", "value",
    ]
    result: subprocess.CompletedProcess[str] = subprocess.run(
        args=command,
        cwd=repo,
        env={
            **os.environ,
            FLYTE_RECOVER_SHUTDOWN_CRASH: str(enabled),
            LOGGING_FMT_ENV_VAR: "json",
            LOGGING_RICH_FMT_ENV_VAR: "0",
            "SHUTDOWN_MODE": mode,
            "SHUTDOWN_SIGNAL": crash_signal.name,
        },
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        timeout=45,
    )
    (case / "launcher.log").write_text(result.stdout)
    print(f"{label}: entrypoint={revision} recovery={enabled} exit={result.returncode}", flush=True)
    for line in result.stdout.splitlines():
        if "EXIT_HOOK" in line or "restoring" in line.lower():
            print(line, flush=True)
    output_file: Path = output_prefix / OUTPUT_FILE_NAME
    error_file: Path = output_prefix / ERROR_FILE_NAME
    if expect_output:
        decoded: LiteralMap = LiteralMap.FromString(output_file.read_bytes())
        observed: int = decoded.literals["o0"].scalar.primitive.integer
        print(f"outputs.pb: o0={observed}; error.pb: {'present' if error_file.is_file() else 'absent'}", flush=True)
        assert observed == 42
    else:
        print(f"outputs.pb: {'present' if output_file.is_file() else 'absent'}; error.pb: {'present' if error_file.is_file() else 'absent'}", flush=True)
        assert not output_file.is_file()
    assert not error_file.is_file()
    assert result.returncode == expected_exit, result.stdout
print(f"Full launcher logs: {root}", flush=True)

Screenshots

Not applicable to this runtime change.

Check all the applicable boxes

  • I updated the documentation accordingly.
  • All new and existing tests passed.
  • All commits are signed-off.

Related PRs

#2305 preserves automatic Ray shutdown so worker logs can finish flushing.

#3029 proposes a separate launcher startup optimization.

Docs link

https://github.com/flyteorg/flytekit/blob/25c96e2cd3ba7e6997681db0fb9c83b528e985e5/docs/source/design/execution.rst#recovering-completed-results-after-a-shutdown-crash

Preserve the completed SDK status after selected native shutdown signals. Keep cancellation, incomplete uploads, invalid completion records, and explicit task failures unrecovered.

Signed-off-by: 1fanwang <1fannnw@gmail.com>
@codecov

codecov Bot commented Nov 5, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.58%. Comparing base (d0c64d5) to head (25c96e2).
⚠️ Report is 3 commits behind head on master.

Additional details and impacted files
@@             Coverage Diff             @@
##           master    #3483       +/-   ##
===========================================
+ Coverage   46.34%   83.58%   +37.24%     
===========================================
  Files         327        3      -324     
  Lines       28868      195    -28673     
  Branches     3039        0     -3039     
===========================================
- Hits        13379      163    -13216     
+ Misses      15388       32    -15356     
+ Partials      101        0      -101     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant