Reduce task queueing latency when using Cosmos

[tatiana]

Context

This issue happened before the last release (1.4) - and can also be reproduced with Cosmos 1.4.1.

Users have observed long task queueing times for Cosmos tasks:

This is not observed when using, for instance, BashOperator task instances:

The task queueing time for the Cosmos DAG is consistently 5s, while it is close to 0s for the BashOperator one.

How to reproduce

Example Cosmos DAG:

from datetime import datetime

from cosmos import DbtDag, ProjectConfig, ProfileConfig, RenderConfig
from cosmos.constants import TestBehavior

from include.constants import jaffle_shop_path, venv_execution_config

dbt_profile_example = DbtDag(
    # dbt/cosmos-specific parameters
    project_config=ProjectConfig(jaffle_shop_path),
    profile_config=ProfileConfig(
        # these map to dbt/jaffle_shop/profiles.yml
        profile_name="airflow_db",
        target_name="bq",
        profiles_yml_filepath=jaffle_shop_path / "profiles.yml",
    ),
    render_config=RenderConfig(test_behavior=TestBehavior.NONE),
    execution_config=venv_execution_config,
    # normal dag parameters
    schedule_interval="@daily",
    start_date=datetime(2023, 1, 1),
    catchup=False,
    dag_id="bq_profile_example",
    tags=["profiles"],
)

Example BashOperator DAG:

from pendulum import datetime

from airflow import DAG
from airflow.operators.bash_operator import BashOperator

from include.constants import jaffle_shop_path as DBT_PROJECT_DIR
from include.constants import dbt_executable

DBT_BIN_PATH = dbt_executable

BASE_DBT_CMD = f"{DBT_BIN_PATH} %s --profiles-dir {DBT_PROJECT_DIR} --project-dir {DBT_PROJECT_DIR} --profile airflow_db --target bq"

DBT_DEPS =  BASE_DBT_CMD % " deps"
DBT_LS = BASE_DBT_CMD % " ls"
BASE_DBT_CMD = DBT_DEPS + " && " + DBT_LS + " && " + BASE_DBT_CMD


with DAG(
    "bash_bq_example",
    start_date=datetime(2020, 12, 23),
    description="A sample Airflow DAG to invoke dbt runs using a BashOperator",
    schedule_interval=None,
    catchup=False,
) as dag:
    raw_payments_seed = BashOperator(
        task_id="raw_payments_seed",
        bash_command=BASE_DBT_CMD % "seed --models raw_payments",
    )
    raw_orders_seed = BashOperator(
        task_id="raw_orders_seed",
        bash_command=BASE_DBT_CMD % "seed --models raw_orders",
    )
    raw_customers_seed = BashOperator(
        task_id="raw_customers_seed",
        bash_command=BASE_DBT_CMD % "seed --models raw_customers",
    )

    run_stg_payments = BashOperator(
        task_id="run_stg_payments",
        bash_command=BASE_DBT_CMD % "run --models stg_payments"
    )

    run_stg_orders = BashOperator(
        task_id="run_stg_orders",
        bash_command=BASE_DBT_CMD % "run --models stg_orders"
    )

    run_stg_customers = BashOperator(
        task_id="run_stg_customers",
        bash_command=BASE_DBT_CMD % "run --models stg_customers"
    )

    run_orders = BashOperator(
        task_id="run_orders",
        bash_command=BASE_DBT_CMD % "run --models orders"
    )

    run_customers = BashOperator(
        task_id="run_customers",
        bash_command=BASE_DBT_CMD % "run --models customers"
    )

    raw_payments_seed >> run_stg_payments
    raw_orders_seed >> run_stg_orders
    raw_customers_seed >> run_stg_customers

    [run_stg_payments, run_stg_orders, run_stg_customers] >> run_customers
    [run_stg_payments, run_stg_orders] >> run_orders

It is expected that some delay will be introduced by Cosmos, but it is too long ATM even for relatively small dbt projects.

Possible solution

• Cache TaskGroup/DAG regardless of the load_method #926

Acceptance criteria

• Reduce the task queueing time when using Cosmos

[tatiana]

Progress can be seen in draft PR: #992

[tatiana]

Further progress:

• Cache TaskGroup/DAG regardless of the load_method #926 (comment)
• Cache TaskGroup/DAG regardless of the load_method #926 (comment)

[tatiana]

When using the current approach in a distributed environment, there are two challenges:

