[TEST] Move temp files into tempdir

python/tvm/contrib/utils.py

@@ -136,6 +136,12 @@ def __truediv__(self, other):
 
         return self.path / other
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, ptype, value, trace):
+        self.remove()
+
     def __del__(self):
         temp_dirs = getattr(self, "TEMPDIRS", None)
         if temp_dirs is None:

tests/python/codegen/test_target_codegen_llvm.py

@@ -805,7 +805,8 @@ def test_llvm_crt_static_lib():
         target=tvm.target.Target("llvm"),
     )
     module.get_source()
-    module.save("test.o")
+    with utils.tempdir() as temp:
+        module.save(temp.relpath("test.o"))
 
 
 @tvm.testing.requires_llvm

tests/python/relax/test_transform_codegen_pass.py

@@ -24,6 +24,7 @@
 import tvm
 import tvm.testing
 from tvm import relax, tir
+from tvm.contrib import utils
 from tvm.relax.dpl import is_op, wildcard
 from tvm.relax.testing import transform
 from tvm.script import ir as I
@@ -58,9 +59,9 @@ def check_executable(exec, dev, inputs, expected, entry_func_name):
 
 
 def check_roundtrip(exec0, dev, inputs, expected, entry_func_name="main"):
-    exec0.mod.export_library("exec.so")
-    exec1 = tvm.runtime.load_module("exec.so")
-    os.remove("exec.so")
+    with utils.tempdir() as temp:
+        exec0.mod.export_library(temp.relpath("exec.so"))
+        exec1 = tvm.runtime.load_module(temp.relpath("exec.so"))
     assert exec0.stats() == exec1["stats"]()
     assert exec0.as_text() == exec1["as_text"]()
 

tests/python/relax/test_vm_build.py

@@ -663,8 +663,10 @@ def te_func(A):
     target = tvm.target.Target("llvm", host="llvm")
     ex = relax.build(mod, target, exec_mode=exec_mode)
 
-    ex.export_library("exec.so")
-    vm = relax.VirtualMachine(tvm.runtime.load_module("exec.so"), tvm.cpu())
+    with utils.tempdir() as temp:
+        ex.export_library(temp.relpath("exec.so"))
+        vm = relax.VirtualMachine(tvm.runtime.load_module(temp.relpath("exec.so")), tvm.cpu())
+
     inp = tvm.nd.array(np.random.rand(2).astype(np.float32))
     inp2 = tvm.nd.array(np.random.rand(3).astype(np.float32))
 
@@ -956,7 +958,7 @@ def test_vm_mul(x: T.handle, y: T.handle, z: T.handle):
     # test returning a tuple
     @R.function
     def test_vm_tuple(
-        x: R.Tensor((), "int32")
+        x: R.Tensor((), "int32"),
     ) -> R.Tuple(R.Tensor((), "int32"), R.Tensor((), "int32")):
         return (x, x)