Support for PythonOperatorPerArtifact flavor in Airflow DAG file

[yoonspark]

Overview

Our team has been working on supporting new graph refactor for pipeline file generation such that duplicate code blocks are factored out to reduce redundant computations (which might be expensive). Our previous implementation (relevant discussion here) was limited in that it produced Airflow DAG file that has operators (i.e., tasks) at the session level: each session may contain multiple artifacts and we often want individual artifact-level control/observability, which demands operators/tasks at the artifact level. This is what we recently worked on, and we've merged a PR implementing PythonOperatorPerArtifact Airflow DAG flavor (previously supporting PythonOperatorPerSession only).

For instance, this is an example of an Airflow DAG file under the PythonOperatorPerSession flavor:

import iris_module
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.utils.dates import days_ago

default_dag_args = {
    "owner": "airflow",
    "retries": 2,
    "start_date": days_ago(1),
}

with DAG(
    dag_id="iris_dag",
    schedule_interval="*/15 * * * *",
    max_active_runs=1,
    catchup=False,
    default_args=default_dag_args,
) as dag:

    run_session_including_iris_model = PythonOperator(
        task_id="run_session_including_iris_model_task",
        python_callable=iris_module.run_session_including_iris_model,
    )

As you can see, all artifacts are "wrapped" into one single operator, which reduces tracking and control over each artifact. Contrast this with the following equivalent version that gets generated under the newly available PythonOperatorPerArtifact flavor:

import pathlib
import pickle

import iris_module
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.utils.dates import days_ago


def dag_setup():
    pickle_folder = pathlib.Path("/tmp").joinpath("iris")
    if not pickle_folder.exists():
        pickle_folder.mkdir()


def dag_teardown():
    pickle_files = pathlib.Path("/tmp").joinpath("iris").glob("*.pickle")
    for f in pickle_files:
        f.unlink()


def task_url1_for_artifact_iris_model_and_downstream():

    url1 = iris_module.get_url1_for_artifact_iris_model_and_downstream()

    pickle.dump(url1, open("/tmp/iris/variable_url1.pickle", "wb"))


def task_iris_model():

    url1 = pickle.load(open("/tmp/iris/variable_url1.pickle", "rb"))

    mod = iris_module.get_iris_model(url1)

    pickle.dump(mod, open("/tmp/iris/variable_mod.pickle", "wb"))


def task_iris_petal_length_pred():

    mod = pickle.load(open("/tmp/iris/variable_mod.pickle", "rb"))

    url1 = pickle.load(open("/tmp/iris/variable_url1.pickle", "rb"))

    petal_length_pred = iris_module.get_iris_petal_length_pred(mod, url1)

    pickle.dump(
        petal_length_pred, open("/tmp/iris/variable_petal_length_pred.pickle", "wb")
    )


default_dag_args = {
    "owner": "airflow",
    "retries": 2,
    "start_date": days_ago(1),
}

with DAG(
    dag_id="iris_dag",
    schedule_interval="*/15 * * * *",
    max_active_runs=1,
    catchup=False,
    default_args=default_dag_args,
) as dag:

    setup = PythonOperator(
        task_id="dag_setup",
        python_callable=dag_setup,
    )

    teardown = PythonOperator(
        task_id="dag_teardown",
        python_callable=dag_teardown,
    )

    url1_for_artifact_iris_model_and_downstream = PythonOperator(
        task_id="url1_for_artifact_iris_model_and_downstream_task",
        python_callable=task_url1_for_artifact_iris_model_and_downstream,
    )

    iris_model = PythonOperator(
        task_id="iris_model_task",
        python_callable=task_iris_model,
    )

    iris_petal_length_pred = PythonOperator(
        task_id="iris_petal_length_pred_task",
        python_callable=task_iris_petal_length_pred,
    )

    url1_for_artifact_iris_model_and_downstream >> iris_model

    iris_model >> iris_petal_length_pred

    setup >> url1_for_artifact_iris_model_and_downstream

    iris_petal_length_pred >> teardown

It's longer code but it now opens room for engineers to peak and control at the artifact level, which allows for customization.

Note that this update is not yet visible to the user since the new graph refactor project has not yet been exposed to the public API. We are wrapping up remaining tasks for the final "replacement" though, so stay tuned!

Reference

• LIN-531 Implement PythonOperatorPerArtifact #748