Skip to content

Integrate Run:AI with Apache Airflow

Airflow is a platform to programmatically author, schedule, and monitor workflows. Specifically, it is used in Machine Learning to create pipelines.

Airflow DAG

In Airflow, a DAG – or a Directed Acyclic Graph – is a collection of all the tasks you want to run, organized in a way that reflects their relationships and dependencies.

A DAG is defined in a Python script, which represents the DAGs structure (tasks and their dependencies) as code.

For example, a simple DAG could consist of three tasks: A, B, and C. It could say that A has to run successfully before B can run, but C can run anytime. It could say that task A times out after 5 minutes, and B can be restarted up to 5 times in case it fails. It might also say that the workflow will run every night at 10pm, but shouldn’t start until a certain date.

Airflow tasks are sent for execution. Specifically, the Airflow - Kubernetes integration allows Airflow tasks to be scheduled on a Kubernetes cluster.

Run:AI - Airflow Integration

DAGs are defined in Python. Airflow tasks based on Kubernetes are defined via the KubernetesPodOperator class. To run an Airflow task with Run:AI you must provide additional Run:AI-related properties to

dag = DAG(...)

resources = {
  "limit_gpu": <number-of-GPUs>

job = KubernetesPodOperator(
    labels={"project": '<project-name>'},
The code:

  • Specifies the runai-scheduler which directs the task to be scheduled with the Run:AI scheduler
  • Specifies a Run:AI Project . A Project in Run:AI specifies guaranteed GPU & CPU quota.

Once you run the DAG, you can see Airflow tasks showing in the Run:AI UI.

Last update: April 14, 2021