#1725 Airflow agent Pull request opened by pingsutw TL;DR Airflow agent allows you to seamlessly run Airflow tasks in the Flyte workflow without changing code. All the airflow tasks will be run on an airflow agent (long-running server) instead of launching a new pod to run it, which dramatically reduces overhead. • Compile Airflow tasks to Flyte tasks • Use Airflow sensors/operators in Flyte workflows • Add support running Airflow tasks locally without running a cluster Type ☐ Bug Fix ☑︎ Feature ☐ Plugin Are all requirements met? ☐ Code completed ☐ Smoke tested ☐ Unit tests added ☐ Code documentation added ☐ Any pending items have an associated Issue Complete description DataprocOperator

from datetime import timedelta

from airflow.utils import trigger_rule

from flytekit import task, workflow, ImageSpec
from airflow.sensors import time_sensor
from airflow.providers.google.cloud.operators.dataproc import DataprocCreateClusterOperator, DataprocDeleteClusterOperator, DataprocSubmitSparkJobOperator

airflow_plugin = "git+<https://github.com/flyteorg/flytekit.git@487438ab59147879eded897674593a1eaee1c78b#subdirectory=plugins/flytekit-airflow>"
image_spec = ImageSpec(base_image="pingsutw/flytekit:v1", packages=["apache-airflow", airflow_plugin], apt_packages=["git"], registry="pingsutw")


@workflow
def wf():
    create_cluster = DataprocCreateClusterOperator(
        task_id="create_dataproc_cluster",
        image_version="2.0.27-debian10",
        storage_bucket="opta-gcp-dogfood-gcp",
        master_machine_type="n1-highmem-32",
        master_disk_size=1024,
        num_workers=2,
        worker_machine_type="n1-highmem-64",
        worker_disk_size=1024,
        region="us-west1",
        cluster_name="flyte-dataproc",
        project_id="dogfood-gcp-dataplane",
    )

    run_spark = DataprocSubmitSparkJobOperator(
        job_name="spark_pi",
        task_id="run_spark",
        dataproc_jars=["file:///usr/lib/spark/examples/jars/spark-examples.jar"],
        main_class="org.apache.spark.examples.JavaWordCount",
        arguments=["<gs://opta-gcp-dogfood-gcp/spark/file.txt>"],
        cluster_name="flyte-dataproc",
        region="us-west1",
        project_id="dogfood-gcp-dataplane",
    )

    delete_cluster = DataprocDeleteClusterOperator(
        task_id="create_dataproc_cluster",
        project_id="dogfood-gcp-dataplane",
        cluster_name="flyte-dataproc",
        region="us-west1",
        retries=3,
        retry_delay=timedelta(minutes=5),
        email_on_failure=True,
        trigger_rule=trigger_rule.TriggerRule.ALL_DONE
    )

    create_cluster >> run_spark >> delete_cluster


if __name__ == '__main__':
    wf()

bashSensor

from airflow.sensors.bash import BashSensor
from flytekit import task, workflow, ImageSpec


airflow_plugin = "git+<https://github.com/flyteorg/flytekit.git@487438ab59147879eded897674593a1eaee1c78b#subdirectory=plugins/flytekit-airflow>"
image_spec = ImageSpec(base_image="pingsutw/flytekit:v1", packages=["apache-airflow", airflow_plugin], apt_packages=["git"], registry="pingsutw")


@task(container_image=image_spec)
def t1():
    print("flyte")


@workflow
def wf():
    sensor = BashSensor(task_id="Sensor_succeeds", bash_command="echo hello")
    sensor >> t1()


if __name__ == '__main__':
    wf()

HiveSensor

from airflow.providers.apache.hive.sensors.hive_partition import HivePartitionSensor

from flytekit import task, workflow, ImageSpec


airflow_plugin = "git+<https://github.com/flyteorg/flytekit.git@487438ab59147879eded897674593a1eaee1c78b#subdirectory=plugins/flytekit-airflow>"
image_spec = ImageSpec(base_image="pingsutw/flytekit:v1", packages=["apache-airflow", airflow_plugin], apt_packages=["git"], registry="pingsutw")


@task(container_image=image_spec)
def t1():
    print("flyte")


@workflow
def wf():
    sensor = HivePartitionSensor(table="flyte", schema="person")
    sensor >> t1()


if __name__ == '__main__':
    wf()

FileSensor

from airflow.sensors.filesystem import FileSensor
from flytekit import task, workflow, ImageSpec


