Add Autoload Testcase for OpenReg

can-gaa-hou · can-gaa-hou · commit bdd7ca9a688d · 2025-08-07T14:21:06.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
@@ -923,9 +923,33 @@ def run_test_with_openreg(test_module, test_directory, options):
     install_dir, return_code = install_cpp_extensions(openreg_dir)
     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,7 +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_transformers_privateuse1": run_test_with_openreg,
+    "test_transformers_privateuse1": run_test_with_openreg_autoload,
 }
 
 
diff --git a/test/test_openreg.py b/test/test_openreg.py
@@ -10,7 +10,6 @@
 
 import numpy as np
 import psutil
-import torch_openreg  # noqa: F401
 
 import torch
 from torch.serialization import safe_globals
@@ -24,6 +23,17 @@
 )
 
 
+class TestAutoload(TestCase):
+    """Tests of autoloading the OpenReg backend"""
+
+    def test_autoload_disable(self):
+        # Test that the backend is loaded automatically
+        self.assertFalse(hasattr(torch, "openreg"))
+        import torch_openreg  # noqa: F401
+        self.assertTrue(torch.openreg.is_available())
+
+
+
 class TestPrivateUse1(TestCase):
     """Tests of third-parth device integration mechinasm based PrivateUse1"""
 
diff --git a/test/test_transformers_privateuse1.py b/test/test_transformers_privateuse1.py
@@ -4,8 +4,6 @@
 from collections import namedtuple
 from functools import partial
 
-import torch_openreg  # noqa: F401
-
 import torch
 from torch.nn.attention import SDPBackend
 from torch.testing._internal.common_nn import NNTestCase
