Source: examples/airflow

Running SkyPilot tasks in Airflow with the SkyPilot API Server

In this guide, we show how a training workflow involving data preprocessing, training and evaluation can be first easily developed with SkyPilot, and then orchestrated in Airflow.

This example uses a remote SkyPilot API Server to manage shared state across invocations, and includes a failure callback to tear down the SkyPilot cluster on task failure.

💡 Tip: SkyPilot also supports defining and running pipelines without Airflow. Check out Jobs Pipelines for more information.

Why use SkyPilot with Airflow?

In AI workflows, the transition from development to production is hard.

Workflow development happens ad-hoc, with a lot of interaction required with the code and data. When moving this to an Airflow DAG in production, managing dependencies, environments and the infra requirements of the workflow gets complex. Porting the code to an airflow requires significant time to test and validate any changes, often requiring re-writing the code as Airflow operators.

SkyPilot seamlessly bridges the dev -> production gap.

SkyPilot can operate on any of your infra, allowing you to package and run the same code that you ran during development on a production Airflow cluster. Behind the scenes, SkyPilot handles environment setup, dependency management, and infra orchestration, allowing you to focus on your code.

Here’s how you can use SkyPilot to take your dev workflows to production in Airflow:

1. Define and test your workflow as SkyPilot tasks.

   • Use sky launch and Sky VSCode integration to run, debug and iterate on your code.

2. Orchestrate SkyPilot tasks in Airflow by invoking sky launch on their YAMLs as a task in the Airflow DAG.

   • Airflow does the scheduling, logging, and monitoring, while SkyPilot handles the infra setup and task execution.

Prerequisites

• Airflow installed locally (SequentialExecutor)

• SkyPilot API server endpoint to send requests to.

  • If you do not have one, refer to the API server docs to deploy one.

  • For this specific example: the API server should have AWS/GCS access to create buckets to store intermediate task outputs.

Configuring the API server endpoint

Once your API server is deployed, you will need to configure Airflow to use it. Set the SKYPILOT_API_SERVER_ENDPOINT variable in Airflow - it will be used by the run_sky_task function to send requests to the API server:

airflow variables set SKYPILOT_API_SERVER_ENDPOINT https://<api-server-endpoint>

You can also use the Airflow web UI to set the variable:

Defining the tasks

We will define the following tasks to mock a training workflow:

1. data_preprocessing.yaml: Generates data and writes it to a bucket.

2. train.yaml: Trains a model on the data in the bucket.

3. eval.yaml: Evaluates the model and writes evaluation results to the bucket.

We have defined these tasks in this directory and uploaded them to a Git repository.

When developing the workflow, you can run the tasks independently using sky launch:

# Run the data preprocessing task, replacing <bucket-name> with the bucket you created above
sky launch -c data --env DATA_BUCKET_NAME=<bucket-name> --env DATA_BUCKET_STORE_TYPE=s3 data_preprocessing.yaml

The train and eval step can be run in a similar way:

# Run the train task
sky launch -c train --env DATA_BUCKET_NAME=<bucket-name> --env DATA_BUCKET_STORE_TYPE=s3 train.yaml

Hint: You can use ssh and VSCode to interactively develop and debug the tasks.

Note: eval can be optionally run on the same cluster as train with sky exec.

Writing the Airflow DAG

Once we have developed the tasks, we can seamlessly run them in Airflow.

1. No changes required to our tasks - we use the same YAMLs we wrote in the previous step to create an Airflow DAG in sky_train_dag.py.

2. Airflow native logging - SkyPilot logs are written to container stdout, which is captured as task logs in Airflow and displayed in the UI.

3. Easy debugging - If a task fails, you can independently run the task using sky launch to debug the issue. SkyPilot will recreate the environment in which the task failed.

Here’s a snippet of the DAG declaration in sky_train_dag.py:

