Introduce @retry_on_oom_or_skip_test

[jacobhinkle]

Inspired by #3174

This is an alternative to #3238.

Previously we were manually resetting the cuda cache whenever the usage was above 80%. This is not ideal since we could have 79% usage and a test that requires 25% and that would fail. We also might clear the cache unnecessarily sometimes: e.g. we are using 81% but only need a few percent for the remainder of tests.

This PR cleans this up by introducing a new test decorator @retry_on_oom_or_skip_test. This decorator must be placed innermost, underneath the other decorators. It will execute the test inside a try block. If the test fails due to torch.OutOfMemoryError, we clear the cuda cache and retry the test. If it fails again due to torch.OutOfMemoryError, then we skip the test.

I updated the python benchmarks to apply this decorator automatically, and to remove the manual clear_cuda_cache() calls.

[jacobhinkle]

!build

[Priya2698]

This is great idea.
We may be able to move this decorator or wrap it into a pytest fixture and turn on autouse: https://docs.pytest.org/en/stable/how-to/fixtures.html#autouse-fixtures-fixtures-you-don-t-have-to-request.
This would no longer require us to individually mark all tests and do not have to think about the order of the decorators in the tests.

[jacobhinkle]

This is great idea. We may be able to move this decorator or wrap it into a pytest fixture and turn on autouse: https://docs.pytest.org/en/stable/how-to/fixtures.html#autouse-fixtures-fixtures-you-don-t-have-to-request. This would no longer require us to individually mark all tests and do not have to think about the order of the decorators in the tests.

Interesting idea. I can't see yet how to make it actually retry the calling function, but it might be possible. Another option might be https://github.com/str0zzapreti/pytest-retry, but I couldn't see how to get that to run the gc/cache clear step before the retry.

[jacobhinkle]

!build

[Priya2698]

Will the changes in the this file be superseded by the other PR?

[jacobhinkle]

I will merge it manually

[jacobhinkle]

Merged

[jacobhinkle]

!build

[Priya2698]

LGTM -- this is a great improvement!
test_matmul still seems to have some changes though which are likely not part of this PR that may need to be cleaned manually.

[jacobhinkle]

test_matmul still seems to have some changes though which are likely not part of this PR that may need to be cleaned manually.

It does? I just removed the try block. Is there something else you noticed?

[Priya2698]

test_matmul still seems to have some changes though which are likely not part of this PR that may need to be cleaned manually.

It does? I just removed the try block. Is there something else you noticed?

Oh you're right. I misunderstood.
The PR can be merged as-is.

[jacobhinkle]

!build

[jacobhinkle]

The jit_python_bc_advisory_17_A100 failure is real. I will fix it before merging

[jacobhinkle]

!build

[jacobhinkle]

The jit_python_bc_advisory_17_A100 failure is real. I will fix it before merging

The error seems strange, like CI was using an older test_ops.py which was trying to import clear_cuda_cache. I'm trying it again. I also fixed test_matmul.py.

[jacobhinkle]

!build

[jacobhinkle]

The error seems strange, like CI was using an older test_ops.py which was trying to import clear_cuda_cache. I'm trying it again. I also fixed test_matmul.py.

This is actually intended behavior. That job runs the tests from merge-base to detect changes in Python API. In this case it detects that we have changed nvfuser/python_utils.py to remove clear_cuda_cache. Since that is a user-facing library, I bumped version.txt in the latest push.