pytorch · jbschlosser · May 9, 2025 · May 9, 2025 · May 13, 2025 · May 19, 2025
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
@@ -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(), only_for="cuda")
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/test_cpp_extensions_aot.py b/test/test_cpp_extensions_aot.py
@@ -2,10 +2,8 @@
 
 import os
 import re
-import subprocess
 import unittest
 from itertools import repeat
-from pathlib import Path
 from typing import get_args, get_origin, Union
 
 import torch
@@ -183,38 +181,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):