Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Data] Avoid memory leak with pyarrow.infer_type on datetime arrays #50403

Merged
merged 32 commits into from
Feb 13, 2025
Merged

[Data] Avoid memory leak with pyarrow.infer_type on datetime arrays #50403

merged 32 commits into from
Feb 13, 2025

Conversation

bveeramani
Copy link
Member

Why are these changes needed?

pyarrow.infer_type leaks memory when inferring a datetime array, and this leak can contribute to OOMs. This PR worksaround the issue by manually inferring the type for datetime arrays.

Related issue number

apache/arrow#45493

Checks

  • I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
  • I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
    • I've added any new APIs to the API Reference. For example, if I added a
      method in Tune, I've added it in doc/source/tune/api/ under the
      corresponding .rst file.
  • I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
  • Testing Strategy
    • Unit tests
    • Release tests
    • This PR is not tested :(

@bveeramani
Initial commit
cb830b0 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
1d9dd2e 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
bed2558 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
c40b06d 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
49a8d71 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
cad41d5 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
17df280 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
bb0e778 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
2620239 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
4dff29a 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
0393f97 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
2a92d8e 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
7d912a8 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
e38bf85 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
01c8911 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
a7e6cb2 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
53dfa51 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
27b1315 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
ba39108 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
6f9d698 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
11c498f 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
ed5fd50 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
1209f80 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
f992321 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
59a732b 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Merge branch 'master' of https://github.com/ray-project/ray
29e12e4 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Initial commit
8b60041 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Fix typo
f0d2973 
Signed-off-by: Balaji Veeramani <[email protected]>
alexeykudinkin
alexeykudinkin reviewed Feb 11, 2025
View reviewed changes
python/ray/air/util/tensor_extensions/arrow.py Show resolved Hide resolved
python/ray/air/tests/test_arrow.py
Comment on lines +56 to +71
def test_infer_type_does_not_leak_memory(dtype):
# Test for https://github.com/apache/arrow/issues/45493.
column_values = np.zeros(923040, dtype=dtype) # A ~7 MiB column

process = psutil.Process()
gc.collect()
pa.default_memory_pool().release_unused()
before = process.memory_info().rss

_infer_pyarrow_type(column_values)

gc.collect()
pa.default_memory_pool().release_unused()
after = process.memory_info().rss

assert after - before < 1024 * 1024, after - before
Copy link
Contributor

Choose a reason for hiding this comment

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

Does it fail w/o your changes? What's amount of memory leaked per call?

Copy link
Member Author

Choose a reason for hiding this comment

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

Yeah. Like 130 MiB lol

Copy link
Member Author

Choose a reason for hiding this comment

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

If you haven't already, see the repro in the commented issue: apache/arrow#45493

@bveeramani
Fix bug
0cdef0d 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani
Address review comments
beb229b 
Signed-off-by: Balaji Veeramani <[email protected]>
@bveeramani bveeramani enabled auto-merge (squash) February 12, 2025 20:55
@github-actions github-actions bot added the go add ONLY when ready to merge, run all tests label Feb 12, 2025
@bveeramani
Fix bug
ec2cf96 
Signed-off-by: Balaji Veeramani <[email protected]>
@github-actions github-actions bot disabled auto-merge February 12, 2025 23:11
@bveeramani bveeramani enabled auto-merge (squash) February 12, 2025 23:12
@bveeramani
Fix test
7ad4a1d 
Signed-off-by: Balaji Veeramani <[email protected]>
@github-actions github-actions bot disabled auto-merge February 12, 2025 23:15
@bveeramani bveeramani enabled auto-merge (squash) February 12, 2025 23:15
@bveeramani bveeramani merged commit f2f5e8e into master Feb 13, 2025
6 checks passed
@bveeramani bveeramani deleted the infer-type-leak branch February 13, 2025 00:09
@bveeramani bveeramani mentioned this pull request Feb 13, 2025
Open
8 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@alexeykudinkin alexeykudinkin alexeykudinkin approved these changes

Assignees
No one assigned
Labels
go add ONLY when ready to merge, run all tests
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@bveeramani @alexeykudinkin