Crash recovery for YARN cluster mode in SparkSubmitOperator built on AIP-103

[amoghrajesh]

---

Was generative AI tooling used to co-author this PR?

• Yes: Claude Sonnet 4.6

Built on top of a few PRs: #67118 and even #65991

What problem are we solving?

SparkSubmitOperator in YARN cluster mode submits the Spark driver as a YARN AM — it runs independently on the cluster. If the Airflow worker dies mid-poll, the YARN app keeps running but Airflow loses track of it. The only recovery today is to submit a brand new application, wasting already-done compute or causing conflicts if the job is not idempotent.

Proposed change

Extends ResumableJobMixin crash recovery (introduced in #67118 for Spark standalone cluster mode) to YARN cluster mode.

On first run:

• Injects spark.yarn.submit.waitAppCompletion=false so spark-submit exits after YARN accepts the app
• Parses the YARN application ID from spark-submit output and persists it to task_state
• Polls the YARN ResourceManager REST API (GET /ws/v1/cluster/apps/{appId}) until terminal, built on top of Spark: track YARN cluster apps via RM REST API to free JVM #65991

On retry (worker crash):

• Reads the application ID from task_state
• Queries YARN resourcemanager for current status, reconnects if still running, skips polling if already succeeded, resubmits only if failed/killed

Status mapping:

The hook queries GET /ws/v1/cluster/apps/{appId} and synthesizes YARN API's two-field response into a single string for the mixin interface:

YARN state	finalStatus	Synthesized status	Mixin behaviour
NEW / NEW_SAVING / SUBMITTED / ACCEPTED / RUNNING	UNDEFINED	state as-is	Poll again
FINISHED	SUCCEEDED	SUCCEEDED	Return result
FINISHED	FAILED / KILLED / UNDEFINED	FAILED	Resubmit
FAILED	FAILED / KILLED	FAILED	Resubmit
KILLED	KILLED	FAILED	Resubmit

state determines whether the app is still running; finalStatus resolves the outcome when state=FINISHED. Using finalStatus alone would be unreliable because YARN can report UNDEFINED for a dead app if the RM recovered from a crash and lost the final status.

Observability: Logs status transitions such as ACCEPTED → RUNNING and a heartbeat every 10 polls so users can see the job is alive without log spam.

on_kill: Since spark-submit has already exited (waitAppCompletion=false), the hook's CLI kill has nothing to terminate. YARN cluster mode on_kill uses the Resource Manager REST API instead (PUT /ws/v1/cluster/apps/{appId}/state).

What changes from the standalone PR (#67118)

The standalone PR covered spark:// masters only. This PR adds the parallel YARN path. The two paths are kept as separate if branches in execute() for readability.

User impact and backcompat

New behaviour (opt-in by existing default):

• Any SparkSubmitOperator with --master yarn --deploy-mode cluster now gets crash recovery automatically. reconnect_on_retry=True is the default
• Requires one new connection extra: yarn_resourcemanager_webapp_address (e.g. http://rm.example.com:8088). Without it, the resumable path raises ValueError at submit time with a clear message.
• Trade-off to be aware of: Full spark-submit log streaming (infrastructure-level YARN orchestrator logs) is replaced by RM REST API polling. Actual driver/executor logs are unaffected — they remain in YARN log aggregation (yarn logs -applicationId ...). Most users will not notice; users relying on orchestrator log streaming should set reconnect_on_retry=False

Unaffected paths:

• YARN client mode — unaffected, still uses the blocking hook.submit() path
• Kubernetes — unaffected, still raises NotImplementedError (sister PR)
• reconnect_on_retry=False — skips crash recovery, submits and polls without task_state persistence

Testing Details

Setup

Using this docker compose setup to spin up a 4 node / container Hadoop 3.2.1 cluster: NN, DN, RM, NM.

services:
namenode:
image: bde2020/hadoop-namenode:2.0.0-hadoop3.2.1-java8
container_name: yarn-namenode
ports:
- "9870:9870"
- "9000:9000"
volumes:
- namenode_data:/hadoop/dfs/name
environment:
- CLUSTER_NAME=yarn-test
env_file:
- hadoop.env
networks:
- yarn-net
datanode:
image: bde2020/hadoop-datanode:2.0.0-hadoop3.2.1-java8
container_name: yarn-datanode
volumes:
- datanode_data:/hadoop/dfs/data
env_file:
- hadoop.env
depends_on:
- namenode
networks:
- yarn-net
resourcemanager:
image: bde2020/hadoop-resourcemanager:2.0.0-hadoop3.2.1-java8
container_name: yarn-resourcemanager
ports:
- "8088:8088"
- "8030:8030"
- "8031:8031"
- "8032:8032"
env_file:
- hadoop.env
depends_on:
- namenode
- datanode
networks:
- yarn-net
nodemanager:
image: bde2020/hadoop-nodemanager:2.0.0-hadoop3.2.1-java8
container_name: yarn-nodemanager
env_file:
- hadoop.env
depends_on:
- resourcemanager
networks:
- yarn-net
volumes:
namenode_data:
datanode_data:
networks:
yarn-net:
driver: bridge

Prepared a script to setup few things / bootstrap them for running spark on hadoop -- setup Java, stages Spark JARs to HDFS, write core-site.xml and yarn-site.xml config, register spark_yarn airflow connection with the right yarn_resourcemanager_webapp_address:

Connection

airflow connections add spark_yarn \
--conn-type spark \
--conn-host yarn \
--conn-extra '{"deploy-mode": "cluster", "yarn_resourcemanager_webapp_address": "http://host.docker.internal:8088"}'

Using this DAG:

from datetime import datetime
from airflow.sdk import DAG
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
with DAG(
dag_id="spark_yarn_mode_repro",
start_date=datetime(2024, 1, 1),
schedule=None,
catchup=False,
tags=["aip-103", "yarn", "repro"],
) as dag:
SparkSubmitOperator(
task_id="submit_long_running_job",
conn_id="spark_yarn",
application="/opt/airflow/dev/spark-cluster-yarn/spark-examples.jar",
java_class="org.apache.spark.examples.SparkPi",
application_args=["100"],
conf={
"spark.yarn.archive": "hdfs:///spark-jars.zip",
"spark.executor.memory": "512m",
"spark.driver.memory": "512m",
"spark.executor.cores": "1",
},
env_vars={
"HADOOP_CONF_DIR": "/tmp/hadoop-conf",
"HADOOP_USER_NAME": "root",
},
num_executors=1,
retries=2,
retry_delay=5,
yarn_track_via_rm_api=True
)

The DAG submits SparkPi with 10,000 iterations (~2 min runtime) — long enough to kill the worker mid-poll and verify reconnection on retry. Check the YARN UI to confirm only one application was submitted across multiple Airflow task attempt.

This is how a successful run will look like:

Steps to Verify reconnect behaviour

1. Trigger the dag and wait for yarn application to spin up

2. Kill the worker mid way and watch the logs

Log:

[2026-06-04 13:15:45] ERROR - Process terminated by signal. For more information, see https://airflow.apache.org/docs/apache-airflow/stable/troubleshooting.html#process-terminated-by-signal. signal=-2 signal_name=SIGINT

3. Observe on yarn UI, application continues running and completes

4. Resume worker once yarn app is complete to observe no resubmission

---

• 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.

[amoghrajesh]

Running w full tests

[amoghrajesh]

cc @nailo2c

[potiuk]

Some checks are.failing though

[amoghrajesh]

I'll be on it, @nailo2c is going to help me test this one soon too

[aaron-y-chen]

I tested the SparkSubmitOperator.on_kill() path against a Kerberized Dataproc/YARN cluster (on GCP), and the YARN kill request failed because the RM REST PUT was sent without Kerberos auth.

What passed:

• Kerberos RM REST polling for a successful YARN app.
• Crash recovery when the original YARN app already succeeded: retry skipped resubmission.
• Crash recovery while the original YARN app was still running: retry reconnected to the same app.
• Crash recovery after the original YARN app was killed/failed: retry submitted a fresh app.

What failed:

YARN application submitted: application_1780697276912_0017
Tracking YARN application application_1780697276912_0017 via ResourceManager REST API polling
YARN application application_1780697276912_0017 status: RUNNING
YARN kill request for application_1780697276912_0017 returned HTTP 401
YARN application application_1780697276912_0017 is still RUNNING
YARN application application_1780697276912_0017 status: FINISHED

The RM state confirmed that the app was not killed; it completed naturally:

{
"id": "application_1780697276912_0017",
"state": "FINISHED",
"finalStatus": "SUCCEEDED",
"name": "airflow-67473-kill"
}

So the local task process received SIGTERM, Airflow attempted to kill the YARN app, but the Kerberized RM returned 401. The YARN app kept running and eventually finished successfully instead of becoming KILLED.

This appears to come from SparkSubmitOperator.on_kill() calling the public hook.kill_yarn_application(), while that public method does not pass auth=self._resolved_yarn_rm_auth to the REST PUT. The existing private _kill_yarn_application() path does pass the resolved auth.

I also verified the proposed fix locally by making the public method delegate to
the existing private method:

defkill_yarn_application(self, application_id: str) ->None:
"""Public alias for ResumableJobMixin / operator on_kill paths."""self._kill_yarn_application(application_id)

With that change in place, the same SIGTERM / on_kill() path sent an
authenticated RM REST request and killed the YARN app:

YARN kill request for application_1780697276912_0019 returned HTTP 202
YARN application application_1780697276912_0019 status: KILLED
RuntimeError: YARN application application_1780697276912_0019 ended with state: KILLED, final status: KILLED

{
"id": "application_1780697276912_0019",
"state": "KILLED",
"finalStatus": "KILLED",
"name": "airflow-67473-kill"
}

It would also be worth adding an operator-level on_kill() test, since a hook-level on_kill() test does not cover this operator path.

[amoghrajesh]

Thanks for your comments @nailo2c! I have handled them now

[aaron-y-chen]

Hi @amoghrajesh, I've tested your latest code, and it worked well on my dataproc yarn cluster, LGTM :)

[amoghrajesh]

Addressed all review comments. Here is a summary of whats in the final state:

Hook (spark_submit.py)

• _is_yarn_cluster_mode flag added (YARN + deploy_mode=cluster)
• submit() no longer blocks for RM API polling when yarn_track_via_rm_api=True. It now exits after capturing _yarn_application_id. Polling responsibility moved to the operator.
• _track_yarn_application renamed to _start_yarn_application_status_tracking with state-change logging and a periodic heartbeat every 10 polls
• query_yarn_application_status() added: normalizes the YARN (state, finalStatus) tuple to a single string for the ResumableJobMixin interface
• kill_yarn_application() public wrapper removed. on_kill() in the operator calls _kill_yarn_application() directly so Kerberos auth is applied

Operator (spark_submit.py)

Four paths for YARN cluster mode in execute():

1. reconnect_on_retry=True + yarn_track_via_rm_api=True -- full crash recovery via execute_resumable()
2. reconnect_on_retry=True + yarn_track_via_rm_api=False -- raises ValueError at startup (no way to resume without RM API)
3. reconnect_on_retry=False + yarn_track_via_rm_api=True -- submit + poll without task_state persistence
4. reconnect_on_retry=False + yarn_track_via_rm_api=False -- falls through to hook.submit(), spark-submit blocks with waitAppCompletion=true (unchanged legacy behavior)

on_kill() kills via REST API for YARN cluster mode since spark-submit has already exited at that point.

[amoghrajesh]

Merging this one now, thanks @nailo2c for production testing.

[github-actions]

Backport failed to create: v3-2-test. View the failure log Run details

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
❌	v3-2-test

You can attempt to backport this manually by running:

cherry_picker f3ca1c2 v3-2-test

This should apply the commit to the v3-2-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.

After you have resolved the conflicts, you can continue the backport process by running:

cherry_picker --continue

If you don't have cherry-picker installed, see the installation guide.

[amoghrajesh]

No need to backport

[aaron-y-chen]

Merging this one now, thanks @nailo2c for production testing.

My pleasure 😄