airflow-state-store

Airflow 任务状态存储（AIP-103）

Airflow 3.3 提供两个键值存储，以及一个用于提交外部作业的运算符的崩溃安全混入类。

task_state_store、asset_state_store 和 ResumableJobMixin 的崩溃安全保证需要 Airflow 3.3+。 请先检查：

af config version

低于 3.3：task_state_store/asset_state_store 不可用，且 durable=True 不起作用——提供程序运算符附带了一个 3.3 之前版本的 ResumableJobMixin 垫片，始终会重新提交（参见第 5 节）。请告知用户这些特定功能目前不可用，并链接到 AIP-103 跟踪问题。这不会影响第 6 节关于 Triggerer 与 mode="reschedule" 的选择，也不会影响“提交成功 ≠ 作业成功”这一通用反模式——这些内容适用于任何 Airflow 版本。对于运行在 3.3 之前版本上的 DAG，请完整提供相关建议；仅省略“durable=True 可增加崩溃安全性”这一部分。

---

第 1 节 — 选择合适的原语

我需要……	使用
持久化游标、偏移量或作业 ID，以便重试时恢复，而不是重新开始	task_state_store
在单个任务的多次重试之间传递少量协调状态（而不是在任务之间传递）	task_state_store
为每个资产存储水位标记或最后处理时间戳，并使其跨 DAG 运行保留	asset_state_store
缓存资产级元数据（清单哈希、行数、架构版本）	asset_state_store
让现有的非可延迟运算符在向外部系统提交作业时具备崩溃安全性	task_state_store 或 ResumableJobMixin

以下情况不要使用这些：

• 在任务之间传递数据 -> 使用 XCom
• 大型负载（模型权重、DataFrame）-> 使用带有对象存储后端的 XCom
• 在多个 DAG 之间共享配置或密钥 -> 使用 Variables 或 Connections

---

第 2 节 — 按需检测现有 DAG 中的反模式

用户要求审查 DAG，或询问“有没有更好的方法”时，请检查以下模式并指出问题：

DAG 中出现的模式	问题	建议
在 @task 函数体中使用 Variable.get(...) / Variable.set(...) 存储每次运行的状态	Variables 是全局共享的；不会按任务实例或重试范围隔离	task_state_store
使用 context["ti"].xcom_push(key="job_id", ...) 在重试后保留状态	XCom 的作用范围是一次 DAG 运行，而不是一次重试；每次重试都会分配新的 ti_id	task_state_store 或 ResumableJobMixin
手动编写 if Variable.get("job_id"): reconnect else: submit 重试恢复逻辑	重新实现了 ResumableJobMixin 已提供的功能，却没有崩溃安全保证	ResumableJobMixin
使用 Variable.set("last_processed_at", ...) 存储水位标记	是全局的；任何 DAG 或任务都可以覆盖它；不会按资产范围隔离	asset_state_store
单独的 submit 任务（wait_for_termination=False / 发起后不等待）+ 第二个传感器/轮询任务等待同一个外部作业（Databricks、Snowflake、BigQuery、Redshift、Spark 等）	提交任务变为绿色只代表作业已被接受，不代表作业已成功——只有传感器任务的结果才能反映实际情况。	参见第 6 节“提交并轮询的 DAG：一个任务还是两个？”——正确选择取决于 Triggerer 是否可用以及作业持续时间，而不是固定答案。

指出问题时，请展示修改前后的代码片段。修改后的示例请使用第 3–5 步中的规范示例。

即使传感器代码没有任何问题，拆分提交任务和传感器任务也值得提出意见——审查传感器代码质量（是否使用正确的 mode=、是否正确处理终止状态、是否缓存 hook），与判断两任务拆分是否是合适的架构，是两个独立的问题。请按照**第 6 节“提交并轮询的 DAG：一个任务还是两个？”**作出判断；不要在此重复推导。

---

第 3 节 — task_state_store：每个任务的协调状态

task_state_store 是一个键值存储，作用范围限定为单个任务实例标识（dag_id + run_id + task_id + map_index）。它会在重试期间保留——同一任务的新一轮重试会读取同一个存储。

from airflow.sdk import dag, task
from pendulum import datetime

@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
def etl_with_checkpoint():

    @task(retries=3)
    def process_records(**context):
        task_state_store = context["task_state_store"]  # injected by Airflow, no setup needed
        cursor = task_state_store.get("last_cursor", default=0)
        records = fetch_records_after(cursor)
        for record in records:
            process(record)
            cursor = record["id"]
            task_state_store.set("last_cursor", cursor)   # checkpoint after each record

    process_records()

