airflow-state-store

Airflow タスク状態ストア（AIP-103）

Airflow 3.3 では、2 つのキー／バリューストアと、外部ジョブを送信するオペレーター向けのクラッシュセーフティミックスインが提供されます。

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
外部システムへの送信を行う既存の非デフerr可能オペレーターに、クラッシュセーフティを持たせる	task_state_store または ResumableJobMixin

次の場合は使用しないでください。

• タスク 間 のデータ受け渡し → XCom を使用します
• 大きなペイロード（モデルの重み、データフレーム） → オブジェクトストレージバックエンドとともに 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 など）を待つ 2 つ目のセンサー／ポーリングタスクを分ける	送信タスクが成功しても意味するのはジョブが 受け付けられた ことだけであり、ジョブが 成功した ことではありません。実際の状態を反映するのはセンサータスクの結果だけです。	**セクション 6「送信とポーリングを行う DAG：1 つのタスクか、2 つか？」**を参照してください。適切な判断は Triggerer の利用可否とジョブの実行時間によって異なり、固定の答えはありません。

指摘する際は、変更前と変更後のスニペットを示してください。変更後の例には、ステップ 3～5 の標準的な例を使用してください。

センサーのコードに問題がなくても、送信タスクとセンサータスクの分割には触れる価値があります。 センサーのコード品質（mode= の正しさ、終端状態の適切な処理、フックのキャッシュ）をレビューすることと、2 タスクへの分割が適切なアーキテクチャかどうかは別の問題です。この判断には **セクション 6「送信とポーリングを行う DAG：1 つのタスクか、2 つか？」**に従ってください。ここで独自に判断し直さないでください。

---

セクション 3 — task_state_store：タスクごとの調整状態

task_state_store は、単一のタスクインスタンス ID（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(...)）、各マップインデックスには、独自の map_index をスコープとする分離された task_state_store が割り当てられます。インデックス間で状態は共有されません。

@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()

キーのルール：

• asset_state_store は Airflow によって名前付きのキーワード引数として注入されます。def my_task(asset_state_store=None) のように宣言してください。**context と組み合わせないでください。Airflow はこれを別途注入します。
• タイムスタンプには datetime.now(tz=timezone.utc).isoformat() を使用してください。タイムゾーン対応ではない datetime.utcnow() は使用しないでください。
• task_state_store と同じく、値は JSON にシリアライズ可能である必要があります。
• キーごとの有効期限はありません。アセットの状態ストアのエントリーには TTL がありません（アセットは単一の実行より長く存続します）。
• アセットを inlet または outlet として宣言するすべての DAG から読み取り可能です。

マップされたタスク — 最後の書き込みが優先されます：

asset_state_store はマップインデックスではなく、アセットをスコープとします。複数のマップインデックスが同じキーに同時に書き込む場合、最後の書き込みが優先されます。インデックスごとに異なるキーを使用するか、指定したキーに書き込むインデックスを 1 つだけにしてください。

@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)

---

Section 5 — ResumableJobMixin：クラッシュに強い外部ジョブの送信

外部システム（Spark、Databricks、dbt Cloud、AWS Batch など）にジョブを送信するタスクで、リトライ時に重複ジョブを送信する可能性がある場合に使用します。同じ execute() 呼び出し内で完了までポーリングする場合も、ジョブ ID を別の後続のポーリング／センサータスクに渡す場合も対象です。このミックスイン（または DatabricksSubmitRunOperator の durable=True というデフォルトのように、すでに組み込まれているプロバイダーオペレーター）を使わない場合、送信後にワーカーがクラッシュすると、次回の送信タスクのリトライで重複ジョブが送信されます。このリスクは、ポーリングを同じタスクで行うか、別のタスクで行うかにかかわらず存在します。

推奨する前にスコープを確認してください： 送信とポーリングがすでに別々のタスクになっている場合（たとえば、submit タスクが XCom 経由でジョブ ID を後続のセンサー／ポーリングタスクに渡す場合）、ポーリング／センサータスクにはこのミックスインは不要です。何も送信しないため、リトライしてもすでに安全です。送信タスクには依然として必要です。ポーリング手順を切り離しても、送信タスク自体の重複送信リスクはなくなりません。下の表の該当する行を参照してください。

ResumableJobMixin を使用しない場合：

