diff --git a/test/test_extension.py b/test/test_extension.py new file mode 100644 index 00000000000..1be54e2310c --- /dev/null +++ b/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] diff --git a/torchvision/_meta_registrations.py b/torchvision/_meta_registrations.py index cb2602acc7e..29fc14e5fbd 100644 --- a/torchvision/_meta_registrations.py +++ b/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") diff --git a/torchvision/extension.py b/torchvision/extension.py index ddac86e70b6..5896efc5231 100644 --- a/torchvision/extension.py +++ b/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():