Managing Airflow DAGs for Data Science tasks in GCP Composer

Question

Pulumi · Accepted Answer

Google Cloud Platform (GCP) Composer is a fully managed workflow orchestration service built on Apache Airflow. It allows you to create, schedule, monitor, and manage workflows. Airflow DAGs (Directed Acyclic Graphs) represent a collection of tasks that you want to run, organized in a way that reflects their relationships and dependencies.

To manage Airflow DAGs with GCP Composer using Pulumi, you'll create an instance of the Composer Environment where you can run your DAGs. This environment comprises several elements, including the configuration of the GKE (Google Kubernetes Engine) cluster it will run on, the location where DAGs should be stored, and various settings for the Airflow environment.

The following program in Python uses Pulumi to create a simple GCP Composer Environment to host Apache Airflow for running your DAGs. Before diving into the code, make sure you have the Pulumi CLI installed and configured with GCP credentials. We'll go step-by-step to help you understand the process.

First, we'll need to import the `pulumi_gcp` module which provides the necessary classes for creating GCP resources:

- The `gcp.composer.Environment` class is what we'll use to create the Composer environment.
- We'll configure the environment settings, like machine type for the nodes, cloud storage bucket for DAGs, and the location of the GCP environment.

Once the environment is up and running, you can deploy your DAGs to the specified GCS bucket and manage them in the Airflow web UI that Composer provides.

Now, let's jump into the code:

```python
import pulumi
import pulumi_gcp as gcp

# Define the region where our GCP resources will be located.
region = 'us-central1'

# Configure the environment for the Composer and Airflow settings. 
# Replace `your_project_id` with your actual GCP project ID.
composer_environment_name = "airflow-environment"
composer_environment_config = gcp.composer.EnvironmentConfigArgs(
    node_count=3,
    node_config=gcp.composer.EnvironmentConfigNodeConfigArgs(
        machine_type='n1-standard-1',
        network='default',  # Use the default network, but a custom network can be specified.
        subnetwork='default',  # Use the default subnetwork, but a custom subnetwork can be specified.
        disk_size_gb=20,  # Disk size for each VM node in the environment.
    ),
    software_config=gcp.composer.EnvironmentConfigSoftwareConfigArgs(
        image_version='composer-1.17.1-airflow-2.2.2',  # Specify the Airflow version.
        # Set Airflow environment variables if needed. For example, to increase parallelism:
        # env_variables={
        #    'AIRFLOW__CORE__PARALLELISM': '5'
        # },
    ),
    # Indicates where the DAGs are stored in GCS and can be found by Airflow.
    dag_gcs_prefix=f'gs://{composer_environment_name}-dags',  
)

# Create the Composer environment
composer_environment = gcp.composer.Environment(
    composer_environment_name,
    config=composer_environment_config,
    location=region,
    project='your_project_id'  # Make sure to replace with your GCP project ID.
)

# Export the Airflow web UI URL so we can easily access it.
pulumi.export('airflow_web_ui', composer_environment.config.apply(lambda conf: conf.airflow_uri))
```

In this program, we create a Composer environment with basic settings. The environment comprises three nodes running on `n1-standard-1` machine types within the default network. We specify the Airflow version we want to use by setting the `image_version` in the `software_config`.

The `dag_gcs_prefix` tells Composer where to find the DAGs in Google Cloud Storage. Airflow will look into the specified GCS bucket for any DAGs to run.

Finally, we export the URI for the Airflow web UI, which you can use to access the Airflow interface and manage your DAGs visually.

Remember to replace `your_project_id` with the GCP project ID you're working with.

After running this script with Pulumi, you'll have an operational Airflow environment within GCP Composer. You can then proceed to write your DAGs in Python and deploy them to the GCS bucket configured as `dag_gcs_prefix`.

To deploy DAGs, you typically upload your `.py` files containing DAG definitions to the GCS bucket. The Airflow scheduler running in Composer will pick up the new DAGs and begin orchestrating your workflows based on the schedules and dependencies you've defined.