From c5f3cdd90f4bb9d066911709b0152b715822fb85 Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Wed, 5 Aug 2026 08:36:56 -0700 Subject: [PATCH] Restore gc and re imports dropped in the nvfuser_direct port Two stdlib imports were lost when code was copied from the legacy nvfuser package into nvfuser_direct, and both call sites now raise NameError. python/nvfuser_direct/pytorch_utils.py calls gc.collect() in retry_on_oom_or_skip_test but never imports gc. The legacy python/nvfuser/pytorch_utils.py had import gc. That decorator is applied to every collected python benchmark in benchmarks/python/conftest.py and to an opinfo test in tests/python/opinfo/test_direct_ops.py, so an OOM there currently dies with NameError: name 'gc' is not defined instead of clearing the cache, retrying, and skipping. The retry never runs at all. python/nvfuser_direct/__init__.py calls re.sub twice in repro_script_for but never imports re. The legacy python/nvfuser/__init__.py had import re. That code path is the non-tensor input branch, so any fd.repro_script_for([..., scalar]) or last_repro_script() after execute(save_repro_inputs=True) with a scalar input raises NameError instead of emitting a repro script. Neither was caught by lint because .flake8 ignores F821, and __init__.py also carries noqa F403 for its star import. Adds two regression tests in tests/python/direct/test_python_direct.py. Both need only the built extension, no GPU: torch.OutOfMemoryError is raised directly and torch.cuda.empty_cache() is a no-op when CUDA is not initialized, and repro_script_for does not require the inputs to match the fusion. Signed-off-by: Aditya Singh --- python/nvfuser_direct/__init__.py | 1 + python/nvfuser_direct/pytorch_utils.py | 3 +- tests/python/direct/test_python_direct.py | 40 +++++++++++++++++++++++ 3 files changed, 43 insertions(+), 1 deletion(-) diff --git a/python/nvfuser_direct/__init__.py b/python/nvfuser_direct/__init__.py index 5dd20d078e0..f5b9c717252 100644 --- a/python/nvfuser_direct/__init__.py +++ b/python/nvfuser_direct/__init__.py @@ -7,6 +7,7 @@ import warnings from typing import Iterable, Optional import functools +import re if "nvfuser" in sys.modules: warnings.warn( diff --git a/python/nvfuser_direct/pytorch_utils.py b/python/nvfuser_direct/pytorch_utils.py index a85c0bbcfc0..11387eb5def 100644 --- a/python/nvfuser_direct/pytorch_utils.py +++ b/python/nvfuser_direct/pytorch_utils.py @@ -6,8 +6,9 @@ from ._C_DIRECT import DataType import ctypes -from typing import Type, Union, Tuple import functools +import gc +from typing import Type, Union, Tuple NumberTypeType = Union[Type[bool], Type[int], Type[float], Type[complex]] diff --git a/tests/python/direct/test_python_direct.py b/tests/python/direct/test_python_direct.py index fdc1d4be0dd..2fe2f2ff27a 100644 --- a/tests/python/direct/test_python_direct.py +++ b/tests/python/direct/test_python_direct.py @@ -342,6 +342,26 @@ def nvfuser_fusion(fd : FusionDefinition) -> None : assert repro_with_inputs == last_repro +def test_repro_script_for_non_tensor_inputs(): + # test_repro_script_for above only passes tensors, so the non-tensor branch + # of repro_script_for is never exercised. That branch rewrites inf and nan + # into float("inf") and float("nan") so the emitted script is valid Python. + with FusionDefinition() as fd: + tv0 = fd.define_tensor( + shape=[-1], + contiguity=[True], + dtype=DataType.Float, + ) + s0 = fd.define_scalar(dtype=DataType.Float) + fd.add_output(fd.ops.mul(tv0, s0)) + + repro = fd.repro_script_for([2.5, float("inf"), float("nan"), -float("inf")]) + assert " 2.5,\n" in repro + assert ' float("inf"),\n' in repro + assert ' float("nan"),\n' in repro + assert ' -float("inf"),\n' in repro + + def test_define_tensor(): with FusionDefinition() as fd: tv0 = fd.define_tensor( @@ -689,3 +709,23 @@ def nvfuser_fusion_id1(fd: FusionDefinition) -> None: assert torch.allclose( fd1.execute([full_input, bcast_input])[0], bcast_input - full_input ) + + +def test_retry_on_oom_or_skip_test(): + # The decorator wraps every python benchmark in benchmarks/python/conftest.py + # and the opinfo tests, so its recovery path has to work. Raise + # torch.OutOfMemoryError once and check the wrapped function is retried + # rather than the decorator itself blowing up. + from nvfuser_direct.pytorch_utils import retry_on_oom_or_skip_test + + calls = [] + + @retry_on_oom_or_skip_test + def flaky(): + calls.append(1) + if len(calls) == 1: + raise torch.OutOfMemoryError("simulated OOM") + return "second attempt" + + assert flaky() == "second attempt" + assert len(calls) == 2