Improve error when torchvision C++ ops fail to load

test/test_extension.py

@@ -0,0 +1,63 @@
+import inspect
+
+import pytest
+import torchvision._meta_registrations as meta_registrations
+from torchvision import extension
+
+
+def _raise_missing_module(lib_name):
+    raise ImportError(f"Could not find module '{lib_name}'")
+
+
+def test_load_library_caches_error(monkeypatch):
+    monkeypatch.setattr(extension, "_EXTENSION_LOAD_ERROR", None)
+    monkeypatch.setattr(extension, "_get_extension_path", _raise_missing_module)
+
+    assert extension._load_library("_C_stable") is False
+    assert isinstance(extension._EXTENSION_LOAD_ERROR, ImportError)
+    assert "Could not find module '_C_stable'" in str(extension._EXTENSION_LOAD_ERROR)
+
+
+def test_c_stable_load_error_not_hidden_by_later_library(monkeypatch):
+    monkeypatch.setattr(extension, "_EXTENSION_LOAD_ERROR", None)
+    monkeypatch.setattr(extension, "_get_extension_path", _raise_missing_module)
+
+    assert extension._load_library("_C_stable") is False
+    assert extension._load_library("image_stable") is False
+    assert "Could not find module '_C_stable'" in str(extension._EXTENSION_LOAD_ERROR)
+
+
+def test_assert_has_ops_includes_underlying_error(monkeypatch):
+    load_error = ImportError("Could not find module '_C_stable' in /tmp/torchvision")
+    monkeypatch.setattr(extension, "_has_ops", lambda: False)
+    monkeypatch.setattr(extension, "_EXTENSION_LOAD_ERROR", load_error)
+
+    with pytest.raises(RuntimeError, match=r"Couldn't load custom C\+\+ ops") as exc_info:
+        extension._assert_has_ops()
+
+    message = str(exc_info.value)
+    assert "Underlying error: ImportError: Could not find module '_C_stable'" in message
+
+
+def test_assert_has_ops_without_cached_error(monkeypatch):
+    monkeypatch.setattr(extension, "_has_ops", lambda: False)
+    monkeypatch.setattr(extension, "_EXTENSION_LOAD_ERROR", None)
+
+    with pytest.raises(RuntimeError, match=r"Couldn't load custom C\+\+ ops") as exc_info:
+        extension._assert_has_ops()
+
+    assert "Underlying error:" not in str(exc_info.value)
+
+
+def test_register_fake_is_guarded_when_ops_are_missing():
+    # register_fake("torchvision::nms"|qnms) must sit inside an _has_ops() guard
+    # so importing an unbuilt source tree does not raise "operator does not exist".
+    source = inspect.getsource(meta_registrations)
+    nms_idx = source.index('@torch.library.register_fake("torchvision::nms")')
+    qnms_idx = source.index('@torch.library.register_fake("torchvision::qnms")')
+    guard = "if torchvision.extension._has_ops():"
+    nearest_guard = source.rfind(guard, 0, nms_idx)
+    assert nearest_guard != -1
+    assert nearest_guard < nms_idx < qnms_idx
+    # The nms/qnms decorators share that same guard block.
+    assert guard not in source[nms_idx:qnms_idx]

torchvision/_meta_registrations.py

