Task state store in Apache Airflow®

The task state store, added in Apache Airflow® 3.3, allows you to save information that persists between task retries. This makes the task state store suitable for storing external job IDs and intra-task checkpoints.

In this guide, you’ll learn:

• When to use the task state store, and how it compares to XCom, deferrable operators, and the asset state store.
• How to set, get, delete, and clear values in the task state store.
• How to view and edit entries in the Airflow UI and through the Airflow REST API.
• How to create a custom operator that uses the task state store through the ResumableJobMixin.
• How to configure the task state store.

Assumed knowledge

To get the most out of this guide, you should have an understanding of:

• Basic Airflow concepts. See Introduction to Apache Airflow.
• Airflow operators. See Airflow operators.
• Airflow decorators. See Introduction to the TaskFlow API and Airflow decorators.

When to use the task state store

Use the task state store to persist information about a task instance’s own working state so that a later attempt of the same task instance, for example after a worker crash or retry, can access it. You can:

• Store the ID of a long-running external process, such as a Spark job or a model fine-tuning job, so a retry reconnects to the running job instead of submitting a duplicate.
• Store an intra-task checkpoint, such as the last page or offset processed, so a retry does not restart from the beginning.
• Store several pieces of state for one task instance, such as a status and a row count, and update them as the task runs.

When using the default task state store backend, Airflow saves entries to the metadata database, and stored values must be JSON serializable (except None). See task state store configuration for information on implementing a custom task state store backend.

Task state store vs. XCom

Both the task state store and XCom store small pieces of data attached to an individual task instance. The main difference is that Airflow clears XComs on retries, while task state entries persist.

Task state store vs. deferrable operators

The task state store, together with the ResumableJobMixin, and deferrable operators are both used in the context of long-running external jobs. Which to choose is often a question of personal preference. Deferrable operators free the worker slot while waiting and require a triggerer. The ResumableJobMixin keeps the worker slot but makes a synchronous operator retry-safe by reconnecting to a running job.

Task state store vs. asset state store

The task state store and the asset state store have a similar key-value interface, but they are scoped differently.

For more information on the asset state store, see the asset state store section of the advanced asset-based scheduling guide.

How to use the task state store

You have several options to interact with the task state store:

• In your task code using the methods available in the task_state_store key of the Airflow context in:
  • A @task-decorated function using the TaskFlow API.
  • A python_callable passed to a PythonOperator.
  • The execute method of a custom operator, created by subclassing BaseOperator or any existing traditional operator.
• In the Airflow UI
• With the Airflow REST API

Task state store methods

The task state store exposes four methods: set, get, delete, and clear.

set(key, value, *, retention=None) writes or overwrites the value for a key, in this case the PID of a long-running remote shell command. The keyword-only retention argument controls when the entry is automatically deleted:

• A timedelta, to delete the entry after the given duration from the time of the write. Airflow computes the expiry timestamp on the worker before sending the value to the API server.
• NEVER_EXPIRE, to keep the entry indefinitely. The cleanup process skips the entry, regardless of the environment default.
• None, the default, to fall back to the environment’s AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS configuration (default: 30 days).

1
context["task_state_store"].set("remote_pid", 239, retention=timedelta(days=28))

get(key, default=None) returns the stored value, or default if the key does not exist. For example, if a previous try of a task instance already started a long running shell command and saved the PID in the task state store, a later retry can get the PID in order to check if the command is still running, and if that is the case, wait on the existing command to finish instead of rerunning it from the start.

1
remote_pid = context["task_state_store"].get("remote_pid", default=None)

delete(key) deletes a single key. It does nothing if the key does not exist.

1
context["task_state_store"].delete("remote_pid")

clear() deletes all task state store keys for the task instance.

1
context["task_state_store"].clear()

Task state store examples

The following examples start a long-running shell command on a remote host and store its process ID (PID) in the task state store with the set method. If the first try fails and the task is retried, it uses the get method to reconnect to the still-running command instead of launching a new one.

