Cache TaskGroup/DAG regardless of the load_method

[tatiana]

Context

PR #904 introduced the first caching implementation in Cosmos on top of the previously implemented support for partial parsing (#800).

Although these changes contributed significantly to the performance of Cosmos, this approach does not improve the DAG parsing performance for methods other than dbt ls. In those cases, Cosmos will still:

• parse the project, manifest, or dbt ls command
• select nodes based on selectors given by the user

During DAG parsing - which happens at the Airflow DAG Processor and is executed every time a task is scheduled/run.

This task aims to extend the cache support to allow the cache of the complete DAG/TaskGroup, without the need to do the two previously described steps.

Acceptance criteria

• Create a function to evaluate if there were no changes to the dbt project or to the Airflow DAG
• Allow users to define their own caching function
• Support caching the DAG/TaskGroup into a pickle file
• Use the caching function to decide if Cosmos should load the DAG/TaskGroup from the cache or if the DAG/TaskGroup should be recreated (saving the new cache on disk)
• Run performance test when using load_method=LoadMode.DBT_MANIFEST

[tatiana]

Progress can be seen in draft PR: #992

[tatiana]

I explored and validated a few possible ways of quickly generating the "version" of the cache. Some are safer than others, the idea was to get a sense of cost/benefit:

"""
How to run script:

pip install checksumdir  # https://pypi.org/project/checksumdir/
python hash_comparison.py --path your/dbt/project

python cosmos/hash_comparison.py --path  dags/dbt/perf
Comparing hash methods pm a directory with 5009 files
- create_hash_with_checksumdir_md5        : 0.2171382334
- create_hash_with_checksumdir_sha1       : 0.1730406
- create_hash_with_checksumdir_sha256     : 0.17227659160000003
- create_hash_with_dir_size               : 0.05859905000000003
- create_hash_with_file_changes           : 0.03936836659999994
- create_hash_with_most_recently_changed  : 0.10243044999999995
- create_hash_with_sql_changes            : 0.03878682500000004
- create_hash_with_subprocess_du          : 0.014803133400000057
- create_hash_with_walk_files_md5         : 0.14606525
- create_hash_with_walk_files_size        : 0.024191108199999967
"""
import argparse
import hashlib
import inspect
import os
import subprocess
import sys
import timeit
from pathlib import Path

import checksumdir


ROOT_DIR = Path(".")
ROOT_DIR = Path("./dags/dbt/jaffle_shop")
REPETITIONS = 5


def create_hash_with_dir_size(dir_path: Path) -> int:
    return sum(f.stat().st_size for f in dir_path.glob('**/*') if f.is_file())


def create_hash_with_most_recently_changed(dir_path: Path) -> str:
    value = subprocess.run(
        f'find {dir_path} -type f -print0 | xargs -0 stat -f "%m %N" | sort -rn | head -1',
        text=True,
        stdout=subprocess.PIPE,
        check=True,
        shell=True
    )
    return value

def create_hash_with_subprocess_du(dir_path: Path) -> str:
    value = subprocess.run(
        f"du -s {dir_path}",
        text=True,
        stdout=subprocess.PIPE,
        check=True,
        shell=True
    )
    return value.stdout.strip()


def create_hash_with_walk_files_size(dir_path: Path):
    return sum(
        sum(
            os.path.getsize(os.path.join(walk_result[0], element))
            for element in walk_result[2]
        )
        for walk_result in os.walk(dir_path)
   )


def create_hash_with_walk_files_md5(dir_path: Path) -> str:
    hasher = hashlib.md5()
    for walk_result in os.walk(dir_path):
        for element in walk_result[2]:
            filepath = os.path.join(walk_result[0], element)
            with open(str(filepath), 'rb') as fp:
                buf = fp.read()
                hasher.update(buf)
    return hasher.hexdigest()


def create_hash_with_file_changes(dir_path: Path) -> float:
    return sum([path.stat().st_mtime for path in dir_path.glob('**/*')])


def create_hash_with_sql_changes(dir_path: Path) -> float:
    return sum([path.stat().st_mtime for path in dir_path.glob('**/*.sql')])


create_hash_with_checksumdir_md5 = lambda dir_path: checksumdir.dirhash(str(dir_path), 'md5')
create_hash_with_checksumdir_sha1 = lambda dir_path: checksumdir.dirhash(str(dir_path), 'sha1')
create_hash_with_checksumdir_sha256 = lambda dir_path: checksumdir.dirhash(str(dir_path), 'sha256')


if __name__ == "__main__":
    parser = argparse.ArgumentParser("hash_comparison")
    parser.add_argument("--path", help=".", type=Path, default=ROOT_DIR)
    args = parser.parse_args()

    total_files = len([path for path in args.path.glob('**/*') if path.is_file()])
    print(f"Comparing hash methods pm a directory with {total_files} files")


    current_module = sys.modules[__name__]
    functions_to_test = inspect.getmembers(current_module, inspect.isfunction)
    for function_name, function in functions_to_test:
        average_time = timeit.timeit(
            lambda: function(args.path),
            number=REPETITIONS
        ) / REPETITIONS
        print(f"- {function_name: <40}: {average_time}")

So far, it seems create_hash_with_file_changes is one of the most accurate and light weight approaches. I strongly believe we should allow users to customise the method used to calculate the version, but this may be a good default for Cosmos.

[tatiana]

I extended the initial implementation to:

• Version the DbtDag cache in 0db7121
• Extend the caching mechanism to DbtTaskGroup in daad0b3

I also updated the PR description with what is missing:

• Allow users to customise their own calculate_current_version
• Refactor so there isn't so much code duplicate between DbtDag and TaskGroup
• Allow users to purge the cache (via env var or param passed in the Airflow UI - second option is probably preferable)
• Unify _create_cache_identifier and create_cache_identifier_v2
• Write tests

[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: