Add warning when importing caffe2 on build without BUILD_CAFFE2=1

Confusing backtraces are issued to users when they run Caffe2 scripts (or tests) on PyTorch builds without Caffe2 enabled through `BUILD_CAFFE2=1` This PR adds warnings (in more than one place) to return a friendly message for the user, helping them to overcome the problem by themselves Pull Request resolved: pytorch#73770 Approved by: https://github.com/BowenBao, https://github.com/malfet, https://github.com/garymm
bowu1999 · Apr 21, 2022 · e071340 · e071340
1 parent bf730e5
commit e071340
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 4 deletions.
diff --git a/caffe2/__init__.py b/caffe2/__init__.py
@@ -0,0 +1,6 @@
+import warnings
+from torch.onnx import _CAFFE2_ATEN_FALLBACK
+
+if not _CAFFE2_ATEN_FALLBACK:
+    warnings.warn("Caffe2 support is not fully enabled in this PyTorch build. "
+                  "Please enable Caffe2 by building PyTorch from source with `BUILD_CAFFE2=1` flag.")
diff --git a/caffe2/proto/__init__.py b/caffe2/proto/__init__.py
@@ -1,3 +1,6 @@
+import warnings
+
+
 # NOTE: we have to import python protobuf here **before** we load cpp extension.
 # Otherwise it breaks under certain build conditions if cpp implementation of
 # protobuf is used. Presumably there's some registry in protobuf library and
@@ -8,7 +11,13 @@
 #    expected caffe2.NetDef got caffe2.NetDef."
 #
 # This has to be done for all python targets, so listing them here
-from caffe2.proto import caffe2_pb2, metanet_pb2, torch_pb2
+try:
+    from caffe2.proto import caffe2_pb2, metanet_pb2, torch_pb2
+except ImportError:
+    warnings.warn('Caffe2 support is not enabled in this PyTorch build. '
+                  'Please enable Caffe2 by building PyTorch from source with `BUILD_CAFFE2=1` flag.')
+    raise
+
 try:
     from caffe2.caffe2.fb.session.proto import session_pb2
 except ImportError:

diff --git a/caffe2/python/__init__.py b/caffe2/python/__init__.py
@@ -1,7 +1,15 @@
-
-from caffe2.proto import caffe2_pb2
 import os
 import sys
+import warnings
+
+
+try:
+    from caffe2.proto import caffe2_pb2
+except ImportError:
+    warnings.warn('Caffe2 support is not enabled in this PyTorch build. '
+                  'Please enable Caffe2 by building PyTorch from source with `BUILD_CAFFE2=1` flag.')
+    raise
+
 # TODO: refactor & remove the following alias
 caffe2_pb2.CPU = caffe2_pb2.PROTO_CPU
 caffe2_pb2.CUDA = caffe2_pb2.PROTO_CUDA

diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
@@ -766,7 +766,7 @@ def _check_module_exists(name: str) -> bool:
 
 TEST_LIBROSA = _check_module_exists('librosa')
 
-BUILD_WITH_CAFFE2 = _check_module_exists("caffe2.python.caffe2_pybind11_state")
+BUILD_WITH_CAFFE2 = torch.onnx._CAFFE2_ATEN_FALLBACK
 
 # Python 2.7 doesn't have spawn
 NO_MULTIPROCESSING_SPAWN = os.environ.get('NO_MULTIPROCESSING_SPAWN', '0') == '1'