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]

Greptile Summary

Restores the missing re and gc standard-library imports in nvfuser_direct, fixing the non-tensor repro-script and OOM-retry paths.

• Adds regression coverage for special floating-point values in generated repro scripts.
• Adds CPU-compatible coverage confirming that an OOM causes one cleanup-and-retry cycle.

Confidence Score: 5/5

The PR appears safe to merge, with the restored imports directly satisfying existing reachable call sites and the regression tests covering both failures.

The changes only restore standard-library bindings already required by existing code, and the added tests exercise the previously failing branches without introducing a concrete behavioral regression.

Important Files Changed

Filename	Overview
python/nvfuser_direct/init.py	Restores the re import required by special-value rewriting in FusionDefinition.repro_script_for.
python/nvfuser_direct/pytorch_utils.py	Restores the gc import required by the existing OOM recovery decorator.
tests/python/direct/test_python_direct.py	Adds focused regression tests for non-tensor repro serialization and retry-after-OOM behavior.

_{Reviews (1): Last reviewed commit: "Restore gc and re imports dropped in the..." | Re-trigger Greptile}