Emit Airflow metrics to support analysing Cosmos performance

[tatiana]

Context

We want more visibility on how much Cosmos spends while parsing the dbt project and building the Airflow DAG.

We'd like to leverage Airflow Metrics collection system by using:

Stats.timer("ol.emit.attempts")

To collect the following metrics:

• cosmos.load_method_custom.duration: time taken to run DbtGraph.load_via_custom_parser
• cosmos.load_method_dbt_ls.duration: time taken to run DbtGraph.load_via_dbt_ls
• cosmos.load_method_dbt_ls_file.duration: time taken to run DbtGraph.load_via_dbt_ls_file
• cosmos.load_method_manifest.duration: time taken to run DbtGraph.load_from_dbt_manifest
• cosmos.convert_to_airflow.duration: time taken to run `build_airflow_graph``
• cosmos.dag_init.duration: time taken to initialise the Airflow DAG
• cosmos.dag_new.duration: time taken to create the Airflow DAG
• cosmos.task_group_init.duration: time taken to initialise the Airflow DAG (__init__)
• cosmos.task_group_new.duration: time taken to create the Airflow DAG (__new__)

Relevant parts of the code:

astronomer-cosmos/cosmos/dbt/graph.py

Lines 168 to 171 in cda2a50

	LoadMode.CUSTOM: self.load_via_custom_parser,
	LoadMode.DBT_LS: self.load_via_dbt_ls,
	LoadMode.DBT_LS_FILE: self.load_via_dbt_ls_file,
	LoadMode.DBT_MANIFEST: self.load_from_dbt_manifest,

astronomer-cosmos/cosmos/airflow/graph.py

Line 215 in cda2a50

def build_airflow_graph(

https://github.com/astronomer/astronomer-cosmos/blob/main/cosmos/airflow/dag.py
https://github.com/astronomer/astronomer-cosmos/blob/main/cosmos/airflow/task_group.py

Acceptance criteria

• All these metrics are sent to statsd when running Cosmos DAGs, when Airflow is configured to do so

[dwreeves]

A few questions:

• Is it possible to inject the dag_id and task_group_id into the metric names, when appropriate?
  • DbtDags don't create task groups right? In that case it may be necessary to do something like replace task_group_id with self or dag or something like that, so the metric naming is a little more consistent.
• You have it as cosmos.load_method_custom.duration, cosmos.load_method_dbt_ls.duration, etc. but would it make sense to do something more like cosmos.load_graph.duration or cosmos.graph.{dag_id}.{task_group_id}.duration? My thinking is:
  • these load methods are unlikely to ever be mixed-and-matched, so distinguishing between them within a single Airflow deployment is not likely to be relevant.
  • as we introduce composability into Cosmos (Introducing composability in the middle layer of Cosmos's API #895, also Decouple LoadMode.AUTOMATIC from load() method in DbtGraph #1001), I also think it makes more sense to refer to things "generically" rather than referencing specific implementations. E.g. imagine a user creates a custom DbtGraph with a custom load method.

[tatiana]

Hey, @dwreeves, these are very valid points.

I'm improving the logs on a per DAG/TaskGroup as part of #1014 (e.g., https://github.com/astronomer/astronomer-cosmos/pull/1014/files#diff-61b585fb903927b6868b9626c95e0ec47e3818eb477d795ebd13b0276d4fd76cR293). This will probably be switched to DEBUG and be further improved, but this would help to address the granularity your suggestion. I'll probably create a PR only for this :)

The goal with having the metrics proposed in this PR is to really have a "group" that helps to have an overview of the health of these numbers across multiple DAGs - and help spot overall if any of these metrics are looking more troublesome than others. WDYT?