with DAG(dag_id='sky_train_dag',
         default_args=default_args,
         schedule_interval=None,
         catchup=False) as dag:
    # Path to SkyPilot YAMLs. Can be a git repo or local directory.
    base_path = 'https://github.com/skypilot-org/mock-train-workflow.git'

    # Generate bucket UUID as first task
    bucket_uuid = generate_bucket_uuid()
    
    # Use the bucket_uuid from previous task
    common_envs = {
        'DATA_BUCKET_NAME': f"sky-data-demo-{{{{ task_instance.xcom_pull(task_ids='generate_bucket_uuid') }}}}",
        'DATA_BUCKET_STORE_TYPE': 's3'
    }
    
    preprocess = run_sky_task.override(task_id="data_preprocess")(
        repo_url, 'data_preprocessing.yaml', envs_override=common_envs, git_branch='clientserver_example')
    train_task = run_sky_task.override(task_id="train")(
        repo_url, 'train.yaml', envs_override=common_envs, git_branch='clientserver_example')
    eval_task = run_sky_task.override(task_id="eval")(
        repo_url, 'eval.yaml', envs_override=common_envs, git_branch='clientserver_example')

    # Define the workflow
    bucket_uuid >> preprocess >> train_task >> eval_task

Behind the scenes, the run_sky_task uses the Airflow native Python operator to invoke the SkyPilot API. All SkyPilot API calls are made to the remote API server, which is configured using the SKYPILOT_API_SERVER_ENDPOINT variable.

The task YAML files can be sourced in two ways:

1. From a Git repository (as shown above):

   repo_url = 'https://github.com/skypilot-org/mock-train-workflow.git'
   run_sky_task(...)(repo_url, 'path/to/yaml', git_branch='optional_branch')
   

   The task will automatically clone the repository and checkout the specified branch before execution.

2. From a local path:

   local_path = '/path/to/local/directory'
   run_sky_task(...)(local_path, 'path/to/yaml')
   

   This is useful during development or when your tasks are stored locally.

All clusters are set to auto-down after the task is done, so no dangling clusters are left behind.

Running the DAG

1. Copy the DAG file to the Airflow DAGs directory.

   cp sky_train_dag.py /path/to/airflow/dags                                               
   # If your Airflow is running on Kubernetes, you may use kubectl cp to copy the file to the pod
   # kubectl cp sky_train_dag.py <airflow-pod-name>:/opt/airflow/dags 
   

2. Run airflow dags list to confirm that the DAG is loaded.

