PoC on ExecutionMode.LOCAL_AIRFLOW_ASYNC

[tatiana]

Description co-authored by @pankajkoti @pankajastro

Context

As of Cosmos 1.5, we are using dbt core to run transformations in the database. In other words, the Airflow worker node is being blocked during the execution of the SQL transformation. If the transformation takes 6h, we have a Airflow worker node just waiting.

One way to overcome this is to use Airflow deferrable operators, as opposed to using dbt-core commands to execute the transformation. This approach is not new, and it was discussed in Airflow Summit 2023 by Monzo. We (@tatiana ) had a meeting with Monzo to discuss the details afterward, in London.

The overall idea is:

• Pre-compile the dbt project, so we have the pre-compiled SQL
• Execute the pre-compiled SQL using native Airflow operators, enabling the deferable mode

For a first end-to-end, we'll only support:

• SQL models
• ExecutionMode.LOCAL
• Databricks

This will allow us to measure the impact of using deferrable operators, and then we can extend to:

• Python models
• Other execution modes
• Other SQL backends

Acceptance criteria

• PoC to try to have this feature end-to-end that
  • Introduce ExecutionMode.LOCAL_AIRFLOW_ASYNC (think if we can find a better name)
  • Check how we are populating the compiled_sql template field in Cosmos. If this is being populated after running the task, then, we can simplify the first E2E to use LoadMode.MANIFEST, and assume the compiled SQL is avaialble for Cosmos.
  • If the dbt profile is not related to Databricks, error
  • if it is related to Databricks, than attempt to use the DatabricksSubmitRunOperator to run the previously compiled sql.
• Log follow-up tasks to make this PoC production ready and to expand the scope of the first end-to-end (support Python models, support other backend, e.g. Snowflake or BQ)
• Identify limitations with this approach (e.g. some macros may not work, and how this will play along with Airflow inlets/outlets)

[tatiana]

One of the advantages of this feature would also be to give additional information to end-users that Airflow currently exposes, but dbt does not.

For example, by using dbt, the BigQuery job ID is not printed in the logs, but if we changed Cosmos to use the Airflow BigQuery operator to run the BQ SQL transformations, the logs include the full job ID. We could also aim to expose this information in the Airflow UI.

This was brought up in a discussion with a customer:
https://astronomer.slack.com/archives/C074NP4A8G1/p1722774197657379

[tatiana]

@pankajkoti @pankajastro will be working closely together on discussing possibilities, trying them out and deciding on a path to implement this. They will be probably creating sub-tickets as needed while this work advances

[tatiana]

Another comment: I don't think we need to worry about handling incremental models where the code compiles differently based on whether the model already exists during the PoC. This could be a feature limitation.

[tatiana]

If we want to leverage the interest from Astronomer Cosmos customer, we should do this for BQ as starting point and make sure it works for DBT LS.

Steps discussed during brainstorming

1. Change profile mapping & change this DAG to use a new execution mode ExecutionMode.LOCAL_AIRFLOW_ASYNC
   https://github.com/astronomer/astronomer-cosmos/blob/main/dev/dags/basic_cosmos_dag.py

2. Have the compiled file & make it available to all dbt operators created by Cosmos

We need to confirm in which files we can see the compiled SQL.

a) Run dbt compile in the Airflow scheduler. Introduce LoadMode.DBT_COMPILE. It will work very similar to LoadMode.DBT_LS (from the perspective of running a subprocess with compile) and LoadMode.DBT_MANIFEST (from the perspective we'd parse the manifest file we create and just run the same logic as for LoadModeManifest). In this case, we'd need to upload the created artefacts to an Object Store, and make sure the paths would be available for the workers. This will have the same performance issues as LoadMode.DBT_LS before Cosmos 1.5.

b) Check the effort of resolving the compiled SQL from the Manifest file or another pre-existing dbt file that is created as part of LoadMode.DBT_LS.

c) We don't mess up with generating the compile files. We expect users to give a path to the compiled files the same way they give the path to the manifest. Then, we fetch the compiled SQLs.

d) Run dbt compile as a setup task at the beginning of the DAG. ProjectConfig(pre_compile_sql=True, remote_root_path, remote_target_conn) first task of Cosmos could be a task that runs dbt compile and store it in a remote store. We'd have to share this location with the other tasks - this can be done either upfront based on the user configuration. If we decide to manage the path, we'd be using something like <dbt-project-path> / <target> / <compiled> / * or <dag> / <task group> / <target> / <compiled>.

[tatiana]

We agreed not to do 2(a) because it will revert to performance issues we faced in the past.

[pankajkoti]

To evaluate approach 2(b) mentioned in comment, I explored on where all we could get the compiled SQLs from. I looked into the target directory of the dbt project and the files available in it. It contains the following files I explored for finding compiled SQLs:

1. graph.gpickle -> The purpose of the graph.gpickle file is to store the structure of the project’s dependency graph (DAG), including relationships between models, tests, seeds, and other nodes. It does not contain the compiled SQLs.
2. partial_parse.msgpack -> partial_parse.msgpack file is a binary-encoded file that dbt uses to store the results of previous parsing operations. It allows dbt to skip parsing unchanged files on subsequent runs by loading the cached metadata from this file. This one too does not contain compiled SQLs.
3. manifest.json -> The manifest.json does not always contain compiled SQLs. I ran the dbt compile command and observed that the compiled_sql still not populated for the model nodes. The dbt docs suggests that A few node properties, such as compiled_sql, only appear for executed nodes. That would be mean we would have to run at least once the dbt project, hence this does not seem to be a reliable way for us to start on, IMO.
4. Upon running the dbt compile command, a compiled directory is created in the target directory of the dbt project. This nested folder does contain the compiled SQLs.

Eventually, we can look forward to also supporting approach "2(c)" but since that would mean we're expecting users to supply the path to the compiled SQLs, I'm inclined towards taking the approach "2(d)" mentioned in comment so that Cosmos takes care of generating the compiled SQLs for the user and is an easier start for them then.