Update I/O initialization (#2417)

Summary:
Attempt to load ffmpeg extension at the top level import

Preparation to use ffmpeg-based I/O as a fallback for sox_io backend.

Pull Request resolved: #2417

Reviewed By: carolineechen

Differential Revision: D36736989

Pulled By: mthrok

fbshipit-source-id: 0beb6f459313b5ea91597393ccb12571444c54d9

torchaudio/_extension.py

@@ -53,6 +53,31 @@ def _load_lib(lib: str) -> bool:
     return True
 
 
+_FFMPEG_INITIALIZED = False
+
+
+def _init_ffmpeg():
+    global _FFMPEG_INITIALIZED
+    if _FFMPEG_INITIALIZED:
+        return
+
+    if not torch.ops.torchaudio.is_ffmpeg_available():
+        raise RuntimeError(
+            "torchaudio is not compiled with FFmpeg integration. Please set USE_FFMPEG=1 when compiling torchaudio."
+        )
+
+    try:
+        _load_lib("libtorchaudio_ffmpeg")
+    except OSError as err:
+        raise ImportError("FFmpeg libraries are not found. Please install FFmpeg.") from err
+
+    import torchaudio._torchaudio_ffmpeg  # noqa
+
+    torch.ops.torchaudio.ffmpeg_init()
+
+    _FFMPEG_INITIALIZED = True
+
+
 def _init_extension():
     if not _mod_utils.is_module_available("torchaudio._torchaudio"):
         warnings.warn("torchaudio C++ extension is not available.")
@@ -63,5 +88,14 @@ def _init_extension():
     # This has to happen after the base library is loaded
     from torchaudio import _torchaudio  # noqa
 
+    # Because this part is executed as part of `import torchaudio`, we ignore the
+    # initialization failure.
+    # If the FFmpeg integration is not properly initialized, then detailed error
+    # will be raised when client code attempts to import the dedicated feature.
+    try:
+        _init_ffmpeg()
+    except Exception:
+        pass
+
 
 _init_extension()

torchaudio/csrc/CMakeLists.txt

@@ -108,6 +108,14 @@ if(OpenMP_CXX_FOUND)
     )
 endif()
 
+if(USE_FFMPEG)
+  list(
+    APPEND
+    LIBTORCHAUDIO_COMPILE_DEFINITIONS
+    USE_FFMPEG
+    )
+endif()
+
 #------------------------------------------------------------------------------#
 # END OF CUSTOMIZATION LOGICS
 #------------------------------------------------------------------------------#

torchaudio/csrc/utils.cpp

@@ -20,11 +20,22 @@ bool is_kaldi_available() {
 #endif
 }
 
+// It tells whether torchaudio was compiled with ffmpeg
+// not the runtime availability.
+bool is_ffmpeg_available() {
+#ifdef USE_FFMPEG
+  return true;
+#else
+  return false;
+#endif
+}
+
 } // namespace
 
 TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def("torchaudio::is_sox_available", &is_sox_available);
   m.def("torchaudio::is_kaldi_available", &is_kaldi_available);
+  m.def("torchaudio::is_ffmpeg_available", &is_ffmpeg_available);
 }
 
 } // namespace torchaudio

torchaudio/io/__init__.py

@@ -1,4 +1,5 @@
-_INITIALIZED = False
+import torchaudio
+
 _LAZILY_IMPORTED = [
     "StreamReader",
     "StreamReaderSourceStream",
@@ -8,32 +9,10 @@
 ]
 
 
-def _init_extension():
-    import torch
-    import torchaudio
-
-    try:
-        torchaudio._extension._load_lib("libtorchaudio_ffmpeg")
-        import torchaudio._torchaudio_ffmpeg
-    except OSError as err:
-        raise ImportError(
-            "Stream API requires FFmpeg libraries (libavformat and such). Please install FFmpeg 4."
-        ) from err
-    try:
-        torch.ops.torchaudio.ffmpeg_init()
-    except RuntimeError as err:
-        raise RuntimeError(
-            "Stream API requires FFmpeg binding. Please set USE_FFMPEG=1 when building from source."
-        ) from err
-
-    global _INITIALIZED
-    _INITIALIZED = True
-
-
 def __getattr__(name: str):
     if name in _LAZILY_IMPORTED:
-        if not _INITIALIZED:
-            _init_extension()
+
+        torchaudio._extension._init_ffmpeg()
 
         from . import _stream_reader