1
from pendulum import duration
2
from airflow.sdk import task
3
from airflow.providers.ssh.hooks.ssh import SSHHook
4
5
6
def _ssh_exec(command):
7
    hook = SSHHook(ssh_conn_id="ssh_default")
8
    with hook.get_conn() as client:
9
        _, stdout, _ = client.exec_command(command)
10
        output = stdout.read().decode().strip()
11
        exit_status = stdout.channel.recv_exit_status()
12
    return exit_status, output
13
14
15
@task(retries=10, retry_delay=duration(minutes=10))
16
def run_ssh_command(launch_cmd, check_cmd, out_file, **context):
17
    tss = context["task_state_store"]
18
    remote_pid = tss.get("remote_pid")
19
20
    if remote_pid is None:
21
        _, remote_pid = _ssh_exec(launch_cmd.replace("{out_file}", out_file))
22
        tss.set("remote_pid", remote_pid)
23
24
    print(f"Reconnecting to remote PID {remote_pid} from a previous attempt.")
25
    exit_status, output = _ssh_exec(
26
        check_cmd.replace("{pid}", remote_pid).replace("{out_file}", out_file)
27
    )
28
    print(f"Remote job output: {output}")
29
    if exit_status != 0:
30
        raise RuntimeError(f"Remote job {remote_pid} failed with exit status {exit_status}.")
31
32
33
run_ssh_command(
34
    launch_cmd=(
35
        "nohup sh -c 'sleep 60; rc=$?; echo finished > {out_file}; "
36
        "echo $rc > {out_file}.rc' >/dev/null 2>&1 & echo $!"
37
    ),
38
    check_cmd=(
39
        "while kill -0 {pid} 2>/dev/null; do sleep 2; done; "
40
        "rc=$(cat {out_file}.rc 2>/dev/null); cat {out_file}; "
41
        '[ "$rc" = "0" ]'
42
    ),
43
    out_file="/tmp/log_{{ run_id }}.out",
44
)

1
from pendulum import duration
2
3
from airflow.providers.standard.operators.python import PythonOperator
4
from airflow.providers.ssh.hooks.ssh import SSHHook
5
6
7
def _ssh_exec(command):
8
    hook = SSHHook(ssh_conn_id="ssh_default")
9
    with hook.get_conn() as client:
10
        _, stdout, _ = client.exec_command(command)
11
        output = stdout.read().decode().strip()
12
        exit_status = stdout.channel.recv_exit_status()
13
    return exit_status, output
14
15
16
def _run_ssh_command_func(launch_cmd, check_cmd, out_file, **context):
17
    tss = context["task_state_store"]
18
    remote_pid = tss.get("remote_pid")
19
20
    if remote_pid is None:
21
        _, remote_pid = _ssh_exec(launch_cmd.replace("{out_file}", out_file))
22
        tss.set("remote_pid", remote_pid)
23
24
    print(f"Reconnecting to remote PID {remote_pid} from a previous attempt.")
25
    exit_status, output = _ssh_exec(
26
        check_cmd.replace("{pid}", remote_pid).replace("{out_file}", out_file)
27
    )
28
    print(f"Remote job output: {output}")
29
    if exit_status != 0:
30
        raise RuntimeError(f"Remote job {remote_pid} failed with exit status {exit_status}.")
31
32
33
PythonOperator(
34
    task_id="run_ssh_command",
35
    python_callable=_run_ssh_command_func,
36
    op_kwargs={
37
        "launch_cmd": (
38
            "nohup sh -c 'sleep 60; rc=$?; echo finished > {out_file}; "
39
            "echo $rc > {out_file}.rc' >/dev/null 2>&1 & echo $!"
40
        ),
41
        "check_cmd": (
42
            "while kill -0 {pid} 2>/dev/null; do sleep 2; done; "
43
            "rc=$(cat {out_file}.rc 2>/dev/null); cat {out_file}; "
44
            '[ "$rc" = "0" ]'
45
        ),
46
        "out_file": "/tmp/log_{{ run_id }}.out",
47
    },
48
    retries=10,
49
    retry_delay=duration(minutes=10),
50
)

