Add Autoload Testcase for OpenReg

can-gaa-hou · can-gaa-hou · commit c6b901b29873 · 2025-08-08T01:51:59.000Z
diff --git a/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py b/test/cpp_extensions/open_registration_extension/torch_openreg/setup.py
@@ -14,8 +14,11 @@
 
 
 def get_pytorch_dir():
+    os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
     import torch
 
+    os.environ.pop("TORCH_DEVICE_BACKEND_AUTOLOAD")
+
     return os.path.dirname(os.path.realpath(torch.__file__))
 
 
diff --git a/test/run_test.py b/test/run_test.py
@@ -924,8 +924,32 @@ def run_test_with_openreg(test_module, test_directory, options):
     if return_code != 0:
         return return_code
 
+    os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0"
+
+    with extend_python_path([install_dir]):
+        return run_test(test_module, test_directory, options)
+    
+    os.environ.pop("TORCH_DEVICE_BACKEND_AUTOLOAD")
+
+
+def run_test_with_openreg_autoload(test_module, test_directory, options):
+    # TODO(FFFrog): Will remove this later when windows/macos are supported.
+    if not IS_LINUX:
+        return 0
+
+    openreg_dir = os.path.join(
+        test_directory, "cpp_extensions", "open_registration_extension", "torch_openreg"
+    )
+    install_dir, return_code = install_cpp_extensions(openreg_dir)
+    if return_code != 0:
+        return return_code
+
+    os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "1"
+
     with extend_python_path([install_dir]):
         return run_test(test_module, test_directory, options)
+    
+    os.environ.pop("TORCH_DEVICE_BACKEND_AUTOLOAD")
 
 
 def test_distributed(test_module, test_directory, options):
@@ -1256,6 +1280,7 @@ def run_ci_sanity_check(test: ShardedTest, test_directory, options):
     "test_autoload_disable": test_autoload_disable,
     "test_cpp_extensions_open_device_registration": run_test_with_openreg,
     "test_openreg": run_test_with_openreg,
+    "test_openreg_autoload": run_test_with_openreg_autoload,
     "test_transformers_privateuse1": run_test_with_openreg,
 }
 
diff --git a/test/test_openreg_autoload.py b/test/test_openreg_autoload.py
@@ -0,0 +1,14 @@
+import torch
+from torch.testing._internal.common_utils import TestCase, run_tests
+
+
+class TestAutoload(TestCase):
+    """Tests of autoloading the OpenReg backend"""
+
+    def test_autoload_enable(self):
+        # Test that the backend is loaded automatically
+        self.assertTrue(torch.openreg.is_available())
+
+
+if __name__ == "__main__":
+    run_tests()