etl_with_checkpoint()

API：

from airflow.sdk import NEVER_EXPIRE

task_state_store.get(key, default=None)                        # returns a JsonValue or default
task_state_store.set(key, value)                               # uses default_retention_days
task_state_store.set(key, value, retention=timedelta(days=7))  # per-key TTL override
task_state_store.set(key, value, retention=NEVER_EXPIRE)       # never expires regardless of config
task_state_store.delete(key)                                   # no-op if key does not exist
task_state_store.clear()                                       # delete all keys for this task instance

键规则：

• 值必须可序列化为 JSON（str、int、float、bool、list、dict——不接受 None 值）。
• 默认过期时间由 [state_store] default_retention_days 控制（0 = 永不过期）。
• 对于必须超出默认保留期限的键（例如，多日 Spark 作业的作业 ID），请使用 NEVER_EXPIRE。
• 值大小上限默认为 64 KB；可通过 [state_store] max_value_storage_bytes 配置（0 = 无限制）。对于更大的负载，请配置自定义 [state_store] backend，或通过 [workers] state_store_backend 配置工作器侧后端。

映射任务——每个索引都有自己的命名空间：

动态映射任务（task.expand(...)）中的每个映射索引都会获得一个独立的 task_state_store，其作用范围限定为自己的 map_index。各索引之间不共享状态。

@task(retries=2)
def process_partition(partition_id, **context):
    task_state_store = context["task_state_store"]
    # Scoped to THIS index only — other indices have their own copy
    cursor = task_state_store.get("cursor", default=0)
    task_state_store.set("cursor", new_cursor)

process_partition.expand(partition_id=[0, 1, 2, 3])

clear() 只会清除当前索引的状态。若要清除任务组所有映射索引的状态，请使用 CLI 或核心 API。

之前（反模式）：

@task
def process(**context):
    cursor = Variable.get("etl_cursor", default_var=0)
    # ... process ...
    Variable.set("etl_cursor", new_cursor)  # global, any task can overwrite

之后：

@task(retries=3)
def process(**context):
    task_state_store = context["task_state_store"]
    cursor = task_state_store.get("cursor", default=0)
    # ... process ...
    task_state_store.set("cursor", new_cursor)    # scoped to this task instance

---

第 4 节 — asset_state_store：跨 DAG 运行的每项资产元数据

asset_state_store 的作用范围限定为某项资产，而非任务实例。它会跨 DAG 运行持久保存——任何生成或使用该资产的任务，都可以读取和写入同一资产上的同一个键。

from airflow.sdk import DAG, Asset, task
from datetime import datetime, timezone

ORDERS = Asset(name="orders/daily", uri="s3://warehouse/orders/daily")

with DAG(dag_id="producer", schedule=None, start_date=datetime(2026, 1, 1), catchup=False):

    @task(inlets=[ORDERS], outlets=[ORDERS])
    def load(asset_state_store=None):        # asset_state_store injected by Airflow — declare as a kwarg
        asset_state_store = asset_state_store[ORDERS]

        watermark = asset_state_store.get("watermark", default="2026-01-01T00:00:00+00:00")
        records = fetch_records_since(watermark)

        now = datetime.now(tz=timezone.utc).isoformat()
        asset_state_store.set("watermark", now)
        asset_state_store.set("last_run_summary", {"rows_loaded": len(records), "completed_at": now})

    load()

从使用方 DAG 读取存储：

with DAG(dag_id="consumer", schedule=[ORDERS], start_date=datetime(2026, 1, 1), catchup=False):

    @task(inlets=[ORDERS])
    def consume(asset_state_store=None):
        asset_state_store = asset_state_store[ORDERS]
        summary = asset_state_store.get("last_run_summary") or {}
        print(f"Processing {summary.get('rows_loaded')} rows up to {asset_state_store.get('watermark')}")

    consume()

键规则：

• Airflow 会将 asset_state_store 作为命名关键字参数注入——请将其声明为 def my_task(asset_state_store=None)。不要与 **context 一起使用；Airflow 会单独注入它。
• 时间戳请使用 datetime.now(tz=timezone.utc).isoformat()——不要使用 datetime.utcnow()（它不带时区信息）。
• 与 task_state_store 一样，值必须可序列化为 JSON。
• 没有逐键过期机制——资产状态存储条目没有 TTL（资产的生命周期长于任何单次运行）。
• 任何将该资产声明为入口或出口的 DAG 都可以读取。