1
from pendulum import duration
2
3
from airflow.sdk.bases.operator import BaseOperator
4
from airflow.providers.ssh.hooks.ssh import SSHHook
5
6
7
def _ssh_exec(command):
8
    hook = SSHHook(ssh_conn_id="ssh_default")
9
    with hook.get_conn() as client:
10
        _, stdout, _ = client.exec_command(command)
11
        output = stdout.read().decode().strip()
12
        exit_status = stdout.channel.recv_exit_status()
13
    return exit_status, output
14
15
16
class MyCustomOperator(BaseOperator):
17
18
    template_fields = ("launch_cmd", "check_cmd", "out_file")
19
20
    def __init__(self, *, launch_cmd, check_cmd, out_file, **kwargs):
21
        super().__init__(**kwargs)
22
        self.launch_cmd = launch_cmd
23
        self.check_cmd = check_cmd
24
        self.out_file = out_file
25
26
    def execute(self, context):
27
        tss = context["task_state_store"]
28
        remote_pid = tss.get("remote_pid")
29
30
        if remote_pid is None:
31
            _, remote_pid = _ssh_exec(self.launch_cmd.replace("{out_file}", self.out_file))
32
            tss.set("remote_pid", remote_pid)
33
34
        self.log.info(f"Reconnecting to remote PID {remote_pid} from a previous attempt.")
35
        exit_status, output = _ssh_exec(
36
            self.check_cmd.replace("{pid}", remote_pid).replace("{out_file}", self.out_file)
37
        )
38
        self.log.info(f"Remote job output: {output}")
39
        if exit_status != 0:
40
            raise RuntimeError(f"Remote job {remote_pid} failed with exit status {exit_status}.")
41
42
43
MyCustomOperator(
44
    task_id="run_ssh_command",
45
    launch_cmd=(
46
        "nohup sh -c 'sleep 60; rc=$?; echo finished > {out_file}; "
47
        "echo $rc > {out_file}.rc' >/dev/null 2>&1 & echo $!"
48
    ),
49
    check_cmd=(
50
        "while kill -0 {pid} 2>/dev/null; do sleep 2; done; "
51
        "rc=$(cat {out_file}.rc 2>/dev/null); cat {out_file}; "
52
        '[ "$rc" = "0" ]'
53
    ),
54
    out_file="/tmp/log_{{ run_id }}.out",
55
    retries=10,
56
    retry_delay=duration(minutes=10),
57
)

1
"""
2
Custom extension of the SSHOperator that stores the remote PID
3
in the task state store and reconnects to it on a retry.
4
"""
5
6
from airflow.providers.ssh.operators.ssh import SSHOperator
7
from pendulum import duration
8
import base64
9
10
11
class MyCustomSSHOperator(SSHOperator):
12
13
    template_fields = ("launch_cmd", "check_cmd", "out_file", *SSHOperator.template_fields)
14
15
    def __init__(self, *, launch_cmd, check_cmd, out_file, **kwargs):
16
        super().__init__(**kwargs)
17
        self.launch_cmd = launch_cmd
18
        self.check_cmd = check_cmd
19
        self.out_file = out_file
20
21
    def execute(self, context):
22
        tss = context["task_state_store"]
23
        remote_pid = tss.get("remote_pid")
24
25
        if remote_pid is None:
26
            self.command = self.launch_cmd.replace("{out_file}", self.out_file)
27
            remote_pid = base64.b64decode(super().execute(context)).decode().strip()
28
            tss.set("remote_pid", remote_pid)
29
30
        self.log.info(f"Reconnecting to remote PID {remote_pid} from a previous attempt.")
31
        self.command = self.check_cmd.replace("{pid}", remote_pid).replace("{out_file}", self.out_file)
32
        return super().execute(context)
