Airflow Task State Store (AIP-103)

Airflow 3.3 intègre deux magasins clé/valeur et un mixin de sécurité contre les crashes pour les opérateurs qui soumettent des jobs externes.

Requiert Airflow 3.3+. Vérifiez d'abord :

af config version

Si la version est inférieure à 3.3, indiquez à l'utilisateur que ces fonctionnalités ne sont pas encore disponibles et renvoyez-le au problème de suivi AIP-103.

Section 1 — Choisir la bonne primitive

Quand NE PAS utiliser ces magasins :

• Passer des données entre tasks → utiliser XCom
• Grandes charges utiles (poids de modèle, dataframes) → utiliser XCom avec un backend de stockage d'objets
• Configuration ou secrets partagés entre DAGs → utiliser Variables ou Connections

Section 2 — Détecter les anti-patterns dans les DAGs existants (à la demande)

Quand l'utilisateur demande à revoir un DAG ou demande « y a-t-il une meilleure façon », scannez ces patterns et signalez-les :

Montrez un snippet avant/après quand vous signalez. Utilisez les exemples canoniques aux sections 3–5 comme « après ».

Section 3 — task_state_store : état de coordination par task

task_state_store est un magasin clé/valeur limité à une identité unique de task instance (dag_id + run_id + task_id + map_index). Il survit aux tentatives — une nouvelle tentative sur la même task lit le même magasin.

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"]  # injecté par Airflow, aucune configuration nécessaire
        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 après chaque enregistrement

    process_records()

etl_with_checkpoint()

API :

from airflow.sdk import NEVER_EXPIRE

task_state_store.get(key, default=None)                        # retourne une JsonValue ou la valeur par défaut
task_state_store.set(key, value)                               # utilise default_retention_days
task_state_store.set(key, value, retention=timedelta(days=7))  # override TTL par clé
task_state_store.set(key, value, retention=NEVER_EXPIRE)       # n'expire jamais indépendamment de la config
task_state_store.delete(key)                                   # no-op si la clé n'existe pas
task_state_store.clear()                                       # supprime toutes les clés de cette instance de task

Règles clés :

• Les valeurs doivent être sérialisables en JSON (str, int, float, bool, list, dict — les valeurs None sont rejetées).
• L'expiration par défaut est contrôlée par [state_store] default_retention_days (0 = n'expire jamais).
• Utilisez NEVER_EXPIRE pour les clés qui doivent survivre à la fenêtre de rétention par défaut (p. ex. un ID de job pour un job Spark multi-jours).
• La taille max des valeurs est 64 KB par défaut ; configurable via [state_store] max_value_storage_bytes (0 = pas de limite). Pour les grandes charges utiles, configurez un backend personnalisé [state_store] backend ou un backend côté worker configuré via : [workers] state_store_backend.

Tasks mappées — chaque index a son propre espace de noms :

Quand une task est mappée dynamiquement (task.expand(...)), chaque index de map obtient un task_state_store isolé limité à son propre map_index. Les indices ne partagent pas l'état.

@task(retries=2)
def process_partition(partition_id, **context):
    task_state_store = context["task_state_store"]
    # Limité à CET index seulement — les autres indices ont leur propre copie
    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() efface seulement l'index actuel. Pour effacer l'état entre tous les indices de map d'un groupe de tasks, utilisez le CLI ou l'API core.

Avant (anti-pattern) :

@task
def process(**context):
    cursor = Variable.get("etl_cursor", default_var=0)
    # ... traiter ...
    Variable.set("etl_cursor", new_cursor)  # global, n'importe quelle task peut le réécrire

Après :

@task(retries=3)
def process(**context):
    task_state_store = context["task_state_store"]
    cursor = task_state_store.get("cursor", default=0)
    # ... traiter ...
    task_state_store.set("cursor", new_cursor)    # limité à cette instance de task

Section 4 — asset_state_store : métadonnées par asset entre les runs de DAG

asset_state_store est limité à un asset, pas à une instance de task. Il persiste entre les runs de DAG — la même clé sur le même asset est lisible et inscriptible par n'importe quelle task qui le produit ou le consomme.

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 injecté par Airflow — déclarez comme 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()

Lire le magasin depuis un DAG consommateur :

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

Règles clés :

• asset_state_store est injecté par Airflow comme un kwarg nommé — déclarez-le comme def my_task(asset_state_store=None). Ne combinez PAS avec **context ; Airflow l'injecte séparément.
• Utilisez datetime.now(tz=timezone.utc).isoformat() pour les timestamps — jamais datetime.utcnow() (pas conscient du fuseau horaire).
• Même contrainte de valeur sérialisable JSON que task_state_store.
• Pas d'expiration par clé — les entrées du magasin d'état asset n'ont pas de TTL (l'asset survive à n'importe quel run unique).
• Lisible par n'importe quel DAG qui déclare l'asset comme inlet ou outlet.

Tasks mappées — le dernier écrivain gagne :

asset_state_store est limité à l'asset, pas à l'index de map. Si plusieurs indices mappés écrivent la même clé simultanément, le dernier écrivain gagne. Utilisez des clés distinctes par index ou assurez-vous qu'un seul index écrit à une clé donnée.