1. Concurrent tasks (in the same node) try to create the cache at the same time
2. Tasks running in different nodes (that didn't have the cache) have to generate it

We'll look into improving this.

Examples of the behaviour in a distributed Airflow environment:

[tatiana]

To store the dbt ls output as an Airflow Variable seems to be more promising on a distributed Airflow environment, with the caveats described in the ticket description:
#1014

[tatiana]

Waiting for feedback from end-users on #1014.

Analysed and confirmed the feasibility of this approach for larger dbt projects:

• Another real dbt project with 9,285 models led to a dbt ls output of 8.4 MB, uncompressed.
• Maximum cell size in Postgres: 20MB

[tatiana]

Here are some updates on this task: We received feedback from end-users on #1014, and they were happy with the performance improvements and the overall approach. We monitored the load in their deployment's database, and it was fine, without any significant increases.

The next agreed-upon steps were to rename the variable used to cache to be prefixed with a cosmic identifier and to work on the purging strategy.

We initially implemented CachePurgeConfig in d1a30b6, but we realized that would be more overhead for end-users, which was unnecessary. Initially, there was also the intention to use a hash of DAG file itself to be part of the caching version definition. The users explained they wanted to be able to change other parts of their DAG that would be unrelated to Cosmos dbt ls.

We released astronomer-cosmos==1.5.0a6

https://pypi.org/project/astronomer-cosmos/1.5.0a6/
https://github.com/astronomer/astronomer-cosmos/pull/1014/commits

We addressed the feedback you gave during our Monday session on purging:

The Airflow vars used to cache the dbt ls output are now prefixed with cosmos_cache
If the following arguments are changed, we'll automatically purge the cache:

• RenderConfig.env_vars
• RenderConfig.exclude
• RenderConfig.select
• RenderConfig.selector
• ProjectConfig.dbt_vars
• ProjectConfig.env_vars
• ProjectConfig.partial_parse

The following argument was introduced in case you'd like to define airflow variables that could be used to purge the cache (it expects a list with Airflow variable names)

• RenderConfig.airflow_vars_to_purge_cache

What is missing on purging:

• Take into account "packages.yml" or "dependencies.yml"
• Change how we consider if the dbt project files changed (before, we were taking into account the modified timestamp, but this does not seem to play well with --dags deployment, since the file timestamps change often)

[tatiana]

An example that can be tested using 1.5.0a6:

DbtDag(
    # dbt/cosmos-specific parameters
    project_config=ProjectConfig(
        jaffle_shop_path,
        dbt_vars={"some-var": 2}
    ),
    profile_config=ProfileConfig(
        # these map to dbt/jaffle_shop/profiles.yml
        profile_name="airflow_db",
        target_name="bq",
        profiles_yml_filepath=jaffle_shop_path / "profiles.yml",
    ),
    render_config=RenderConfig(
        test_behavior=TestBehavior.NONE,
        env_vars={"PURGE": os.getenv("PURGE", "0")},
        airflow_vars_to_purge_cache=["purge"],
    ),
    execution_config=venv_execution_config,
    # normal dag parameters
    schedule_interval="@daily",
    start_date=datetime(2023, 1, 1),
    catchup=False,
    dag_id="bq_efficient",
    tags=["profiles"],
)

If the value of dbt_vars changes, the cache for this DAG will be purged.
If the value of the environment variable PURGE is changed, the cache for this DAG will be purged.
If the value for the Airflow var purge is changed, the cache for this DAG will also be purged.

[tatiana]

Today, I've made a change to the dbt project hash. It's now created using sha256 and not the dbt project files modified timestamp (0e0f663). This change allows the solution to work with --dags only deployment.

I'll also consider the sha256 for "packages.yml" or "dependencies.yml" if they exist, and cut a new alpha released once these changes are implemented.

[tatiana]

With the latest changes, the overall DAG duration for a jaffle shop DAG dropped from 00:01:03 (already using partial parsing cache) to 00:00:39, with a consistent small task queue time even after --dags deployments.

Before:

After:

[tatiana]

Some of the last improvements on this workstream:

• I've added several tests to cover the behaviour
• Had very positive feedback from customers
• Placed the ticket Speed up LoadMode.DBT_LS by caching dbt ls output in Airflow Variable #1014 under review, fixing one last broken test

[tatiana]

Fixed issues (deterministic cache hash on different VMs) and created tests that cover 100% of this feature.