33
34
35
MyCustomSSHOperator(
36
    task_id="run_ssh_command",
37
    ssh_conn_id="ssh_default",
38
    launch_cmd=(
39
        "nohup sh -c 'sleep 60; rc=$?; echo finished > {out_file}; "
40
        "echo $rc > {out_file}.rc' >/dev/null 2>&1 & echo $!"
41
    ),
42
    check_cmd=(
43
        "while kill -0 {pid} 2>/dev/null; do sleep 2; done; "
44
        "rc=$(cat {out_file}.rc 2>/dev/null); cat {out_file}; "
45
        '[ "$rc" = "0" ]'
46
    ),
47
    out_file="/tmp/log_{{ run_id }}.out",
48
    cmd_timeout=120,
49
    retries=10,
50
    retry_delay=duration(minutes=10),
51
)

Task state store in the Airflow UI

You can view and edit task state store entries for any task instance in the Airflow UI.

To reach the task state store of a specific task instance, select the task instance square in the grid (1), open the Storage tab (2), then select the Task State Store tab. From here, you can:

• Add an entry with Add Task State Store (3).
• Remove all entries for the task instance with Clear All Task State Store (4).
• Edit an entry (5).
• Delete an entry (6).

When you select Add Task State Store (3), the Add Task State Store window appears. Enter a Key and Value for the entry, then choose when it expires:

• Default (30 days): Use the default retention period for this Airflow environment, set by the default_retention_days configuration.
• Never: The entry never expires.
• Custom: The entry expires at the date and time you select.

Select Save to store the entry.

Task state store API endpoints

The Airflow REST API exposes task state store operations under each task instance’s state-store path:

/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store

The REST API limits each individual value to max_value_storage_bytes (64 KB by default) and rejects a larger value with a 422 error. For the request and response schemas of each operation, see the Airflow REST API reference.

ResumableJobMixin

For custom operators that submit one long-running external job and poll for its completion, you can use the ResumableJobMixin. After submitting the job, the mixin persists the external job ID to the task state store before polling starts. On a retry, the mixin reads the stored ID and reconnects to the running job instead of submitting a duplicate. The SparkSubmitOperator in the Apache Spark provider (6.2.0+) uses this mixin to reconnect to a running Spark application after a retry.

To use the mixin, inherit from ResumableJobMixin, call execute_resumable(context) from your execute method, and write the following six methods to determine how to interact with your external system:

• submit_job(context): Submit the job and return its external ID. The mixin stores this value in the task state store.
• get_job_status(external_id, context): Query the external system and return its raw status string.
• is_job_active(status): Return True if the job is still running and can be reconnected to.
• is_job_succeeded(status): Return True if the job completed successfully.
• poll_until_complete(external_id, context): Block until the job reaches a terminal state, and raise on failure.
• get_job_result(external_id, context): Return the job result after completion, or None if not applicable.

1
from typing import Any
2
3
from airflow.sdk import BaseOperator, ResumableJobMixin
4
5
6
class MyCustomResumableJobOperator(ResumableJobMixin, BaseOperator):
7
    external_id_key = "my_job_id"
8
9
    def execute(self, context):
10
        return self.execute_resumable(context)
11
12
    def submit_job(self, context) -> str:
13
        return self.hook.submit(...)  # Return the external job ID.
14
15
    def get_job_status(self, external_id, context) -> str:
16
        return self.hook.get_status(external_id)
17
18
    def is_job_active(self, status: str) -> bool:
19
        return status in ("RUNNING", "PENDING")
20
21
    def is_job_succeeded(self, status: str) -> bool:
22
        return status == "SUCCEEDED"
23
24
    def poll_until_complete(self, external_id, context) -> None:
25
        self.hook.wait(external_id)
26
27
    def get_job_result(self, external_id, context) -> Any:
28
        return None

On a retry, the mixin reads the stored ID and checks the current job status:

• If the job is still active, the mixin reconnects and continues polling.
• If the job already succeeded, the mixin returns the result without resubmitting.
• If the job is in a terminal failure state, the mixin submits a fresh job.

The external_id_key class attribute sets the key used to store the job ID. The default is remote_job_id.

Task state store configuration

You can configure the task state store with the [state_store] configs.

By default, the task state store and asset state store persist entries in the Airflow metadata database. To store entries elsewhere, you can provide a custom backend. See the Apache Airflow task and asset state store configuration documentation for more information.