3. Find the DAG in the Airflow UI (typically http://localhost:8080) and enable it. The UI may take a couple of minutes to reflect the changes. Force unpause the DAG if it is paused with airflow dags unpause sky_train_dag

4. Trigger the DAG from the Airflow UI using the Trigger DAG button.

5. Navigate to the run in the Airflow UI to see the DAG progress and logs of each task.

If a task fails, task_failure_callback will automatically tear down the SkyPilot cluster.

Future work: a native Airflow Executor built on SkyPilot

Currently this example relies on a helper run_sky_task method to wrap SkyPilot invocation in @task, but in the future SkyPilot can provide a native Airflow Executor.

In such a setup, SkyPilot state management also not be required, as the executor will handle SkyPilot cluster launching and termination.

Included files

data_preprocessing.yaml

resources:
  cpus: 1

envs:
  DATA_BUCKET_NAME: sky-demo-data-test
  DATA_BUCKET_STORE_TYPE: s3

file_mounts:
  /data:
    name: $DATA_BUCKET_NAME
    store: $DATA_BUCKET_STORE_TYPE

setup: |
  echo "Setting up dependencies for data preprocessing..."

run: |
  echo "Running data preprocessing..."
  
  # Generate few files with random data to simulate data preprocessing
  for i in {0..9}; do
      dd if=/dev/urandom of=/data/file_$i bs=1M count=10
  done
  
  echo "Data preprocessing completed, wrote to $DATA_BUCKET_NAME"

eval.yaml

resources:
  cpus: 1
  # Add GPUs here

envs:
  DATA_BUCKET_NAME: sky-demo-data-test
  DATA_BUCKET_STORE_TYPE: s3

file_mounts:
  /data:
    name: $DATA_BUCKET_NAME
    store: $DATA_BUCKET_STORE_TYPE

setup: |
  echo "Setting up dependencies for eval..."

run: |
  echo "Evaluating the trained model..."
  
  # Run a mock evaluation job that reads the trained model from /data/trained_model.txt
  cat /data/trained_model.txt || true
  # Generate a mock accuracy
  ACCURACY=$(shuf -i 90-100 -n 1)
  echo "Metric - accuracy: $ACCURACY%"
  echo "Evaluation report" > /data/evaluation_report.txt
  
  echo "Evaluation completed, report written to $DATA_BUCKET_NAME"

sky_train_dag.py

import os
import uuid

from airflow import DAG
from airflow.decorators import task
from airflow.models import Variable
from airflow.utils.dates import days_ago
import yaml

default_args = {
    'owner': 'airflow',
    'start_date': days_ago(1),
}

# Unique bucket name for this DAG run
DATA_BUCKET_NAME = str(uuid.uuid4())[:4]


def task_failure_callback(context):
    """Callback to shut down SkyPilot cluster on task failure."""
    cluster_name = context['task_instance'].xcom_pull(
        task_ids=context['task_instance'].task_id, key='cluster_name')
    if cluster_name:
        print(
            f"Task failed or was cancelled. Shutting down SkyPilot cluster: {cluster_name}"
        )
        import sky
        down_request = sky.down(cluster_name)
        sky.stream_and_get(down_request)


@task(on_failure_callback=task_failure_callback)
def run_sky_task(base_path: str,
                 yaml_path: str,
                 envs_override: dict = None,
                 git_branch: str = None,
                 **kwargs):
    """Generic function to run a SkyPilot task.

    This is a blocking call that runs the SkyPilot task and streams the logs.
    In the future, we can use deferrable tasks to avoid blocking the worker
    while waiting for cluster to start.

    Args:
        base_path: Base path (local directory or git repo URL)
        yaml_path: Path to the YAML file (relative to base_path)
        envs_override: Dictionary of environment variables to override in the task config
        git_branch: Optional branch name to checkout (only used if base_path is a git repo)
    """
    import subprocess
    import tempfile

    # Set the SkyPilot API server endpoint from Airflow Variables
    endpoint = Variable.get('SKYPILOT_API_SERVER_ENDPOINT', None)
    if not endpoint:
        raise ValueError('SKYPILOT_API_SERVER_ENDPOINT is not set in airflow.')
    os.environ['SKYPILOT_API_SERVER_ENDPOINT'] = endpoint

    original_cwd = os.getcwd()
    try:
        # Handle git repos vs local paths
        if base_path.startswith(('http://', 'https://', 'git://')):
            with tempfile.TemporaryDirectory() as temp_dir:
                # TODO(romilb): This assumes git credentials are available in the airflow worker
                subprocess.run(['git', 'clone', base_path, temp_dir],
                               check=True)

                # Checkout specific branch if provided
                if git_branch:
                    subprocess.run(['git', 'checkout', git_branch],
                                   cwd=temp_dir,
                                   check=True)

                full_yaml_path = os.path.join(temp_dir, yaml_path)
                # Change to the temp dir to set context
                os.chdir(temp_dir)

                # Run the sky task
                return _run_sky_task(full_yaml_path, envs_override or {},
                                     kwargs)
        else:
            full_yaml_path = os.path.join(base_path, yaml_path)
            os.chdir(base_path)

            # Run the sky task
            return _run_sky_task(full_yaml_path, envs_override or {}, kwargs)
    finally:
        os.chdir(original_cwd)


def _run_sky_task(yaml_path: str, envs_override: dict, kwargs: dict):
    """Internal helper to run the sky task after directory setup."""
    import sky

    with open(os.path.expanduser(yaml_path), 'r', encoding='utf-8') as f:
        task_config = yaml.safe_load(f)

    # Initialize envs if not present
    if 'envs' not in task_config:
        task_config['envs'] = {}

    # Update the envs with the override values
    # task.update_envs() is not used here, see https://github.com/skypilot-org/skypilot/issues/4363
    task_config['envs'].update(envs_override)

    task = sky.Task.from_yaml_config(task_config)
    cluster_uuid = str(uuid.uuid4())[:4]
    task_name = os.path.splitext(os.path.basename(yaml_path))[0]
    cluster_name = f'{task_name}-{cluster_uuid}'
    kwargs['ti'].xcom_push(key='cluster_name',
                           value=cluster_name)  # For failure callback

    launch_request_id = sky.launch(task, cluster_name=cluster_name, down=True)
    job_id, _ = sky.stream_and_get(launch_request_id)
    # TODO(romilb): In the future, we can use deferrable tasks to avoid blocking
    # the worker while waiting for cluster to start.

    # Stream the logs for airflow logging
    sky.tail_logs(cluster_name=cluster_name, job_id=job_id, follow=True)

    # Terminate the cluster after the task is done
    down_id = sky.down(cluster_name)
    sky.stream_and_get(down_id)


@task
def generate_bucket_uuid(**context):
    bucket_uuid = str(uuid.uuid4())[:4]
    return bucket_uuid


with DAG(dag_id='sky_train_dag',
         default_args=default_args,
         schedule_interval=None,
         catchup=False) as dag:
    # Path to SkyPilot YAMLs. Can be a git repo or local directory.
    base_path = 'https://github.com/skypilot-org/mock-train-workflow.git'

    # Generate bucket UUID as first task
    # See https://stackoverflow.com/questions/55748050/generating-uuid-and-use-it-across-airflow-dag
    bucket_uuid = generate_bucket_uuid()

    # Use the bucket_uuid from previous task
    common_envs = {
        'DATA_BUCKET_NAME': f"sky-data-demo-{{{{ task_instance.xcom_pull(task_ids='generate_bucket_uuid') }}}}",
        'DATA_BUCKET_STORE_TYPE': 's3'
    }

    preprocess = run_sky_task.override(task_id="data_preprocess")(
        base_path,
        'data_preprocessing.yaml',
        envs_override=common_envs,
        git_branch='clientserver_example')
    train_task = run_sky_task.override(task_id="train")(
        base_path,
        'train.yaml',
        envs_override=common_envs,
        git_branch='clientserver_example')
    eval_task = run_sky_task.override(task_id="eval")(
        base_path,
        'eval.yaml',
        envs_override=common_envs,
        git_branch='clientserver_example')

    # Define the workflow
    bucket_uuid >> preprocess >> train_task >> eval_task

train.yaml

resources:
  cpus: 1
  # Add GPUs here

envs:
  DATA_BUCKET_NAME: sky-demo-data-test
  DATA_BUCKET_STORE_TYPE: s3
  NUM_EPOCHS: 2

file_mounts:
  /data:
    name: $DATA_BUCKET_NAME
    store: $DATA_BUCKET_STORE_TYPE

setup: |
  echo "Setting up dependencies for training..."

run: |
  echo "Running training..."
  
  # Run a mock training job that loops through the files in /data starting with 'file_'
  for (( i=1; i<=NUM_EPOCHS; i++ )); do
    for file in /data/file_*; do
      echo "Epoch $i: Training on $file"
      sleep 2
    done
  done
  
  # Mock checkpointing the trained model to the data bucket
  echo "Trained model" > /data/trained_model.txt
  
  echo "Training completed, model written to to $DATA_BUCKET_NAME"