映射任务——最后写入者胜出：

asset_state_store 的作用范围限定为资产，而非映射索引。如果多个映射索引并发写入同一个键，则最后一次写入生效。请为每个索引使用不同的键，或确保只有一个索引写入指定键。

@task(outlets=[my_asset])
def load_partition(partition_id, asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    # Distinct key per index — no race condition
    asset_state_store.set(f"offset_{partition_id}", new_offset)

之前（反模式）：

Variable.set(f"watermark_{asset_name}", new_offset)   # global, not scoped to asset

之后：

@task(inlets=[my_asset], outlets=[my_asset])
def load(asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    asset_state_store.set("watermark", new_offset)

---

第 5 节 — ResumableJobMixin：防崩溃的外部作业提交

当任务需要向外部系统（Spark、Databricks、dbt Cloud、AWS Batch 等）提交作业，且重试时可能重复提交时，都应使用此混入类——无论该任务是在一次 execute() 调用中同时轮询作业是否完成，还是将作业 ID 交给单独的下游轮询/传感器任务。如果不使用此混入类（或已内置此功能的 Provider 操作符，例如默认设置 durable=True 的 DatabricksSubmitRunOperator），那么工作进程在提交后崩溃，就会导致下一次重试提交任务时重复提交作业——无论轮询是在同一任务中进行，还是由单独的任务完成，这种风险都存在。

推荐使用前先确认适用范围：如果提交和轮询已经是两个独立任务（例如，submit 任务通过 XCom 将作业 ID 传给下游传感器/轮询任务），那么轮询/传感器任务不需要此混入类——它不会提交任何内容，因此重试本身就是安全的。提交任务仍然需要它；拆分出轮询步骤并不能消除提交任务自身重复提交的风险。请参见下表中的对应行。

不应使用 ResumableJobMixin 的情况：

情况	改用	原因
已部署 Triggerer，且存在（或可以编写）可延迟操作符	可延迟操作符	在轮询期间释放工作进程槽位，更节省资源
任务在单次执行中扇出许多并发 I/O 操作	async def 任务 / BaseAsyncOperator	异步适用于高吞吐量 I/O，而非崩溃恢复
retries=0	—	无需恢复崩溃，也就没有需要重新连接的对象
外部系统没有可追踪的作业 ID（submit_job 返回 None）	普通操作符	混入类的防崩溃保证会被静默禁用，因此没有任何价值
提交和轮询已经拆分为两个独立任务（提交任务 → XCom → 传感器/轮询任务）	轮询/传感器任务不需要任何处理，但提交任务自身仍需要 durable=True	轮询/传感器任务不会提交任何内容，因此重试本身是幂等的——无需在此处设检查点。提交任务重试时仍可能重复提交，除非它自身具备防崩溃能力。例如，即使 wait_for_termination=False，DatabricksSubmitRunOperator.execute() 也会通过 execute_resumable() 执行（并在 durable=True 时为运行 ID 设检查点；默认即为此值）——请保留该默认值。手动编写的提交任务（直接调用 hook 的普通 @task）不会自动获得这些功能，应该自行实现 ResumableJobMixin。

ResumableJobMixin 会在整个轮询期间占用工作进程槽位——与标准同步操作符相同。它的优势在于防崩溃和保持作业连续性，而非提高资源效率。

选择退出崩溃恢复：

此混入类默认设置为 durable=True。将其设为 durable=False 可跳过所有 task_state_store 交互，并执行普通的提交/轮询/结果流程——适用于测试环境，或外部系统已有自身去重机制的情况：

MyBatchOperator(task_id="job", durable=False)

# Or via default_args to disable for all tasks in a DAG:
with DAG("my_dag", default_args={"durable": False}):
    ...

实现此混入类

from airflow.sdk import BaseOperator, ResumableJobMixin
from pydantic import JsonValue


class MyBatchOperator(BaseOperator, ResumableJobMixin):

    external_id_key = "batch_job_id"   # key used in task_state_store; set once, never rename

    def execute(self, context):
        return self.execute_resumable(context)  # never call self.execute() — call this

    def submit_job(self, context) -> JsonValue:
        # Submit and return the job identifier. This value is persisted to task_state_store
        # before polling starts. Return None only if the system has no trackable ID
        # (in that case crash-safety is disabled and the job resubmits on every retry).
        return self.hook.submit_batch(...)

    def get_job_status(self, external_id: JsonValue, context) -> str:
        # Query the external system. Return a raw status string.
        return self.hook.get_status(external_id)

    def is_job_active(self, status: str) -> bool:
        # Return True if the job is still running and should be reconnected to.
        return status in ("RUNNING", "PENDING", "QUEUED")

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

    def poll_until_complete(self, external_id: JsonValue, context) -> None:
        # Block until the job reaches a terminal state. Raise on failure.
        self.hook.wait(external_id)

    def get_job_result(self, external_id: JsonValue, context):
        # Return the job result after success. Return None if not applicable.
        return None

重试时会发生什么

重试时的作业状态	混入类的行为
仍在运行	重新连接——调用 poll_until_complete，不再提交
已成功	立即返回 get_job_result
失败 / 未知	提交新作业

external_id_key 警告

**对于已部署且有正在运行任务实例的操作符，绝不要重命名 external_id_key。**旧键会以之前的名称存储在 task_state_store 中。重命名会使混入类将每个活动重试都视为新的提交，从而破坏防崩溃保证。

之前（反模式）：

def execute(self, context):
    job_id = Variable.get("spark_job_id", default_var=None)
    if job_id and self._is_running(job_id):
        self._wait(job_id)
    else:
        job_id = self.hook.submit(...)
        Variable.set("spark_job_id", job_id)   # global, race-prone
        self._wait(job_id)

之后：

class MySparkOperator(BaseOperator, ResumableJobMixin):
    external_id_key = "spark_job_id"
    def execute(self, context): return self.execute_resumable(context)
    def submit_job(self, context): return self.hook.submit(...)
    # ... implement the 5 other methods ...

对于“提交 + 单独的传感器任务”这种 DAG 结构，具体来说是否应将其合并为一个任务，请参见 第 6 节“提交并轮询的 DAG：一个任务还是两个？”。这是一个由 Triggerer 是否可用以及作业运行时长决定的架构选择。它与提交任务本身是否需要 durable=True 来保证崩溃安全是不同的问题（通常需要，无论是否保留拆分）——请参见上方表格行。

---

第 6 节 — 提交并轮询的 DAG：一个任务还是两个？

Databricks/Snowflake/BigQuery/Redshift/Spark 作业常见的一种 DAG 结构是：一个 submit 任务以 wait_for_termination=False（或等效设置）启动作业并立即返回，随后由手动编写的传感器任务轮询同一作业，直至完成。值得特别指出的是，提交任务在外部系统接受作业的那一刻就会显示成功（而不是等作业完成）——仅看提交任务的重试、告警或 SLA，无法了解作业是否真正成功，只有传感器任务的结果才能说明这一点。但是否应取消这种拆分，取决于 Triggerer 是否可用：

如果已部署 Triggerer——将其合并为一个设置了 deferrable=True 的任务，并明确设置 wait_for_termination=True。deferrable=True 本身并不意味着 wait_for_termination=True——对于 DatabricksSubmitRunOperator，只有同时设置 wait_for_termination=True，延迟辅助程序才会进入延迟状态；如果沿用旧提交任务中的 wait_for_termination=False，它就会悄无声息地跳过延迟和轮询，重现本节旨在消除的“发出即忘”问题。

这样仍可在轮询期间释放工作进程，而且任务自身的结果会反映作业的实际结果——但这并不是毫无取舍、严格更优的升级：延迟路径会直接提交作业，不会经过 mixin 的检查点步骤，因此 durable=True 在这里无法提供保护（与下方的同步路径不同）。从工作进程崩溃到 Trigger 接管之间若发生崩溃，重试时仍会重新提交作业。

run_job = DatabricksSubmitRunOperator(
    task_id="run_job",
    tasks=[{"task_key": "job", "notebook_task": {"notebook_path": NOTEBOOK_PATH}}],
    deferrable=True,
    wait_for_termination=True,   # required — deferrable=True alone does not imply this
)

如果未部署 Triggerer——这确实需要权衡，不能自动地一概而论。具体分为两种情况，取决于作业运行期间占用一个工作进程是否可以接受：

• 作业可以同步运行（运行时间较短或中等，或者工作进程容量并不紧张），且运算符支持持久化执行（继承了第 5 节中的 ResumableJobMixin，并提供 durable=True）——建议用运算符的持久化模式将两任务模式替换为一个任务：wait_for_termination=True、durable=True（两者已经是默认值）。这不仅仅是减少了需要管理的环节——运算符内置的轮询循环从提交到轮询全程都由 ResumableJobMixin 保护，因此工作进程在轮询中崩溃时会重新连接，而不是重新提交。手动编写的传感器任务属于独立的非 mixin 代码路径，有自己的重试语义，通常更弱（请检查它是否设置了自己的 retries；很多都没有）。在这种情况下合并任务，删掉的是一个需要持续维护的文件，而不只是一个任务：

  run_job = DatabricksSubmitRunOperator(
      task_id="run_job",
      tasks=[{"task_key": "job", "notebook_task": {"notebook_path": NOTEBOOK_PATH}}],
      wait_for_termination=True,   # synchronous — holds the worker for the run, acceptable here
      durable=True,                # default; ResumableJobMixin checkpoints the run id end-to-end
  )
  

• 作业确实运行时间很长，而且工作进程槽位紧张，占用一个槽位直到作业结束正是需要避免的成本——保留 submit + mode="reschedule" 传感器的拆分。这种情况下，这才是合理且高效利用工作进程的设计，而不是一看到就该删除的反模式。

无论哪种情况：

• 如果保留拆分，请确认下游任务（以及任何告警/SLA）依赖的是传感器任务，而不是 submit。如果下游有什么内容以 submit 成功为依据，那才是真正的问题——应修正依赖关系，而不是架构。
• 在评审中明确说明其中的权衡（工作进程槽位成本和拆分后的结果，与单任务正确性和内置崩溃安全性之间的取舍），不要断言某一方就是“正确模式”。 崩溃安全性与是否折叠有关联，但不是同一个决策。 是否移除拆分取决于工作槽成本（见上文），与 durable=True 无关。无论是否折叠，提交任务通常都应保持启用该设置（请参阅第 5 节中“When NOT to use ResumableJobMixin”表格：轮询／传感器任务不需要此设置，但提交任务自身的重复提交风险不会仅仅因为它与轮询任务拆分开就消失）。如果选择折叠，优先使用算子自身的 durable=True 模式，而不是保留额外接入的手写传感器——内置路径已在一处为提交和轮询都提供了 ResumableJobMixin 的崩溃安全性，而手写传感器无法复现这一点。

---

第 7 节 — 配置参考

[state_store]
# Full dotted path to the storage backend. Default writes to the Airflow metadata DB.
backend = airflow.state.metastore.MetastoreStateStoreBackend

# Days to retain task state store entries after their last update. 0 = disable time-based cleanup.
# Does NOT affect asset_state_store rows — asset state store has no TTL.
default_retention_days = 30

# Rows deleted per batch during cleanup. 0 = no batching (single unbounded delete).
# Tune on large deployments to reduce lock contention.
state_cleanup_batch_size = 0

# Auto-delete all task state store keys when a task succeeds. Default: False.
# Does NOT affect asset_state_store — asset state store persists across runs and must be cleared explicitly.
clear_on_success = False

工作端后端（可选，[workers] 部分）——通过本地后端路由任务状态存储写入，然后再传到 API 服务器。如果大型负载或需要凭据的存储应保留在工作端，此功能会很有用：

[workers]
state_store_backend = mypackage.store.WorkerSideBackend

---

第 8 节 — 安全检查清单

• Airflow 版本 ≥ 3.3（af config version）
• 值可序列化为 JSON（str、int、float、bool、list、dict——不能是 datetime 或自定义对象）
• task_state_store 键是简短且描述性明确的字符串（避免使用点号和斜杠）
• 向 asset_state_store 写入数据的映射任务：每个索引使用不同的键，或接受最后写入者胜出的语义
• 映射任务：全局清除状态时，在下游任务中使用 CLI／核心 API，而不是在任务主体内调用 clear()
• ResumableJobMixin：已设置 external_id_key，且部署后不会重命名
• ResumableJobMixin：execute() 调用 self.execute_resumable(context)，而不是执行自定义逻辑
• ResumableJobMixin：如果已禁用崩溃恢复，使用 durable=False 是有意为之
• 大型负载（超过配置的 max_value_storage_bytes）使用自定义的 [state_store] backend，或通过以下配置启用工作端后端：[workers] state_store_backend

---

相关技能

• authoring-dags — DAG 编写的通用模式和惯例。
• airflow-hitl — 暂停 DAG 以等待人工审批（Airflow 3.1+）。
• airflow — af config、af registry 和 Airflow CLI 通用参考。