airflow_plugin = "git+<https://github.com/flyteorg/flytekit.git@487438ab59147879eded897674593a1eaee1c78b#subdirectory=plugins/flytekit-airflow>"
image_spec = ImageSpec(base_image="pingsutw/flytekit:v1", packages=["apache-airflow", airflow_plugin], apt_packages=["git"], registry="pingsutw")


@task(container_image=image_spec)
def t1():
    print("flyte")


@workflow
def wf():
    sensor = FileSensor(task_id="id", filepath="/tmp/1234")
    sensor >> t1()


if __name__ == '__main__':
    wf()

TimeSensor

from datetime import datetime, timedelta
from pytz import UTC

from airflow.sensors.time_sensor import TimeSensor
from flytekit import task, workflow, ImageSpec

airflow_plugin = "git+<https://github.com/flyteorg/flytekit.git@487438ab59147879eded897674593a1eaee1c78b#subdirectory=plugins/flytekit-airflow>"
image_spec = ImageSpec(base_image="pingsutw/flytekit:v1", packages=["apache-airflow", airflow_plugin], apt_packages=["git"], registry="pingsutw")


@task(container_image=image_spec)
def t1():
    print("flyte")


@workflow
def wf():
    sensor = TimeSensor(task_id="fire_immediately", target_time=(datetime.now(tz=UTC)+timedelta(seconds=5)).time())
    sensor >> t1()


if __name__ == '__main__':
    wf()

PythonSensor

from airflow.sensors.python import PythonSensor
from flytekit import task, workflow, ImageSpec

airflow_plugin = "git+<https://github.com/flyteorg/flytekit.git@487438ab59147879eded897674593a1eaee1c78b#subdirectory=plugins/flytekit-airflow>"
image_spec = ImageSpec(base_image="pingsutw/flytekit:v1", packages=["apache-airflow", airflow_plugin], apt_packages=["git"], registry="pingsutw")


@task(container_image=image_spec)
def t1():
    print("flyte")


def py_func():
    print("airflow python sensor")
    return True


@workflow
def wf():
    sensor = PythonSensor(task_id="fire_immediately", python_callable=py_func)
    sensor >> t1()


if __name__ == '__main__':
    wf()

Tracking Issue NA Follow-up issue NA flyteorg/flytekit GitHub Actions: build-plugins (3.11, flytekit-sqlalchemy) GitHub Actions: build-plugins (3.11, flytekit-snowflake) GitHub Actions: build-plugins (3.11, flytekit-polars) GitHub Actions: build-plugins (3.11, flytekit-papermill) GitHub Actions: build-plugins (3.11, flytekit-pandera) GitHub Actions: build-plugins (3.11, flytekit-kf-tensorflow) GitHub Actions: build-plugins (3.11, flytekit-kf-pytorch) GitHub Actions: build-plugins (3.11, flytekit-kf-mpi) GitHub Actions: build-plugins (3.11, flytekit-k8s-pod) GitHub Actions: build-plugins (3.11, flytekit-huggingface) GitHub Actions: build-plugins (3.11, flytekit-hive) GitHub Actions: build-plugins (3.11, flytekit-envd) GitHub Actions: build-plugins (3.11, flytekit-duckdb) GitHub Actions: build-plugins (3.11, flytekit-dolt) GitHub Actions: build-plugins (3.11, flytekit-deck-standard) GitHub Actions: build-plugins (3.11, flytekit-dbt) GitHub Actions: build-plugins (3.11, flytekit-data-fsspec) GitHub Actions: build-plugins (3.11, flytekit-dask) GitHub Actions: build-plugins (3.11, flytekit-bigquery) GitHub Actions: build-plugins (3.11, flytekit-aws-sagemaker) GitHub Actions: build-plugins (3.11, flytekit-aws-batch) GitHub Actions: build-plugins (3.11, flytekit-aws-athena) GitHub Actions: build-plugins (3.10, flytekit-whylogs) GitHub Actions: build-plugins (3.10, flytekit-vaex) GitHub Actions: build-plugins (3.10, flytekit-sqlalchemy) GitHub Actions: build-plugins (3.10, flytekit-spark) GitHub Actions: build-plugins (3.10, flytekit-ray) GitHub Actions: build-plugins (3.10, flytekit-polars) GitHub Actions: build-plugins (3.10, flytekit-papermill) ✅ 1 other check has passed 1/30 successful checks