Skip to content

Comments

fix: get task dependencies without serializing task tree to string#41494

Merged
mobuchowski merged 1 commit intoapache:mainfrom
mobuchowski:fix-ol-tree-generation
Aug 15, 2024
Merged

fix: get task dependencies without serializing task tree to string#41494
mobuchowski merged 1 commit intoapache:mainfrom
mobuchowski:fix-ol-tree-generation

Conversation

@mobuchowski
Copy link
Contributor

@mobuchowski mobuchowski commented Aug 15, 2024

Currently _get_parsed_dag_tree uses get_tree_view which in degenerated case (like, in test test_get_dag_tree_large_dag can generate string tree representation taking multiple gigabytes of memory.

However, for what it's trying to accomplish, pretty much any temporary allocations are not necessary. This fix constructs task dependency tree without intermediate representation.

Previous memory consumption:

root@a24bae3584cb:/opt/airflow# pytest --memray tests/providers/openlineage/utils/test_utils.py::test_get_dag_tree_large_dag
=========================================================================================================================================================================== test session starts ============================================================================================================================================================================
platform linux -- Python 3.12.5, pytest-8.3.2, pluggy-1.5.0 -- /usr/local/bin/python
cachedir: .pytest_cache
rootdir: /opt/airflow
configfile: pyproject.toml
plugins: memray-1.7.0, timeouts-1.2.1, icdiff-0.9, mock-3.14.0, rerunfailures-14.0, requests-mock-1.12.1, xdist-3.6.1, asyncio-0.23.8, anyio-4.4.0, instafail-0.5.0, cov-5.0.0, time-machine-2.15.0, custom-exit-code-0.3.0
asyncio: mode=Mode.STRICT
setup timeout: 0.0s, execution timeout: 0.0s, teardown timeout: 0.0s
collected 1 item

tests/providers/openlineage/utils/test_utils.py::test_get_dag_tree_large_dag PASSED                                                                                                                                                                                                                                                                                  [100%]


============================================================================================================================================================================== MEMRAY REPORT ===============================================================================================================================================================================
Allocation results for tests/providers/openlineage/utils/test_utils.py::test_get_dag_tree_large_dag at the high watermark

	 📦 Total memory allocated: 5.4GiB
	 📏 Total allocations: 23
	 📊 Histogram of allocation sizes: |▁▁█  |
	 🥇 Biggest allocating functions:
		- _safe_get_dag_tree_view:/opt/airflow/airflow/providers/openlineage/utils/utils.py:446 -> 2.7GiB
		- get_tree_view:/opt/airflow/airflow/models/dag.py:2445 -> 2.7GiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB


=================================================================================================================================================================== Warning summary. Total: 3, Unique: 3 ===================================================================================================================================================================
airflow: total 1, unique 1
  collect: total 1, unique 1
other: total 2, unique 2
  collect: total 2, unique 2
Warnings saved into /opt/airflow/tests/warnings.txt file.
============================================================================================================================================================================ 1 passed in 8.60s =============================================================================================================================================================================

current memory consumption:

root@2788b43cb914:/opt/airflow# pytest --memray tests/providers/openlineage/utils/test_utils.py::test_get_dag_tree_large_dag
=========================================================================================================================================================================== test session starts ============================================================================================================================================================================
platform linux -- Python 3.12.5, pytest-8.3.2, pluggy-1.5.0 -- /usr/local/bin/python
cachedir: .pytest_cache
rootdir: /opt/airflow
configfile: pyproject.toml
plugins: memray-1.7.0, timeouts-1.2.1, icdiff-0.9, mock-3.14.0, rerunfailures-14.0, requests-mock-1.12.1, xdist-3.6.1, asyncio-0.23.8, anyio-4.4.0, instafail-0.5.0, cov-5.0.0, time-machine-2.15.0, custom-exit-code-0.3.0
asyncio: mode=Mode.STRICT
setup timeout: 0.0s, execution timeout: 0.0s, teardown timeout: 0.0s
collected 1 item

tests/providers/openlineage/utils/test_utils.py::test_get_dag_tree_large_dag PASSED                                                                                                                                                                                                                                                                                  [100%]


============================================================================================================================================================================== MEMRAY REPORT ===============================================================================================================================================================================
Allocation results for tests/providers/openlineage/utils/test_utils.py::test_get_dag_tree_large_dag at the high watermark

	 📦 Total memory allocated: 16.0MiB
	 📏 Total allocations: 30
	 📊 Histogram of allocation sizes: |▁█ ▃▇|
	 🥇 Biggest allocating functions:
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB
		- __setattr__:/opt/airflow/airflow/models/baseoperator.py:1191 -> 1.3MiB


=================================================================================================================================================================== Warning summary. Total: 3, Unique: 3 ===================================================================================================================================================================
airflow: total 1, unique 1
  collect: total 1, unique 1
other: total 2, unique 2
  collect: total 2, unique 2
Warnings saved into /opt/airflow/tests/warnings.txt file.
============================================================================================================================================================================ 1 passed in 2.49s =============================================================================================================================================================================

kacpermuda
kacpermuda approved these changes Aug 15, 2024
View reviewed changes
Copy link
Contributor

@kacpermuda kacpermuda left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks so much better than my previous implementation 🤦 Thanks @mobuchowski 🚀 !

@mobuchowski mobuchowski merged commit 0e7c757 into apache:main Aug 15, 2024
@jlaneve
Copy link
Contributor

jlaneve commented Aug 15, 2024

Nice! Do you know if there's an issue open to track the underlying function? This fixed it for OL but not for other consumers of that function

@mobuchowski mobuchowski mentioned this pull request Aug 15, 2024
Closed
2 tasks
@mobuchowski
Copy link
Contributor Author

mobuchowski commented Aug 15, 2024

@jlaneve just created it. #41505

Artuz37 pushed a commit to Artuz37/airflow that referenced this pull request Aug 19, 2024
@mobuchowski @Artuz37
fix: get task dependencies without serializing task tree to string (a…
4ea24cd 
…pache#41494)

Signed-off-by: Maciej Obuchowski <[email protected]>
@eladkal eladkal mentioned this pull request Aug 19, 2024
Closed
43 tasks
@kacpermuda kacpermuda mentioned this pull request Aug 19, 2024
Merged
romsharon98 pushed a commit to romsharon98/airflow that referenced this pull request Aug 20, 2024
@mobuchowski @romsharon98
fix: get task dependencies without serializing task tree to string (a…
2db1ef9 
…pache#41494)

Signed-off-by: Maciej Obuchowski <[email protected]>
@eladkal eladkal mentioned this pull request Aug 25, 2024
Closed
20 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@utkarsharma2 utkarsharma2 utkarsharma2 approved these changes

+2 more reviewers

@JDarDagran JDarDagran JDarDagran approved these changes

@kacpermuda kacpermuda kacpermuda approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

area:providers provider:openlineage AIP-53

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@mobuchowski @jlaneve @JDarDagran @utkarsharma2 @kacpermuda