Restore gc and re imports dropped in the nvfuser_direct port

[adityasingh2400]

Two stdlib imports were lost when code was copied from the legacy nvfuser package into nvfuser_direct. Both call sites raise NameError today.

gc in python/nvfuser_direct/pytorch_utils.py

retry_on_oom_or_skip_test calls gc.collect() at line 64, and the file's complete import list is torch, DataType from ._C_DIRECT, ctypes, typing, and functools. There is no import gc. The import is not coming in sideways either, since the extension import is an explicit from ._C_DIRECT import DataType rather than a star import.

The legacy python/nvfuser/pytorch_utils.py had import gc on line 10, alongside the identical gc.collect() call. It was dropped when the function was ported.

The blast radius is the whole benchmark suite. benchmarks/python/conftest.py line 201 wraps every collected item with this decorator, and tests/python/opinfo/test_direct_ops.py line 219 applies it to an op test. So today an OOM in any of them raises NameError: name 'gc' is not defined instead of clearing the cache, retrying, and skipping on a second OOM. The retry never runs at all, which means the decorator's entire purpose is defeated and the failure it produces points at the wrong thing.

re in python/nvfuser_direct/__init__.py

FusionDefinition.repro_script_for calls re.sub at lines 483 and 486 and the file never imports re. The legacy python/nvfuser/__init__.py had import re on line 16 with the same two calls.

That code is the non-tensor input branch, the one that rewrites inf and nan into float("inf") and float("nan") so the emitted script is valid Python. Any fd.repro_script_for([..., scalar]), or fd.last_repro_script() after fd.execute(inputs, save_repro_inputs=True) with a scalar input, raises NameError instead of returning a repro script. The existing test_repro_script_for passes only CUDA tensors plus a constant scalar, which define_scalar does not register as a fusion input, so this branch has no coverage.

Why lint did not catch either

.flake8 lists F821 among its ignores, so undefined names are invisible to lintrunner. __init__.py additionally carries # noqa: F401,F403 on its from ._C_DIRECT import * line, and F403 is precisely the code meaning flake8 can no longer detect undefined names in that file.

I swept the whole Python tree with an AST pass that reports any stdlib module used as mod.attr with no matching binding in scope, which avoids the star-import false positives that make a raw F821 run unusable here. Across python/, tests/python/, benchmarks/python/, and tools/ it reports exactly these two and nothing else, so this is the complete set.

Tests

Two regression tests in tests/python/direct/test_python_direct.py.

test_retry_on_oom_or_skip_test raises torch.OutOfMemoryError on the first call and asserts the wrapped function is invoked a second time and its value returned.

test_repro_script_for_non_tensor_inputs passes 2.5, float("inf"), float("nan"), and -float("inf") to repro_script_for and asserts each is emitted in a form that is valid Python.

Neither needs a GPU. torch.OutOfMemoryError is raised directly rather than provoked, torch.cuda.empty_cache() is a no-op when CUDA is not initialized, repro_script_for does not require its inputs argument to match the fusion, and the torch.cuda.device_count() loop at the top of that method simply iterates zero times. They do need the built extension, like the rest of the file.

Verification

Both fixes were verified fail-before and pass-after on this machine with no GPU and no CUDA build, by loading the real source files through importlib against a stub _C_DIRECT module. With python/nvfuser_direct/__init__.py and python/nvfuser_direct/pytorch_utils.py checked out from main, the decorator raises NameError: name 'gc' is not defined after exactly one call, and repro_script_for raises NameError: name 're' is not defined. With the patched files, the decorator retries and returns on the second attempt, and repro_script_for emits all four non-tensor inputs with inf and nan correctly rewritten.

Formatting checked with black 23.3.0 and flake8 6.1.0 against .flake8, the versions pinned in .lintrunner.toml. Both are clean and black leaves all three files unchanged. Re-running flake8 --select=F821 over python/nvfuser_direct/ now reports nothing.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge.

Summary

Restores missing gc and re standard-library imports in nvfuser_direct, preventing NameError failures in OOM recovery and scalar repro-script generation.

• Adds re for non-tensor repro input serialization.
• Adds gc for the OOM retry cleanup path.
• Adds regression coverage for both restored imports.

_{Reviews (2) · Last reviewed commit: "Make the repro regression test inputs ma..."}

[sylvesterkaczmarek]

Could the repro regression use inputs that actually match the fusion it constructs?

The fusion has one tensor input plus one scalar input, but the test calls repro_script_for([2.5, inf, nan, -inf]). That is enough to exercise the missing re import, but the resulting repro script is not a valid reproduction of this fusion because its input list has the wrong shape and arity.

Using a matching tensor plus a non-tensor special value would still cover the inf/nan serialization path while also ensuring the test represents a script a user could actually execute.

[adityasingh2400]

@sylvesterkaczmarek agreed, the old input list had the wrong arity for the fusion it built. Fixed in b6ab75f.

The test now passes one (4,) float32 CUDA tensor plus one scalar, which is what the fusion declares, and loops the scalar over 2.5, inf, nan and -inf. For each case it first calls fd.execute(inputs) so the list is proven valid for the fusion, then asserts the exact inputs = [...] block the script emits: the make_tensor((4,), ...) line followed by 2.5, float("inf"), float("nan") or -float("inf"). That still goes through the non-tensor branch of repro_script_for, which is where the restored re import is used.

On verification: I do not have a GPU build, so this was not run under pytest. I checked the file with py_compile and black, and reproduced the non-tensor formatting path in isolation with the same two re.sub calls to confirm the four expected strings. The tensor line matches the format test_repro_script_for already asserts for a contiguous CUDA tensor.

[sylvesterkaczmarek]

Re-reviewed current b6ab75f. The regression now uses a valid fusion input list: one CUDA tensor plus one scalar, executes it first, and then checks the emitted scalar serialization for finite, Inf, NaN, and -Inf. That resolves my earlier concern about testing an invalid reproduction. No remaining blocker from my review.