Skip to content

Replace reattach_on_restart with durable execution in KPO#69914

Open
amoghrajesh wants to merge 2 commits into
apache:mainfrom
astronomer:kpo-crash-recovery
Open

Replace reattach_on_restart with durable execution in KPO#69914
amoghrajesh wants to merge 2 commits into
apache:mainfrom
astronomer:kpo-crash-recovery

Conversation

@amoghrajesh

Copy link
Copy Markdown
Contributor

Was generative AI tooling used to co-author this PR?
  • Yes Sonnet 4.5

Why?

KPO's existing reattach_on_restart relies on a label search on every retry, which can raise FoundMoreThanOnePodFailure when more than one matching pod exists, and gives no way to reconnect to a specific pod unambiguously. On Airflow 3.3+, persisting the running pod's identity to task state store lets a retry reconnect directly to that exact pod instead, removing the ambiguity window entirely. reattach_on_restart is deprecated in favor of a new durable flag that supersedes it; existing configurations map onto durable with zero behavior change, and pre-3.3 environments keep today's label-search behavior unchanged.

Summary

  • Adds a durable parameter to KubernetesPodOperator (default True), superseding the
    deprecated reattach_on_restart parameter.
  • On Airflow 3.3+, durable=True persists the running pod's identity (name, namespace, uid)
    to task state store before the operator waits on it. On retry, the operator reconnects
    directly to that pod via a direct read_namespaced_pod lookup instead of a label search --
    removing the FoundMoreThanOnePodFailure ambiguity window a label search can hit when more
    than one matching pod exists.
  • The existing label-search mechanism (find_pod()) is retained as a one-time bootstrap
    fallback: a pod created by a pre-upgrade provider version (no persisted identity yet) is
    still found via labels on the first retry after upgrading, and that retry starts persisting
    state for every subsequent one. No in-flight task or existing pod breaks across an upgrade.
  • reattach_on_restart still works, mapping its value onto durable with zero behavior
    change; passing it emits AirflowProviderDeprecationWarning on Airflow 3.3+ only (below
    3.3, durable/task state store don't functionally exist yet, so warning would be
    premature).
  • durable=False opts out of task state store entirely, preserving today's reattach_on_restart=False
    behavior exactly, including its pre-existing incidental-cleanup quirk (an orphaned pod from
    a killed prior attempt can still get found and deleted via unconditional find_pod()
    refetch calls elsewhere in execute_sync/execute_async -- unrelated to and unchanged by
    this work).
  • Docs: added a "Durable execution" section to KPO's operator docs, cross-linking task state
    store.

Testing

Using this dag:

from datetime import datetime, timedelta

from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from airflow.sdk import DAG

with DAG(
    dag_id="try_kpo",
    schedule=None,
    start_date=datetime(2024, 1, 1),
    catchup=False,
) as dag:
    run_pod = KubernetesPodOperator(
        task_id="run_pod",
        name="try-kpo",
        namespace="default",
        image="ubuntu:24.04",
        cmds=["bash", "-cx"],
        arguments=["sleep 60 && echo done"],
        in_cluster=False,
        do_xcom_push=False,
        retries=2,
        retry_delay=timedelta(seconds=5),
    )

Before my changes:

Worker killed:

image

But pod still exists:

image

Reconnected when worker came back up:

image

After my changes (changing to pass durable=True)

    run_pod = KubernetesPodOperator(
        task_id="run_pod",
        name="try-kpo",
        namespace="default",
        image="ubuntu:24.04",
        cmds=["bash", "-cx"],
        arguments=["sleep 60 && echo done"],
        in_cluster=False,
        do_xcom_push=False,
        retries=2,
        retry_delay=timedelta(seconds=5),
        durable=True,
    )

WOrker killed:

image

Task store persisted:

image
[Breeze:3.10.20] root@0014bf761b80:/opt/airflow$ cat /tmp/airflow_state/ti_run_pod/pod_identifier.json
{"name": "try-kpo-vlap9a9c", "namespace": "default", "uid": "878b321f-fb21-4779-94c4-1811f7ced720"}[Breeze:3.10.20]

Pod present and reconnected to via state store:

image
  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@raphaelauv

Copy link
Copy Markdown
Contributor

with reattach_on_restart the operator check at every restart if a matching pod exist

but with durable , the operator only check if the value on the state is not empty ? So in case of hard-crash of the worker just after calling k8S to run the pod , the operator will not try to fund the pod ?

thanks

@amoghrajesh

Copy link
Copy Markdown
Contributor Author

@raphaelauv good comment, but that's covered already. _get_pod_from_task_state_store returning None (nothing persisted yet) falls straight through to find_pod(), the same label search reattach_on_restart always did.

So, the crash-before-persist window youre describing falls to the exact old behavior for that one retry & the operator still checks for a matching pod via labels, same as reattach_on_restart always did. It only skips the label search once the identity is actually persisted on a later retry.

@raphaelauv

Copy link
Copy Markdown
Contributor

thanks for the quick answer :)

That's perfect ,also it would be great to make it explicit in the doc

@amoghrajesh

Copy link
Copy Markdown
Contributor Author

@raphaelauv thanks, handled it in 7a75ba6

Comment on lines +179 to +181
:param reattach_on_restart: deprecated, use ``durable`` instead. If the worker dies while the pod
is running, reattach and monitor during the next try. If False, always create a new pod for
each try.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
:param reattach_on_restart: deprecated, use ``durable`` instead. If the worker dies while the pod
is running, reattach and monitor during the next try. If False, always create a new pod for
each try.
:param reattach_on_restart: deprecated for Airflow 3.3+, use ``durable`` instead. If the worker dies while the pod
is running, reattach and monitor during the next try. If False, always create a new pod for
each try.

Comment on lines +182 to +186
:param durable: if the worker dies while the pod is running, reattach and monitor during the next
try instead of creating a duplicate pod. If False, always create a new pod for each try.
Supersedes ``reattach_on_restart``; on Airflow 3.3+ the reconnection uses a persisted pod
identity in task state store instead of a label search, removing the ambiguity failure a label
search can hit when more than one matching pod exists. Defaults to ``True``.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
:param durable: if the worker dies while the pod is running, reattach and monitor during the next
try instead of creating a duplicate pod. If False, always create a new pod for each try.
Supersedes ``reattach_on_restart``; on Airflow 3.3+ the reconnection uses a persisted pod
identity in task state store instead of a label search, removing the ambiguity failure a label
search can hit when more than one matching pod exists. Defaults to ``True``.
:param durable: if the worker dies while the pod is running, reattach and monitor during the next
try instead of creating a duplicate pod. If False, always create a new pod for each try.
Supersedes ``reattach_on_restart``; on Airflow 3.3+ the reconnection uses a persisted pod
identity in task state store instead of a label search, removing the ambiguity failure a label
search can hit when more than one matching pod exists. Defaults to ``True``. Only available for Airflow 3.3+

Comment on lines +392 to +403
if "reattach_on_restart" in kwargs:
# Dropped from the named signature entirely so the warning fires on any explicit use,
# `True` or `False`.
reattach_on_restart = kwargs.pop("reattach_on_restart")
if AIRFLOW_V_3_3_PLUS:
warnings.warn(
"`reattach_on_restart` is deprecated and will be removed in a future release. "
"Use `durable` instead.",
AirflowProviderDeprecationWarning,
stacklevel=2,
)
durable = reattach_on_restart

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we need also the handle the case of durable with Airflow<3.3
I think for that case we want to just make it work as reattach_on_restart so we want ask them to change parameter name but just handle this for them under the hood?

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

Labels

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

3 participants