@@ -160,32 +160,36 @@ def meta_ps_roi_pool_backward(
     return grad.new_empty((batch_size, channels, height, width))
 
 
-@torch.library.register_fake("torchvision::nms")
-def meta_nms(dets, scores, iou_threshold):
-    torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
-    torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
-    torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
-    torch._check(
-        dets.size(0) == scores.size(0),
-        lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
-    )
-    ctx = torch._custom_ops.get_ctx()
-    num_to_keep = ctx.create_unbacked_symint()
-    return dets.new_empty(num_to_keep, dtype=torch.long)
-
-
-@torch.library.register_fake("torchvision::qnms")
-def meta_qnms(dets, scores, iou_threshold):
-    torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
-    torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
-    torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
-    torch._check(
-        dets.size(0) == scores.size(0),
-        lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
-    )
-    ctx = torch._custom_ops.get_ctx()
-    num_to_keep = ctx.create_unbacked_symint()
-    return dets.new_empty(num_to_keep, dtype=torch.long)
+# register_fake requires the C++ operator to already exist. Skip when the
+# extension failed to load so `import torchvision` still succeeds; nms()/qnms()
+# then raise via _assert_has_ops() with the cached load error.
+if torchvision.extension._has_ops():
+
+    @torch.library.register_fake("torchvision::nms")
+    def meta_nms(dets, scores, iou_threshold):
+        torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
+        torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
+        torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
+        torch._check(
+            dets.size(0) == scores.size(0),
+            lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
+        )
+        ctx = torch._custom_ops.get_ctx()
+        num_to_keep = ctx.create_unbacked_symint()
+        return dets.new_empty(num_to_keep, dtype=torch.long)
+
+    @torch.library.register_fake("torchvision::qnms")
+    def meta_qnms(dets, scores, iou_threshold):
+        torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
+        torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
+        torch._check(scores.dim() == 1, lambda: f"scores should be a 1d tensor, got {scores.dim()}")
+        torch._check(
+            dets.size(0) == scores.size(0),
+            lambda: f"boxes and scores should have same number of elements in dimension 0, got {dets.size(0)} and {scores.size(0)}",
+        )
+        ctx = torch._custom_ops.get_ctx()
+        num_to_keep = ctx.create_unbacked_symint()
+        return dets.new_empty(num_to_keep, dtype=torch.long)
 
 
 @register_meta("qroi_align")

torchvision/extension.py

@@ -4,17 +4,27 @@
 
 from ._internally_replaced_utils import _get_extension_path
 
+# Last ImportError/OSError raised while loading a native extension. Surfaced by
+# _assert_has_ops() so users see the missing .so / ABI mismatch instead of only
+# a generic "couldn't load custom C++ ops" message.
+_EXTENSION_LOAD_ERROR = None
+
 
 def _load_library(lib_name):
     """Load a library, optionally warning on failure based on env variable.
 
     Returns True if the library was loaded successfully, False otherwise.
     """
+    global _EXTENSION_LOAD_ERROR
     try:
         lib_path = _get_extension_path(lib_name)
         torch.ops.load_library(lib_path)
         return True
     except (ImportError, OSError) as e:
+        # Prefer the C++ ops library error; later optional extensions (e.g.
+        # image_stable) must not hide the reason _has_ops() is False.
+        if _EXTENSION_LOAD_ERROR is None or lib_name == "_C_stable":
+            _EXTENSION_LOAD_ERROR = e
         if os.environ.get("TORCHVISION_WARN_WHEN_EXTENSION_LOADING_FAILS"):
             import warnings
 
@@ -34,7 +44,7 @@ def _has_ops():  # noqa: F811
 
 def _assert_has_ops():
     if not _has_ops():
-        raise RuntimeError(
+        msg = (
             "Couldn't load custom C++ ops. This can happen if your PyTorch and "
             "torchvision versions are incompatible, or if you had errors while compiling "
             "torchvision from source. For further information on the compatible versions, check "
@@ -44,6 +54,9 @@ def _assert_has_ops():
             "please reinstall torchvision so that it matches your PyTorch install. "
             "Set TORCHVISION_WARN_WHEN_EXTENSION_LOADING_FAILS=1 and retry to get more details."
         )
+        if _EXTENSION_LOAD_ERROR is not None:
+            msg += f" Underlying error: {type(_EXTENSION_LOAD_ERROR).__name__}: {_EXTENSION_LOAD_ERROR}"
+        raise RuntimeError(msg)
 
 
 def _check_cuda_version():