Unify cpp_extension build directory removal (pytorch#136059)

Keeps existing default directory clearing logic, even though it fails when TORCH_EXTENSIONS_DIR is set. To properly clear, we'd need to track all the folders we compiled the extensions to.

Pull Request resolved: pytorch#136059
Approved by: https://github.com/ezyang, https://github.com/albanD

test/inductor/test_extension_backend.py

@@ -1,6 +1,5 @@
 # Owner(s): ["module: inductor"]
 import os
-import shutil
 import sys
 import unittest
 
@@ -49,15 +48,6 @@
 TestCase = test_torchinductor.TestCase
 
 
-def remove_build_path():
-    if sys.platform == "win32":
-        # Not wiping extensions build folder because Windows
-        return
-    default_build_root = torch.utils.cpp_extension.get_default_build_root()
-    if os.path.exists(default_build_root):
-        shutil.rmtree(default_build_root, ignore_errors=True)
-
-
 @unittest.skipIf(IS_FBCODE, "cpp_extension doesn't work in fbcode right now")
 class ExtensionBackendTests(TestCase):
     module = None
@@ -67,7 +57,7 @@ def setUpClass(cls):
         super().setUpClass()
 
         # Build Extension
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
         source_file_path = os.path.dirname(os.path.abspath(__file__))
         source_file = os.path.join(
             source_file_path, "extension_backends/cpp/extension_device.cpp"
@@ -86,7 +76,7 @@ def tearDownClass(cls):
         cls._stack.close()
         super().tearDownClass()
 
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
     def setUp(self):
         torch._dynamo.reset()

test/profiler/test_cpp_thread.py

@@ -1,8 +1,6 @@
 # Owner(s): ["oncall: profiler"]
 
 import os
-import shutil
-import subprocess
 from unittest import skipIf
 
 import torch
@@ -11,17 +9,6 @@
 from torch.testing._internal.common_utils import IS_WINDOWS, run_tests, TestCase
 
 
-def remove_build_path():
-    default_build_root = torch.utils.cpp_extension.get_default_build_root()
-    if os.path.exists(default_build_root):
-        if IS_WINDOWS:
-            # rmtree returns permission error: [WinError 5] Access is denied
-            # on Windows, this is a word-around
-            subprocess.run(["rm", "-rf", default_build_root], stdout=subprocess.PIPE)
-        else:
-            shutil.rmtree(default_build_root)
-
-
 if is_fbcode():
     import caffe2.test.profiler_test_cpp_thread_lib as cpp  # @manual=//caffe2/test:profiler_test_cpp_thread_lib
 else:
@@ -93,7 +80,7 @@ def setUpClass(cls) -> None:
     @classmethod
     def tearDownClass(cls):
         if not is_fbcode():
-            remove_build_path()
+            torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
     def setUp(self) -> None:
         if not torch.cuda.is_available():

test/test_cpp_extensions_jit.py

@@ -35,18 +35,7 @@
 IS_LINUX = sys.platform.startswith("linux")
 
 
-def remove_build_path():
-    default_build_root = torch.utils.cpp_extension.get_default_build_root()
-    if os.path.exists(default_build_root):
-        if IS_WINDOWS:
-            # rmtree returns permission error: [WinError 5] Access is denied
-            # on Windows, this is a word-around
-            subprocess.run(["rm", "-rf", default_build_root], stdout=subprocess.PIPE)
-        else:
-            shutil.rmtree(default_build_root)
-
-
-# There's only one test that runs gracheck, run slow mode manually
+# There's only one test that runs gradcheck, run slow mode manually
 @torch.testing._internal.common_utils.markDynamoStrictTest
 class TestCppExtensionJIT(common.TestCase):
     """Tests just-in-time cpp extensions.
@@ -67,11 +56,11 @@ def tearDown(self):
 
     @classmethod
     def setUpClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
     @classmethod
     def tearDownClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
     def test_jit_compile_extension(self):
         module = torch.utils.cpp_extension.load(

test/test_cpp_extensions_mtia_backend.py

@@ -1,8 +1,6 @@
 # Owner(s): ["module: mtia"]
 
 import os
-import shutil
-import sys
 import tempfile
 import unittest
 
@@ -25,15 +23,6 @@
 TEST_CUDA = TEST_CUDA and CUDA_HOME is not None
 
 
-def remove_build_path():
-    if sys.platform == "win32":
-        # Not wiping extensions build folder because Windows
-        return
-    default_build_root = torch.utils.cpp_extension.get_default_build_root()
-    if os.path.exists(default_build_root):
-        shutil.rmtree(default_build_root, ignore_errors=True)
-
-
 @unittest.skipIf(
     IS_ARM64 or not IS_LINUX or TEST_CUDA or TEST_PRIVATEUSE1 or TEST_ROCM or TEST_XPU,
     "Only on linux platform and mutual exclusive to other backends",
@@ -58,11 +47,11 @@ def tearDown(self):
 
     @classmethod
     def tearDownClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
     @classmethod
     def setUpClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
         build_dir = tempfile.mkdtemp()
         # Load the fake device guard impl.
         cls.module = torch.utils.cpp_extension.load(

test/test_cpp_extensions_open_device_registration.py

@@ -2,8 +2,6 @@
 
 import _codecs
 import os
-import shutil
-import sys
 import tempfile
 import types
 import unittest
@@ -30,15 +28,6 @@
 TEST_ROCM = TEST_CUDA and torch.version.hip is not None and ROCM_HOME is not None
 
 
-def remove_build_path():
-    if sys.platform == "win32":
-        # Not wiping extensions build folder because Windows
-        return
-    default_build_root = torch.utils.cpp_extension.get_default_build_root()
-    if os.path.exists(default_build_root):
-        shutil.rmtree(default_build_root, ignore_errors=True)
-
-
 def generate_faked_module():
     def device_count() -> int:
         return 1
@@ -98,7 +87,7 @@ def tearDown(self):
 
     @classmethod
     def setUpClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
         cls.module = torch.utils.cpp_extension.load(
             name="custom_device_extension",

test/test_cpp_extensions_stream_and_event.py

@@ -1,8 +1,6 @@
 # Owner(s): ["module: mtia"]
 
 import os
-import shutil
-import sys
 import tempfile
 import unittest
 
@@ -26,15 +24,6 @@
 TEST_CUDA = TEST_CUDA and CUDA_HOME is not None
 
 
-def remove_build_path():
-    if sys.platform == "win32":
-        # Not wiping extensions build folder because Windows
-        return
-    default_build_root = torch.utils.cpp_extension.get_default_build_root()
-    if os.path.exists(default_build_root):
-        shutil.rmtree(default_build_root, ignore_errors=True)
-
-
 # Since we use a fake MTIA device backend to test generic Stream/Event, device backends are mutual exclusive to each other.
 # The test will be skipped if any of the following conditions are met:
 @unittest.skipIf(
@@ -67,11 +56,11 @@ def tearDown(self):
 
     @classmethod
     def tearDownClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
 
     @classmethod
     def setUpClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
         build_dir = tempfile.mkdtemp()
         # Load the fake device guard impl.
         src = f"{os.path.abspath(os.path.dirname(__file__))}/cpp_extensions/mtia_extension.cpp"

test/test_transformers.py

@@ -55,7 +55,6 @@
 
 if not IS_FBCODE:
     from test_cpp_extensions_open_device_registration import (
-        remove_build_path,
         generate_faked_module
     )
 
@@ -3849,7 +3848,7 @@ def test_is_causal_and_mask_fails(self, device):
 class TestSDPAPrivateUse1Only(NNTestCase):
     @classmethod
     def setUpClass(cls):
-        remove_build_path()
+        torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
         cls.module = torch.utils.cpp_extension.load(
             name="custom_device_extension",
             sources=[

torch/testing/_internal/common_utils.py

@@ -5359,3 +5359,17 @@ def match_obj(obj):
 
     finally:
         gc.set_debug(0)
+
+
+def remove_cpp_extensions_build_root():
+    """
+    Removes the default root folder under which extensions are built.
+    """
+    default_build_root = torch.utils.cpp_extension.get_default_build_root()
+    if os.path.exists(default_build_root):
+        if IS_WINDOWS:
+            # rmtree returns permission error: [WinError 5] Access is denied
+            # on Windows, this is a workaround
+            subprocess.run(["rm", "-rf", default_build_root], stdout=subprocess.PIPE)
+        else:
+            shutil.rmtree(default_build_root, ignore_errors=True)