Arize AX Airflow Provider

The official Apache Airflow provider for Arize AX — schedule, automate, and orchestrate your LLMOps workflows directly from Airflow.

Build DAGs that evaluate prompts continuously, compare experiments before deploying, detect drift, curate datasets from production traces, and gate releases on evaluation scores, all using purpose-built operators that wrap the Arize AX platform.

Features

• 107 operators across 15 domains: datasets, experiments, projects, spans, traces, evaluators, prompts, tasks, annotations, integrations, AI integrations, API keys, spaces, resource restrictions, ML
• 8 sensors for waiting on dataset readiness, experiment completion, span ingestion, evaluation scores, task runs, and more
• Built-in gates fail_on_regression=True, fail_on_drift=True, min_score thresholds raise AirflowException directly, no glue code needed
• Idempotent operations if_exists="skip" on creates, ignore_if_missing=True on deletes
• Clean XCom values operators return scalar IDs by default, full results available via named XCom keys
• Continuous evaluation tasks with Eval Hub for live production monitoring
• Human-in-the-loop annotations through annotation queues
• Cloud storage out of the box — every read-side and write-side operator (Parquet / JSONL / dataset examples / experiment runs / annotated spans / ML exports) accepts path="s3://...", "gs://...", "abfs://..." URIs. Pass credentials inline via storage_options= or by name via cloud_conn_id="aws_default" (Airflow Connection). Local paths bypass cloud code paths entirely.
• 23 example DAGs covering CI/CD gates, drift detection, prompt lifecycle, RAG evaluation, fine-tuning data pipelines, a self-optimizing loop, and a cloud-export demo (MinIO/S3/GCS/Azure)

Installation

pip install arize-ax-airflow-provider

