Improve performance by 22-35% or more by caching partial parse artefact

[tatiana]

Improve the performance to run the benchmark DAG with 100 tasks by 34% and the benchmark DAG with 10 tasks by 22%, by persisting the dbt partial parse artifact in Airflow nodes. This performance can be even higher in the case of dbt projects that take more time to be parsed.

With the introduction of #800, Cosmos supports using dbt partial parsing files. This feature has led to a substantial performance improvement, particularly for large dbt projects, both during Airflow DAG parsing (using LoadMode.DBT_LS) and also Airflow task execution (when using ExecutionMode.LOCAL and ExecutionMode.VIRTUALENV).

There were two limitations with the initial support to partial parsing, which the current PR aims to address:

1. DAGs using Cosmos ProfileMapping classes could not leverage this feature. This is because the partial parsing relies on profile files not changing, and by default, Cosmos would mock the dbt profile in several parts of the code. The consequence is that users trying Cosmos 1.4.0a1 will see the following message:

13:33:16  Unable to do partial parsing because profile has changed
13:33:16  Unable to do partial parsing because env vars used in profiles.yml have changed

2. The user had to explicitly provide a partial_parse.msgpack file in the original project folder for their Airflow deployment - and if, for any reason, this became outdated, the user would not leverage the partial parsing feature. Since Cosmos runs dbt tasks from within a temporary directory, the partial parse would be stale for some users, it would be updated in the temporary directory, but the next time the task was run, Cosmos/dbt would not leverage the recently updated partial_parse.msgpack file.

The current PR addresses these two issues respectfully by:

1. Allowing users that want to leverage Cosmos ProfileMapping and partial parsing to use RenderConfig(enable_mock_profile=False)

2. Introducing a Cosmos cache directory where we are persisting partial parsing files. This feature is enabled by default, but users can opt out by setting the Airflow configuration [cosmos][enable_cache] = False (exporting the environment variable AIRFLOW__COSMOS__ENABLE_CACHE=0). Users can also define the temporary directory used to store these files using the [cosmos][cache_dir] Airflow configuration. By default, Cosmos will create and use a folder cosmos inside the system's temporary directory: https://docs.python.org/3/library/tempfile.html#tempfile.gettempdir .

This PR affects both DAG parsing and task execution. Although it does not introduce an optimisation per se, it makes the partial parse feature implemented #800 available to more users.

Closes: #722

I updated the documentation in the PR: #898

Some future steps related to optimization associated to caching to be addressed in separate PRs:
i. Change how we create mocked profiles, to create the file itself in the same way, referencing an environment variable with the same name - and only changing the value of the environment variable (#924)
ii. Extend caching to the profiles.yml created by Cosmos in the newly introduced tmp/cosmos without the need to recreate it every time (#925).
iii. Extend caching to the Airflow DAG/Task group as a pickle file - this approach is more generic and would work for every type of DAG parsing and executor. (#926)
iv. Support persisting/fetching the cache from remote storage so we don't have to replicate it for every Airflow scheduler and worker node. (#927)
v. Cache dbt deps lock file/avoid installing dbt steps every time. We can leverage package-lock.yml introduced in dbt t 1.7 (https://docs.getdbt.com/reference/commands/deps#predictable-package-installs), but ideally, we'd have a strategy to support older versions of dbt as well. (#930)
vi. Support caching partial_parse.msgpack even when vars change: https://medium.com/@sebastian.daum89/how-to-speed-up-single-dbt-invocations-when-using-changing-dbt-variables-b9d91ce3fb0d
vii. Support partial parsing in Docker and Kubernetes Cosmos executors (#929)
viii. Centralise all the Airflow-based config into Cosmos settings.py & create a dedicated docs page containing information about these (#928)

How to validate this change

Run the performance benchmark against this and the main branch, checking the value of /tmp/performance_results.txt.

Example of commands run locally:

# Setup
AIRFLOW_HOME=`pwd` AIRFLOW_CONN_AIRFLOW_DB="postgres://postgres:[email protected]:5432/postgres" PYTHONPATH=`pwd` AIRFLOW_HOME=`pwd` AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT=20000 AIRFLOW__CORE__DAG_FILE_PROCESSOR_TIMEOUT=20000 hatch run tests.py3.11-2.7:test-performance-setup

# Run test for 100 dbt models per DAG:
MODEL_COUNT=100 AIRFLOW_HOME=`pwd` AIRFLOW_CONN_AIRFLOW_DB="postgres://postgres:[email protected]:5432/postgres" PYTHONPATH=`pwd` AIRFLOW_HOME=`pwd` AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT=20000 AIRFLOW__CORE__DAG_FILE_PROCESSOR_TIMEOUT=20000 hatch run tests.py3.11-2.7:test-performance

An example of output when running 100 with the main branch:

NUM_MODELS=100
TIME=114.18614888191223
MODELS_PER_SECOND=0.8757629623135543
DBT_VERSION=1.7.13

And with the current PR:

NUM_MODELS=100
TIME=75.17766404151917
MODELS_PER_SECOND=1.33018232576064
DBT_VERSION=1.7.13

[netlify]

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Name	Link
🔨 Latest commit	2a3f007
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/663251880d28e10008c3746b

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.69%. Comparing base (808c9b9) to head (47aa75d).
Report is 1 commits behind head on main.

❗ Current head 47aa75d differs from pull request most recent head 2a3f007. Consider uploading reports for the commit 2a3f007 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #904      +/-   ##
==========================================
+ Coverage   94.79%   95.69%   +0.90%     
==========================================
  Files          57       59       +2     
  Lines        2767     2858      +91     
==========================================
+ Hits         2623     2735     +112     
+ Misses        144      123      -21     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[jlaneve]

do we want to cache this locally or to s3/cloud store? it'd be really cool to do to blob store as well, especially for KE, etc

[kaxil]

do we want to cache this locally or to s3/cloud store? it'd be really cool to do to blob store as well, especially for KE, etc

Worth using ObjectStore for the newer Airflow version too, that way we can support more Cloud Stores.

And not great, but a "hacky" approach can be to see if users have a "XCom backend" configured, and if yes, you can store the cache in XCom. From 2.9, you can also use ObjectStore as XCom backend: https://airflow.apache.org/blog/airflow-2.9.0/#object-storage-as-xcom-backend

It can happen in a separate PR too, so that you can add a caching layer in this one, and extend it to remote store (or XCom backend) in the next one.

And once the PR is out of draft and ready for prime time :) & reviews, worth adding what performance issues this solves & what it won't to set the right expectations with users.

[tatiana]

Thanks a lot for the feedback, @pankajkoti ! I addressed all of them except for the minor #904 (comment), which can be improved in a follow-up ticket once you reply. Since you approved the PR, I'll mark this comment as resolved and merging .

[maslick]

Enabling caching can quickly fill up the 30GB ephemeral storage on MWAA workers if there are 500 or more DAGs.