@task(outlets=[my_asset])
def load_partition(partition_id, asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    # Clé distincte par index — pas de condition de course
    asset_state_store.set(f"offset_{partition_id}", new_offset)

Avant (anti-pattern) :

Variable.set(f"watermark_{asset_name}", new_offset)   # global, pas limité à l'asset

Après :

@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 : soumission de job externe crash-safe

À utiliser quand un opérateur soumet un job à un système externe (Spark, Databricks, dbt Cloud, AWS Batch, etc.) puis vote pour achèvement. Sans ce mixin, un crash worker pendant le vote signifie que la tentative suivante soumet un job dupliqué.

Quand NE PAS utiliser ResumableJobMixin :

ResumableJobMixin retient le slot worker pendant toute la durée du vote — identique à un opérateur synchrone standard. Le bénéfice est la sécurité contre les crashes et la continuité du job, pas l'efficacité des ressources.

Refuser la récupération de crash :

Le mixin est livré avec durable=True par défaut. Définissez durable=False pour ignorer toute interaction task_state_store et exécuter un cycle soumission/vote/résultat ordinaire — utile dans les environnements de test ou quand le système externe a sa propre dédup :

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

# Ou via default_args pour désactiver pour toutes les tasks d'un DAG :
with DAG("my_dag", default_args={"durable": False}):
    ...

Implémenter le mixin

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


class MyBatchOperator(BaseOperator, ResumableJobMixin):

    external_id_key = "batch_job_id"   # clé utilisée dans task_state_store ; défini une fois, jamais renommé

    def execute(self, context):
        return self.execute_resumable(context)  # jamais appeler self.execute() — appeler ceci

    def submit_job(self, context) -> JsonValue:
        # Soumettre et retourner l'identifiant du job. Cette valeur est persistée à task_state_store
        # avant le début du vote. Retourner None seulement si le système n'a pas d'ID traçable
        # (dans ce cas la sécurité contre les crashes est désactivée et le job est resoumis à chaque tentative).
        return self.hook.submit_batch(...)

    def get_job_status(self, external_id: JsonValue, context) -> str:
        # Interroger le système externe. Retourner une chaîne d'état brute.
        return self.hook.get_status(external_id)

    def is_job_active(self, status: str) -> bool:
        # Retourner True si le job est toujours en cours et doit être reconnecté.
        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:
        # Bloquer jusqu'à ce que le job atteigne un état terminal. Lever en cas d'échec.
        self.hook.wait(external_id)

    def get_job_result(self, external_id: JsonValue, context):
        # Retourner le résultat du job après succès. Retourner None si non applicable.
        return None

Que se passe-t-il en cas de tentative

Avertissement external_id_key

Ne jamais renommer external_id_key sur un opérateur déjà déployé avec des instances de task en vol. L'ancienne clé est stockée dans task_state_store sous le nom précédent. Un renommage fait que le mixin traite chaque tentative active comme une nouvelle soumission, déjouant la garantie de sécurité contre les crashes.

Avant (anti-pattern) :

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, sujet aux conditions de course
        self._wait(job_id)

Après :

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(...)
    # ... implémenter les 5 autres méthodes ...

Section 6 — Référence de configuration

[state_store]
# Chemin pointillé complet vers le backend de stockage. Par défaut écrit dans la BD de métadonnées Airflow.
backend = airflow.state.metastore.MetastoreStateStoreBackend

# Jours de rétention des entrées du magasin d'état task après leur dernière mise à jour. 0 = désactiver le nettoyage basé sur le temps.
# N'affecte PAS les lignes asset_state_store — le magasin d'état asset n'a pas de TTL.
default_retention_days = 30

# Lignes supprimées par lot pendant le nettoyage. 0 = pas de batching (suppression unique non bornée).
# Affinez sur les grands déploiements pour réduire la contention de verrous.
state_cleanup_batch_size = 0

# Supprimer automatiquement toutes les clés du magasin d'état task quand une task réussit. Par défaut : False.
# N'affecte PAS asset_state_store — le magasin d'état asset persiste entre les runs et doit être effacé explicitement.
clear_on_success = False

Backend côté worker (optionnel, section [workers]) — achemine les écritures du magasin d'état task via un backend local avant qu'elles n'atteigne le serveur API. Utile quand les grandes charges utiles ou le stockage accrédité doivent rester sur le worker :

[workers]
state_store_backend = mypackage.store.WorkerSideBackend

Section 7 — Checklist de sécurité

• [ ] Version Airflow ≥ 3.3 (af config version)
• [ ] Les valeurs sont sérialisables JSON (str, int, float, bool, list, dict — pas de datetime, pas d'objets personnalisés)
• [ ] Les clés task_state_store sont de courtes chaînes descriptives (évitez les points et barres obliques)
• [ ] Tasks mappées écrivant à asset_state_store : utilisez des clés distinctes par index ou acceptez la sémantique dernier-écrivain-gagne
• [ ] Tasks mappées : l'effacement d'état à l'échelle de la flotte utilise CLI/API core depuis une task aval, pas clear() dans le corps de la task
• [ ] ResumableJobMixin : external_id_key est défini et ne sera pas renommé après déploiement
• [ ] ResumableJobMixin : execute() appelle self.execute_resumable(context), pas une logique personnalisée
• [ ] ResumableJobMixin : durable=False est intentionnel si la récupération de crash est désactivée
• [ ] Les grandes charges utiles (> configuré max_value_storage_bytes) utilisent un backend personnalisé [state_store] backend ou un backend côté worker configuré via : [workers] state_store_backend

Compétences connexes

• authoring-dags — patterns et conventions généraux de rédaction de DAG.
• airflow-hitl — pause d'un DAG pour approbation humaine (Airflow 3.1+).
• airflow — af config, af registry et référence CLI Airflow générale.