Task state store in Apache Airflow®

• 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

• 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

• 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 doesn’t 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.

​

Task state store vs. XCom

​

Task state store vs. deferrable operators

​

Task state store vs. asset state store

​

How to use 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

• 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 exclude the entry from the automatic cleanup process and keep it until it is manually deleted.
• None, the default, to fall back to the environment’s AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS configuration (default: 30 days).

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

remote_pid = task_state_store.get("remote_pid", default=None)

task_state_store.delete("remote_pid")

task_state_store.clear()

​

Task state store examples

​

Task state store in the Airflow UI

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

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

​

Task state store API endpoints

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

​

ResumableJobMixin

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

from typing import Any

from airflow.sdk import BaseOperator, ResumableJobMixin


class MyCustomResumableJobOperator(ResumableJobMixin, BaseOperator):
    external_id_key = "my_job_id"

    def execute(self, context):
        return self.execute_resumable(context)

    def submit_job(self, context) -> str:
        return self.hook.submit(...)  # Return the external job ID.

    def get_job_status(self, external_id, context) -> str:
        return self.hook.get_status(external_id)

    def is_job_active(self, status: str) -> bool:
        return status in ("RUNNING", "PENDING")

    def is_job_succeeded(self, status: str) -> bool:
        return status == "SUCCEEDED"

    def poll_until_complete(self, external_id, context) -> None:
        self.hook.wait(external_id)

    def get_job_result(self, external_id, context) -> Any:
        return None

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

​

Task state store configuration