状況	代わりに使用するもの	理由
Triggerer がデプロイされており、遅延可能なオペレーターが存在する（または実装できる）	遅延可能なオペレーター	ポーリング中にワーカースロットを解放できるため、リソース効率が向上します
1 回の実行内でタスクが多くの I/O 操作を並行して実行する	async def タスク／BaseAsyncOperator	非同期処理は高スループットの I/O 向けであり、クラッシュ復旧向けではありません
retries=0	—	再接続先がないため、クラッシュ復旧はできません
外部システムに追跡可能なジョブ ID がない（submit_job が None を返す）	通常のオペレーター	ミックスインのクラッシュ安全性の保証は暗黙のうちに無効になり、価値がありません
送信とポーリングがすでに 2 つの別々のタスクに分かれている（送信タスク → XCom → センサー／ポーリングタスク）	ポーリング／センサータスクには何も不要ですが、送信タスク自体には引き続き durable=True が必要です	ポーリング／センサータスクは何も送信しないため、リトライしてもすでに冪等であり、チェックポイントは不要です。一方、送信タスクは、それ自体がクラッシュに強くなければ、リトライ時に再送信する可能性があります。たとえば、DatabricksSubmitRunOperator.execute() は wait_for_termination=False の場合でも execute_resumable() を経由し、durable=True（デフォルト）で実行 ID をチェックポイントに記録します。このデフォルトは維持してください。フックを直接呼び出す手書きの送信タスク（通常の @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 に保存されています。名前を変更すると、ミックスインは実行中のすべてのリトライを新規送信として扱い、クラッシュ安全性の保証が損なわれます。

Before（アンチパターン）：

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)

After:

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

「submit + 別個のセンサータスク」という DAG 構成を1つのタスクにまとめるべきかどうかについては、**セクション6「Submit-and-poll DAGs: one task or two?」**を参照してください。これは、Triggerer の利用可否とジョブの実行時間に基づくアーキテクチャ上の判断です。クラッシュ安全性のために submit タスク自体に durable=True が必要かどうか（分割を維持するかどうかにかかわらず、通常は必要です）とは別の問題です。上の表の行を参照してください。

---

セクション6 — Submit-and-poll DAG：タスクは1つか2つか？

Databricks／Snowflake／BigQuery／Redshift／Spark のジョブでよく見られる DAG 構成は、wait_for_termination=False（または同等の設定）でジョブを起動してすぐに戻る submit タスクと、同じジョブの完了までポーリングする手作りのセンサータスクを後続させるものです。外部システムがジョブを受け付けた時点で submit タスクが成功状態になり、ジョブが完了した時点ではないことは、必ず指摘してください。submit タスクだけを対象にした再試行、アラート、SLA からは、実際のジョブの成功について何もわかりません。わかるのはセンサータスクの結果だけです。ただし、この分割をなくすべきかどうかは、Triggerer が利用可能かどうかによって異なります。

Triggerer がデプロイされている場合 — deferrable=True を指定して1つのタスクにまとめ、wait_for_termination=True も明示的に設定してください。deferrable=True だけでは wait_for_termination=True になるわけではありません。DatabricksSubmitRunOperator では、遅延処理のヘルパーは wait_for_termination も True の場合にのみ遅延します。古い submit タスクから wait_for_termination=False を引き継ぐと、遅延もポーリングも行われずにスキップされ、このセクションで解消しようとしているまさにその fire-and-forget の空白が再現されます。

これでもポーリング中にワーカーを解放でき、タスク自体の結果に実際のジョブ結果が反映されます。ただし、デメリットが一切ない、単純に優れたアップグレードではありません。遅延パスではジョブを直接送信し、mixin のチェックポイント処理を通らないため、ここでは durable=True は何も保護しません（下記の同期パスとは異なります）。送信後、トリガーが処理を引き継ぐ前にワーカーがクラッシュすると、再試行時に再送信されます。

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 がデプロイされていない場合 — これは自動的にどちらか一方を選ぶのではなく、実際にトレードオフがあります。ジョブの実行時間全体にわたってワーカーを確保することが許容できるかどうかによって、次の2つに分かれます。

• ジョブを同期的に実行しても問題ない（実行時間が短〜中程度、またはワーカー容量に余裕がある）うえ、オペレーターが耐久性のある実行に対応している場合（セクション5の ResumableJobMixin を継承し、durable=True を公開している） — 2タスク構成をオペレーターの耐久モードを使った1つのタスクに置き換えることを推奨します。設定は 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 とは無関係です。通常、ポーリングタスクと分割するかどうかにかかわらず、サブミットタスクでは durable=True を維持してください（セクション5の「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 に書き込むマップ済みタスクでは、インデックスごとに異なるキーを使用するか、最後の書き込みが優先される動作を受け入れる
• マップ済みタスクの場合、全体の状態クリアにはタスク本文内の clear() ではなく、下流タスクから CLI／コア API を使用する
• 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 の一般的なリファレンス。