From 19904cf208a16d8ee2938bf09119c84d5aae7b2c Mon Sep 17 00:00:00 2001 From: Joel Schlosser Date: Fri, 9 May 2025 15:00:51 -0400 Subject: [PATCH] Make python_agnostic cpp extension tests standalone [ghstack-poisoned] --- .../test/test_python_agnostic.py | 65 +++++++++++++++++++ test/test_cpp_extensions_aot.py | 32 --------- 2 files changed, 65 insertions(+), 32 deletions(-) create mode 100644 test/cpp_extensions/python_agnostic_extension/test/test_python_agnostic.py diff --git a/test/cpp_extensions/python_agnostic_extension/test/test_python_agnostic.py b/test/cpp_extensions/python_agnostic_extension/test/test_python_agnostic.py new file mode 100644 index 00000000000000..74968a088fc7db --- /dev/null +++ b/test/cpp_extensions/python_agnostic_extension/test/test_python_agnostic.py @@ -0,0 +1,65 @@ +# Owner(s): ["module: cpp"] + +import os +import shutil +import subprocess +import sys +import unittest +from pathlib import Path + +from torch.testing._internal.common_device_type import ( + instantiate_device_type_tests, + onlyCUDA, +) +from torch.testing._internal.common_utils import IS_LINUX, run_tests, shell, TestCase + + +class TestPythonAgnostic(TestCase): + @classmethod + def setUpClass(cls): + # Wipe the dist dir if it exists + cls.extension_root = Path(__file__).parent.parent + cls.dist_dir = os.path.join(cls.extension_root, "dist") + if os.path.exists(cls.dist_dir): + shutil.rmtree(cls.dist_dir) + + # Build the wheel + wheel_cmd = [sys.executable, "setup.py", "bdist_wheel"] + return_code = shell(wheel_cmd, cwd=cls.extension_root, env=os.environ) + if return_code != 0: + raise RuntimeError("python_agnostic bdist_wheel failed to build") + + @onlyCUDA + @unittest.skipIf(not IS_LINUX, "test requires linux tools ldd and nm") + def test_extension_is_python_agnostic(self, device): + # For this test, run_test.py will call `python setup.py bdist_wheel` in the + # cpp_extensions/python_agnostic_extension folder, where the extension and + # setup calls specify py_limited_api to `True`. To approximate that the + # extension is indeed python agnostic, we test + # a. The extension wheel name contains "cp39-abi3", meaning the wheel + # should be runnable for any Python 3 version after and including 3.9 + # b. The produced shared library does not have libtorch_python.so as a + # dependency from the output of "ldd _C.so" + # c. The .so does not need any python related symbols. We approximate + # this by running "nm -u _C.so" and grepping that nothing starts with "Py" + + matches = list(Path(self.dist_dir).glob("*.whl")) + self.assertEqual(len(matches), 1, msg=str(matches)) + whl_file = matches[0] + self.assertRegex(str(whl_file), r".*python_agnostic-0\.0-cp39-abi3-.*\.whl") + + build_dir = os.path.join(self.extension_root, "build") + matches = list(Path(build_dir).glob("**/*.so")) + self.assertEqual(len(matches), 1, msg=str(matches)) + so_file = matches[0] + lddtree = subprocess.check_output(["ldd", so_file]).decode("utf-8") + self.assertFalse("torch_python" in lddtree) + + missing_symbols = subprocess.check_output(["nm", "-u", so_file]).decode("utf-8") + self.assertFalse("Py" in missing_symbols) + + +instantiate_device_type_tests(TestPythonAgnostic, globals()) + +if __name__ == "__main__": + run_tests() diff --git a/test/test_cpp_extensions_aot.py b/test/test_cpp_extensions_aot.py index ab7184797b4ef2..861ecf6c94986a 100644 --- a/test/test_cpp_extensions_aot.py +++ b/test/test_cpp_extensions_aot.py @@ -183,38 +183,6 @@ def test_cuda_dlink_libs(self): test = cuda_dlink.add(a, b) self.assertEqual(test, ref) - @unittest.skipIf(not TEST_CUDA, "python_agnostic is a CUDA extension + needs CUDA") - @unittest.skipIf(not common.IS_LINUX, "test requires linux tools ldd and nm") - def test_python_agnostic(self): - # For this test, run_test.py will call `python setup.py bdist_wheel` in the - # cpp_extensions/python_agnostic_extension folder, where the extension and - # setup calls specify py_limited_api to `True`. To approximate that the - # extension is indeed python agnostic, we test - # a. The extension wheel name contains "cp39-abi3", meaning the wheel - # should be runnable for any Python 3 version after and including 3.9 - # b. The produced shared library does not have libtorch_python.so as a - # dependency from the output of "ldd _C.so" - # c. The .so does not need any python related symbols. We approximate - # this by running "nm -u _C.so" and grepping that nothing starts with "Py" - - dist_root = os.path.join("cpp_extensions", "python_agnostic_extension", "dist") - matches = list(Path(dist_root).glob("*.whl")) - self.assertEqual(len(matches), 1, msg=str(matches)) - whl_file = matches[0] - self.assertRegex(str(whl_file), r".*python_agnostic-0\.0-cp39-abi3-.*\.whl") - - build_root = os.path.join( - "cpp_extensions", "python_agnostic_extension", "build" - ) - matches = list(Path(build_root).glob("**/*.so")) - self.assertEqual(len(matches), 1, msg=str(matches)) - so_file = matches[0] - lddtree = subprocess.check_output(["ldd", so_file]).decode("utf-8") - self.assertFalse("torch_python" in lddtree) - - missing_symbols = subprocess.check_output(["nm", "-u", so_file]).decode("utf-8") - self.assertFalse("Py" in missing_symbols) - @torch.testing._internal.common_utils.markDynamoStrictTest class TestPybindTypeCasters(common.TestCase):