diff --git a/test/inductor/test_extension_backend.py b/test/inductor/test_extension_backend.py index 6f972e46a1d..4504d472fe2 100644 --- a/test/inductor/test_extension_backend.py +++ b/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 @@ run_and_get_cpp_code = test_torchinductor.run_and_get_cpp_code 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 @@ class ExtensionBackendTests(TestCase): 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 @@ class ExtensionBackendTests(TestCase): cls._stack.close() super().tearDownClass() - remove_build_path() + torch.testing._internal.common_utils.remove_cpp_extensions_build_root() def setUp(self): torch._dynamo.reset() diff --git a/test/profiler/test_cpp_thread.py b/test/profiler/test_cpp_thread.py index 1e7acc155ec..527fc5aca1c 100644 --- a/test/profiler/test_cpp_thread.py +++ b/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._environment import is_fbcode 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 @@ class CppThreadTest(TestCase): @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(): diff --git a/test/test_cpp_extensions_jit.py b/test/test_cpp_extensions_jit.py index 9b190b29f3a..844ca88dd25 100644 --- a/test/test_cpp_extensions_jit.py +++ b/test/test_cpp_extensions_jit.py @@ -35,18 +35,7 @@ IS_WINDOWS = sys.platform == "win32" 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 @@ class TestCppExtensionJIT(common.TestCase): @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( diff --git a/test/test_cpp_extensions_mtia_backend.py b/test/test_cpp_extensions_mtia_backend.py index 3b81344a3cd..6203b799328 100644 --- a/test/test_cpp_extensions_mtia_backend.py +++ b/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_ROCM = TEST_CUDA and torch.version.hip is not None and ROCM_HOME is not Non 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 @@ class TestCppExtensionMTIABackend(common.TestCase): @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( diff --git a/test/test_cpp_extensions_open_device_registration.py b/test/test_cpp_extensions_open_device_registration.py index 4e86ed458b0..e6883b0f2e9 100644 --- a/test/test_cpp_extensions_open_device_registration.py +++ b/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_CUDA = TEST_CUDA and CUDA_HOME is not None 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 @@ class TestCppExtensionOpenRgistration(common.TestCase): @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", diff --git a/test/test_cpp_extensions_stream_and_event.py b/test/test_cpp_extensions_stream_and_event.py index c26e8b2b1a8..f6b2281e171 100644 --- a/test/test_cpp_extensions_stream_and_event.py +++ b/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_ROCM = TEST_CUDA and torch.version.hip is not None and ROCM_HOME is not Non 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 @@ class TestCppExtensionStreamAndEvent(common.TestCase): @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" diff --git a/test/test_transformers.py b/test/test_transformers.py index 0e61a7e1feb..e332342e772 100644 --- a/test/test_transformers.py +++ b/test/test_transformers.py @@ -55,7 +55,6 @@ from torch.testing._internal.common_cuda import ( if not IS_FBCODE: from test_cpp_extensions_open_device_registration import ( - remove_build_path, generate_faked_module ) @@ -3849,7 +3848,7 @@ class TestAttnBias(NNTestCase): 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=[ diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py index 9866ab4b5a6..71805be2287 100644 --- a/torch/testing/_internal/common_utils.py +++ b/torch/testing/_internal/common_utils.py @@ -5359,3 +5359,17 @@ def check_leaked_tensors(limit=1, matched_type=torch.Tensor): 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)