Restore gc and re imports dropped in the nvfuser_direct port

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(

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]]
 

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