Route pipeline failures to a fix or a person

A task failed overnight and the on-call engineer has to read the error, decide whether it was a blip worth rerunning, something a person has to fix now, or noise to ignore, and then act. This Dag hands the reading to a model that returns a pick and how sure it is. Airflow runs only the chosen branch, sends uncertain picks to a human with a deadline, and demands more confidence to page someone than to rerun a task.

What this demonstrates

Run it

  1. Create the decision_default connection for the decision model you run, TypeSafe Jev or a server answering the System One API (Decision models covers both). Both Dags below read it.

  2. Trigger the Dag:

    airflow dags test example_llm_branch_decision_policy

The triage_failure log shows the pick and its confidence. Exactly one of rerun, page_oncall and ignore runs. Below the confidence bar the run pauses for review instead (Airflow 3.1 or later); answer from Required Actions in the UI (see the note on What you can build).

The Dag

airflow/providers/common/ai/example_dags/example_llm_branch.py[source]

@dag(tags=["example"])
def example_llm_branch_decision_policy():
    # A decision model reports how sure it is of each pick; a text model does not, and
    # with a min_confidence set every pick would count as uncertain and go to review.
    route = LLMBranchOperator(
        task_id="triage_failure",
        prompt=(
            "Task load_orders failed: psycopg2.OperationalError: could not connect to server: "
            "Connection timed out. Is the server running on host db.internal (10.0.4.12)?"
        ),
        llm_conn_id="decision_default",
        system_prompt="Pick the remediation that addresses the cause of the failure.",
        branches={
            "rerun": "The failure looks transient: a timeout, a dropped connection, a rate limit.",
            # Paging someone on a wrong pick costs more than an extra rerun, so this branch needs more.
            "page_oncall": BranchOption(
                "Something a person has to fix now: data corruption, an outage, a security issue.",
                min_confidence=0.9,
            ),
            "ignore": "Expected or harmless: a known flaky check, a duplicate alert.",
        },
        decision_policy=DecisionPolicy(min_confidence=0.6, on_uncertain="review"),
        approval_timeout=timedelta(hours=4),
        allow_modifications=True,
    )

    @task
    def rerun():
        return "Clearing the failed task"

    @task
    def page_oncall():
        return "Paging on-call"

    @task
    def ignore():
        return "Leaving it"

    route >> [rerun(), page_oncall(), ignore()]

Classify-then-act variant

When the action depends on the score itself, classify in one task and act in the next (Decision models explains reading the score). It uses the same decision_default connection; run it as airflow dags test example_decision_model_confidence:

airflow/providers/common/ai/example_dags/example_decision_model.py[source]

class FailureCause(UseEnumMemberDocstrings, str, Enum):
    """Why an Airflow task failed."""

    # Each docstring is what the model weighs that option by. Some System One servers
    # refuse an option without one, so describe every option.
    transient = "transient"
    """A fault that clears by itself: a timeout, throttling, a dropped connection."""

    resource = "resource"
    """A dependency is down or unreachable and needs fixing before a retry can work."""

    permanent = "permanent"
    """A bug or bad input that fails the same way however often it is retried."""


@dag(tags=["example", "decision"])
def example_decision_model_confidence():
    """Classify a failure and escalate when the model says it does not know.

    ``LLMBranchOperator`` can gate its pick on one confidence bar through ``decision_policy``.
    Reading the confidence in a task of your own, as here, lets you act differently in more
    than one band.
    """

    @task
    def classify(log_line: str) -> dict:
        agent = PydanticAIHook(llm_conn_id="decision_default").create_agent(
            output_type=FailureCause,
            instructions="Classify why this Airflow task failed.",
        )
        result = agent.run_sync(log_line)
        # Confidence is reported per output field; a bare output type lands under
        # "response". A bounded ``float`` output would report none at all -- there the
        # probability is the answer -- so this ``or 0.0`` would read as no confidence
        # rather than as a missing one. ``LLMOperator`` and ``LLMBranchOperator`` push the
        # same confidence in their ``decision`` XCom.
        details = result.response.provider_details or {}
        return {
            "category": result.output.value,
            "confidence": (details.get("confidence") or {}).get("response"),
        }

    @task
    def act(classification: dict) -> str:
        confidence = classification["confidence"] or 0.0
        if confidence >= ACT_ABOVE:
            return f"Remediating {classification['category']} automatically"
        if confidence >= REVIEW_ABOVE:
            return f"Filing {classification['category']} for review"
        return "Paging a human: the classification was a coin flip"

    act(classify(UNDERDETERMINED_INCIDENT))

Adapting it

  • Replace the hard-coded prompt with the failed task’s own error. An on_failure_callback on the production Dag can trigger this one with the exception text in conf, and the prompt reads {{ dag_run.conf["error"] }}.

  • Make the branch tasks do the work: rerun clears the failed task instance through the REST API, page_oncall posts to your alerting provider.

  • Tune min_confidence per branch. A wrong page costs more than an extra rerun, so page_oncall demands more certainty.

  • To decide retry versus fail inside Airflow’s retry loop, see Retry policies.

Was this entry helpful?