To enable cloud-storage targets on write-side operators (s3://, gs://, abfs://), install the matching extra(s):

pip install 'arize-ax-airflow-provider[s3]'        # adds fsspec + s3fs
pip install 'arize-ax-airflow-provider[gcs]'       # adds fsspec + gcsfs
pip install 'arize-ax-airflow-provider[azure]'     # adds fsspec + adlfs
pip install 'arize-ax-airflow-provider[cloud]'     # all three at once

Local paths work without any extra. See the Cloud storage section below for the storage_options shape (or test against MinIO with zero cloud credentials via example_arize_ax_cloud_export_dag.py).

Requires Python 3.10+, Apache Airflow 3.0+ (tested against 3.0–3.2), and Arize SDK 8.50.0 (the provider exact-pins arize==8.50.0 so installs only use the version we test against; bump alongside the provider release when adopting a newer SDK).

Setup

1. In Airflow UI, go to Admin → Connections → Add
2. Set Connection Id to arize_ax_default
3. Set Connection Type to arize_ax
4. Set Password to your Arize API key
5. (Optional) Set Extra to {"space_id": "your-space-id"} to use it as the default space

Recognised Extra keys: space_id / default_space, api_host, region, api_scheme, max_past_years, and connection_pool_maxsize. Set connection_pool_maxsize when a worker runs many Arize tasks concurrently — it is forwarded to ArizeClient (SDK 8.50.0) and avoids urllib3 "Connection pool is full" warnings and the serialized retries that follow.

Quick start

A server-side CI/CD gate: pull a candidate prompt from Prompt Hub, have Arize Eval Hub run it against an evaluation dataset (no LLM SDK or API key on the Airflow worker), compare against the production baseline, and promote on pass. Requires four Airflow Variables: arize_ax_prompt_name, arize_ax_eval_dataset_id, arize_ax_baseline_experiment_id, arize_ax_ai_integration_id.

from datetime import datetime
from typing import Any

from airflow import DAG
from airflow.models import Variable
from airflow.providers.standard.operators.python import PythonOperator

from airflow.providers.arize_ax.operators.experiments import (
    ArizeAxCompareExperimentsOperator,
)
from airflow.providers.arize_ax.operators.prompts import (
    ArizeAxGetPromptOperator,
    ArizeAxPromotePromptOperator,
)
from airflow.providers.arize_ax.operators.tasks import (
    ArizeAxCreateRunExperimentTaskOperator,
    ArizeAxGetTaskRunOperator,
    ArizeAxTriggerTaskRunOperator,
)
from airflow.providers.arize_ax.sensors.arize_ax import ArizeAxTaskRunSensor


def build_run_config_from_prompt(**ctx) -> dict[str, Any]:
    """Materialize a Prompt Hub prompt version into a server-side run config."""
    prompt = ctx["ti"].xcom_pull(task_ids="fetch_candidate_prompt")
    version = prompt["version"]
    return {
        "experiment_type": "llm_generation",
        "ai_integration_id": Variable.get("arize_ax_ai_integration_id"),
        "model_name": version["model"],
        "messages": version["messages"],
        "input_variable_format": version["input_variable_format"],
        "invocation_parameters": version.get("invocation_params") or {},
        "provider_parameters": version.get("provider_params") or {},
    }


with DAG(
    dag_id="llm_cicd_gate",
    start_date=datetime(2026, 1, 1),
    schedule="@daily",
    catchup=False,
    render_template_as_native_obj=True,  # required: dict XCom must stay a dict
) as dag:

    fetch_candidate_prompt = ArizeAxGetPromptOperator(
        task_id="fetch_candidate_prompt",
        prompt_name="{{ var.value.arize_ax_prompt_name }}",
        version_label="staging",
    )

    build_run_config = PythonOperator(
        task_id="build_run_config",
        python_callable=build_run_config_from_prompt,
    )

    create_task = ArizeAxCreateRunExperimentTaskOperator(
        task_id="create_candidate_task",
        name="candidate-{{ ds_nodash }}",
        dataset_id="{{ var.value.arize_ax_eval_dataset_id }}",
        run_configuration="{{ ti.xcom_pull(task_ids='build_run_config') }}",
        if_exists="skip",
    )

    trigger_run = ArizeAxTriggerTaskRunOperator(
        task_id="trigger_candidate_run",
        task_id_param="{{ ti.xcom_pull(task_ids='create_candidate_task') }}",
        experiment_name="candidate-{{ ds_nodash }}",
    )

    wait_for_run = ArizeAxTaskRunSensor(
        task_id="wait_for_candidate_run",
        run_id="{{ ti.xcom_pull(task_ids='trigger_candidate_run') }}",
        poke_interval=15,
        timeout=900,
        mode="reschedule",
    )

    get_result = ArizeAxGetTaskRunOperator(
        task_id="get_candidate_result",
        run_id="{{ ti.xcom_pull(task_ids='trigger_candidate_run') }}",
    )

    compare = ArizeAxCompareExperimentsOperator(
        task_id="compare_to_baseline",
        baseline_experiment_id="{{ var.value.arize_ax_baseline_experiment_id }}",
        candidate_experiment_id="{{ ti.xcom_pull(task_ids='get_candidate_result')['experiment_id'] }}",
        metric_names=["accuracy"],
        pass_threshold=0.0,
        fail_on_regression=True,
    )

    promote = ArizeAxPromotePromptOperator(
        task_id="promote_to_production",
        prompt_name="{{ var.value.arize_ax_prompt_name }}",
        label="production",
    )

    fetch_candidate_prompt >> build_run_config >> create_task >> trigger_run
    trigger_run >> wait_for_run >> get_result >> compare >> promote

That's a complete server-side CI/CD gate. The fail_on_regression=True flag means compare raises AirflowException when the candidate underperforms the baseline, so promote never runs. The same arize_ax_prompt_name Variable threads through fetch (at the staging label) and promote (to the production label), the gate is the only path that re-labels production.

Prefer running the LLM call yourself? Swap the server-side stack (Get/Create/Trigger/Wait/Get) for a single ArizeAxRunExperimentOperator with a task= callable. See example_arize_ax_llm_cicd_gate_dag.py for the client-side variant.

Operator domains

Domain	Operators
Datasets	List, create, get, delete, list/append/annotate/update/delete examples, export to file, health check, smart refresh
Experiments	List, create, get, delete, run, list/annotate runs, get score, compare, detect drift, calibration, behavioral regression, budget allocator
Projects	List, create, get, delete
Spans	List (with server-side column projection), delete (with time window), log, update evaluations/annotations/metadata, export to DataFrame/Parquet, get metrics, curate to dataset, curate feedback dataset (self-learning agents), export annotated, export to fine-tuning, adaptive sampling
Evaluators	List, create (template or code), get, update, delete, list/get/add versions (template or code)
Prompts	List, create, get, delete, compare, promote, optimize (meta-prompt via Prompt Learning SDK)
Tasks	List, create (evaluation or run-experiment), get, update, delete, list runs, get run, trigger run, cancel run
Annotations	List/create/delete configs, list/get/create/update/delete queues, list/add/delete records, annotate, assign
Traces	List (trace-level listing via client.traces)
Integrations	List, get, create LLM, create agent, update LLM, update agent, delete — successor to AI Integrations, covers LLM and AGENT types
AI Integrations	List, get, create, update, delete (superseded by Integrations)
API Keys	List, create, delete, refresh
Spaces	List, get, create, update (both accept is_private), delete
Resource Restrictions	List — which resources are access restricted
ML	Log batch/stream, export to DataFrame/Parquet

Sensors

Sensor	Purpose
ArizeAxExperimentCompleteSensor	Wait until an experiment reaches a terminal state
ArizeAxDatasetReadySensor	Wait until a dataset has at least N examples
ArizeAxSpanCountSensor	Wait until span count in a project exceeds a threshold
ArizeAxEvaluationScoreSensor	Wait until evaluation score crosses a threshold
ArizeAxExperimentRunCountSensor	Wait until experiment has N runs
ArizeAxSpanIngestionSensor	Wait until span ingestion stabilizes
ArizeAxAnnotationQueueSensor	Wait until annotation queue is configured
ArizeAxTaskRunSensor	Wait until a task run reaches a terminal state

Design patterns

The provider follows established Airflow operator conventions so DAGs read naturally and stay maintainable:

• Idempotent creates Set if_exists="skip" to handle 409 conflicts by resolving the existing resource by name
• Idempotent deletes All Delete operators accept ignore_if_missing=True (default), logging on 404 instead of raising
• Built-in gates Comparison operators (CompareExperiments, DetectEvalDrift, EvaluatorCalibration, BehavioralRegression, GetExperimentScore) accept fail_on_* / min_score params that raise AirflowException on failure
• Param validation Operators validate required space_id / project_id in execute() with clear error messages
• Convenience XCom keys List operators push first_id and first_name for direct chaining via Jinja templates
• Override evaluations ArizeAxTriggerTaskRunOperator(override_evaluations=True) re-evaluates spans that already have labels

Example DAGs

Bundled in provider_pkg/example_dags/:

Pattern	DAG
Self-contained smoke test	example_arize_ax_e2e_dag.py
LLM CI/CD gate	example_arize_ax_llm_cicd_gate_dag.py
Prompt lifecycle (staging → production)	example_arize_ax_prompt_lifecycle_dag.py
Prompt A/B testing	example_arize_ax_prompt_ab_test_dag.py
Drift detection with auto-rollback	example_arize_ax_drift_detection_dag.py
Behavioral regression detection	example_arize_ax_behavioral_regression_dag.py
Evaluator calibration vs human labels	example_arize_ax_evaluator_calibration_dag.py
RAG evaluation pipeline	example_arize_ax_rag_evaluation_dag.py
Production span curation into datasets	example_arize_ax_dataset_curation_dag.py
Fine-tuning data pipeline	example_arize_ax_finetune_data_pipeline_dag.py
Continuous evaluation tasks	example_arize_ax_tasks_dag.py
Annotation queues for HITL	example_arize_ax_annotation_queues_dag.py
Multi-model experiment matrix	example_arize_ax_llm_experiments_dag.py
Self-learning agent (multi-prompt optimization from production feedback)	example_arize_ax_prompt_optimization_with_feedback_dag.py
Self-optimizing loop (self-contained closed loop: baseline → optimize → candidate → gate → promote)	example_arize_ax_self_optimizing_loop_dag.py
Cloud-export demo (export spans to S3 / GCS / Azure or MinIO via storage_options)	example_arize_ax_cloud_export_dag.py
Platform surface (traces, span column projection, integrations, resource restrictions)	example_arize_ax_platform_surface_dag.py

Plus dataset, experiment, evaluator, span, project, ML, and admin demos for individual domain walkthroughs.

Cloud storage

Every operator that takes a path / output_path / examples_path / experiment_runs_path / dataframe_path accepts either a local filesystem path or a remote URI (s3://, gs://, gcs://, abfs://, az://, azure://). This covers write-side exporters and read-side ingestion ops alike:

Direction	Operators
Write	ArizeAxSpansExportToParquetOperator, ArizeAxExportSpansToFineTuningOperator, ArizeAxExportAnnotatedSpansOperator, ArizeAxExportDatasetExamplesToFileOperator, ArizeAxExportExperimentRunsToFileOperator, ArizeAxMLExportToParquetOperator, ArizeAxAdaptiveSamplingOperator (output side)
Read	ArizeAxCreateDatasetOperator, ArizeAxAppendDatasetExamplesOperator, ArizeAxCreateExperimentOperator, ArizeAxMLLogBatchOperator, ArizeAxSpansLogOperator, ArizeAxUpdateSpansEvaluationsOperator, ArizeAxUpdateSpansAnnotationsOperator, ArizeAxUpdateSpansMetadataOperator

Local paths bypass cloud code entirely (no fsspec import, no overhead). For remote URIs the operator writes to a local temp file first and uploads via fsspec (write-side), or downloads to a temp file the SDK can consume (read-side) — multipart / parallel transfers and resumable uploads are handled by the protocol backend (s3fs, gcsfs, adlfs). Pandas-native paths (CSV / parquet) skip the local hop and stream through fsspec directly.

Two ways to pass credentials, both supported on every operator:

# 1) Inline storage_options (raw fsspec kwargs — most explicit, useful for
#    one-off DAG-level credentials or local emulators).
ArizeAxSpansExportToParquetOperator(
    task_id="export_to_s3",
    path="s3://my-bucket/spans/{{ ds_nodash }}.parquet",
    storage_options={
        "key": "{{ var.value.aws_access_key_id }}",
        "secret": "{{ var.value.aws_secret_access_key }}",
    },
    space_id="...", project_name="...", start_time="...", end_time="...",
)

# 2) cloud_conn_id (idiomatic for production — credentials live in Airflow's
#    encrypted Connection store; one place to rotate, audited, no Variables).
ArizeAxCreateDatasetOperator(
    task_id="load_examples_from_s3",
    name="prod-eval-{{ ds_nodash }}",
    examples_path="s3://my-bucket/datasets/eval.json",
    cloud_conn_id="aws_default",   # or "google_cloud_default", "azure_default", etc.
)

The resolver maps any aws, google_cloud_platform, wasb / azure connection's credentials and extras onto the right fsspec backend kwargs (key/secret/client_kwargs for s3fs, token/project for gcsfs, account_name/account_key/sas_token for adlfs). Unknown conn_types fall back to passing the connection extras dict through verbatim. When both are set, storage_options wins as an explicit override.

Testing without any cloud account. Run MinIO standalone (S3-API compatible) and point storage_options.client_kwargs.endpoint_url at it:

docker run -d --name minio -p 9000:9000 -p 9001:9001 \
  -e MINIO_ROOT_USER=minioadmin -e MINIO_ROOT_PASSWORD=minioadmin \
  minio/minio server /data --console-address ":9001"

storage_options={
    "key": "minioadmin", "secret": "minioadmin",
    "client_kwargs": {"endpoint_url": "http://host.docker.internal:9000"},
}

See example_arize_ax_cloud_export_dag.py for a full self-contained demo: export spans → readback through fsspec → stage examples back into the bucket → load them via ArizeAxCreateDatasetOperator (round-trip on the same cloud bucket).

The self-learning agent demo requires the prompt-learning-enhanced SDK (available as a git source only, PyPI does not accept direct-URL deps, so it is not declared as an optional extra here). Install it separately on workers that run the optimization DAG:

pip install 'arize-phoenix-evals>=2.0,<3.0' \
            'prompt-learning-enhanced @ git+https://github.com/Arize-ai/prompt-learning.git'

The arize-phoenix-evals<3.0 pin is required because upstream prompt-learning-enhanced imports phoenix.evals.models, which was removed in 3.0.0. ArizeAxOptimizePromptOperator raises a clear AirflowException with this exact install line when the SDK is missing.

Documentation

• Provider guide: Arize AX Airflow Provider — install, setup, design patterns, example DAGs
• Operator reference: Operators, sensors, and hooks — full reference for all 107 operators
• Arize AX docs: arize.com/docs/ax
• Arize SDK reference: Arize Python SDK v8

Support

• Help & support: arize.com/support
• Slack: Arize Community Slack
• Twitter/X: @ArizeAI