From c5f3cdd90f4bb9d066911709b0152b715822fb85 Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Wed, 5 Aug 2026 08:36:56 -0700 Subject: [PATCH 1/2] 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 From b6ab75f05d6187a5fa896b228daf47a7e5fd000c Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Wed, 16 Sep 2026 22:00:21 -0700 Subject: [PATCH 2/2] Make the repro regression test inputs match the fusion The fusion in test_repro_script_for_non_tensor_inputs takes one tensor and one scalar, but the test passed four scalars, so the emitted script was not a valid reproduction of that fusion. Pass a matching tensor plus one scalar, execute the fusion with that list to prove it is valid, and assert the exact inputs block for 2.5, inf, nan and -inf. Co-Authored-By: Claude Fable 5.1 --- tests/python/direct/test_python_direct.py | 26 ++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/python/direct/test_python_direct.py b/tests/python/direct/test_python_direct.py index 2fe2f2ff27a..5d71d57d312 100644 --- a/tests/python/direct/test_python_direct.py +++ b/tests/python/direct/test_python_direct.py @@ -346,6 +346,8 @@ 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. + # The inputs match the fusion, one tensor and one scalar, so the emitted + # script is a reproduction that can actually be run against this fusion. with FusionDefinition() as fd: tv0 = fd.define_tensor( shape=[-1], @@ -355,11 +357,25 @@ def test_repro_script_for_non_tensor_inputs(): 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 + tensor_inp = torch.randn(4, dtype=torch.float32, device="cuda:0") + for scalar_inp, expected_scalar in [ + (2.5, "2.5"), + (float("inf"), 'float("inf")'), + (float("nan"), 'float("nan")'), + (-float("inf"), '-float("inf")'), + ]: + inputs = [tensor_inp, scalar_inp] + # The input list is valid for the fusion, so the repro is executable. + fd.execute(inputs) + + expected_inputs = ( + "\ninputs = [\n" + " torch.testing.make_tensor((4,), dtype=torch.float32, device='cuda:0'),\n" + f" {expected_scalar},\n" + "]\n" + "fd.execute(inputs)\n" + ) + assert expected_inputs in fd.repro_script_for(inputs) def test_define_tensor():