Refactor extension modules initialization (#2968)

Summary:
* Refactor _extension module so that
  * the implementation of initialization logic and its execution are separated.
    * logic goes to `_extension.utils`
    * the execution is at `_extension.__init__`
    * global variables are defined and modified in `__init__`.
* Replace `is_sox_available()` with `_extension._SOX_INITIALIZED`
* Replace `is_kaldi_available()` with `_extension._IS_KALDI_AVAILABLE`
* Move `requies_sox()` and `requires_kaldi()` to break the circular dependency among `_extension` and `_internal.module_utils`.
* Merge the sox-related initialization logic in `_extension.utils` module.

Pull Request resolved: https://github.com/pytorch/audio/pull/2968

Reviewed By: hwangjeff

Differential Revision: D42387251

Pulled By: mthrok

fbshipit-source-id: 0c3245dfab53f9bc1b8a83ec2622eb88ec96673f

test/torchaudio_unittest/common_utils/case_utils.py

@@ -11,7 +11,7 @@ from itertools import zip_longest
 import torch
 import torchaudio
 from torch.testing._internal.common_utils import TestCase as PytorchTestCase
-from torchaudio._internal.module_utils import is_kaldi_available, is_module_available, is_sox_available
+from torchaudio._internal.module_utils import is_module_available
 
 from .backend_utils import set_audio_backend
 
@@ -209,12 +209,12 @@ skipIfNoCuda = _skipIf(
     key="NO_CUDA",
 )
 skipIfNoSox = _skipIf(
-    not is_sox_available(),
+    not torchaudio._extension._SOX_INITIALIZED,
     reason="Sox features are not available.",
     key="NO_SOX",
 )
 skipIfNoKaldi = _skipIf(
-    not is_kaldi_available(),
+    not torchaudio._extension._IS_KALDI_AVAILABLE,
     reason="Kaldi features are not available.",
     key="NO_KALDI",
 )

torchaudio/_extension/__init__.py

+import logging
+import os
+import sys
+
+from torchaudio._internal.module_utils import fail_with_message, is_module_available, no_op
+
+from .utils import _check_cuda_version, _init_dll_path, _init_ffmpeg, _init_sox, _load_lib  # noqa
+
+_LG = logging.getLogger(__name__)
+
+
+# Note:
+# `_check_cuda_version` is not meant to be used by regular users.
+# Builder uses it for debugging purpose, so we export it.
+# https://github.com/pytorch/builder/blob/e2e4542b8eb0bdf491214451a1a4128bd606cce2/test/smoke_test/smoke_test.py#L80
+__all__ = [
+    "fail_if_no_kaldi",
+    "fail_if_no_sox",
+    "_check_cuda_version",
+    "_IS_TORCHAUDIO_EXT_AVAILABLE",
+    "_IS_KALDI_AVAILABLE",
+    "_SOX_INITIALIZED",
+    "_FFMPEG_INITIALIZED",
+]
+
+
+if os.name == "nt" and (3, 8) <= sys.version_info < (3, 9):
+    _init_dll_path()
+
+
+# When the extension module is built, we initialize it.
+# In case of an error, we do not catch the failure as it suggests there is something
+# wrong with the installation.
+_IS_TORCHAUDIO_EXT_AVAILABLE = is_module_available("torchaudio.lib._torchaudio")
+# Kaldi features are implemented in _torchaudio extension, but it can be individually
+# turned on/off at build time. Available means that _torchaudio is loaded properly, and
+# Kaldi features are found there.
+_IS_KALDI_AVAILABLE = False
+if _IS_TORCHAUDIO_EXT_AVAILABLE:
+    _load_lib("libtorchaudio")
+
+    import torchaudio.lib._torchaudio  # noqa
+
+    _check_cuda_version()
+    _IS_KALDI_AVAILABLE = torchaudio.lib._torchaudio.is_kaldi_available()
+
+
+# Similar to libtorchaudio, sox-related features should be importable when present.
+#
+# Note: This will be change in the future when sox is dynamically linked.
+# At that point, this initialization should handle the case where
+# sox integration is built but libsox is not found.
+_SOX_INITIALIZED = False
+if is_module_available("torchaudio.lib._torchaudio_sox"):
+    _init_sox()
+    _SOX_INITIALIZED = True
+
+
+# Initialize FFmpeg-related features
+_FFMPEG_INITIALIZED = False
+if is_module_available("torchaudio.lib._torchaudio_ffmpeg"):
+    try:
+        _init_ffmpeg()
+        _FFMPEG_INITIALIZED = True
+    except Exception:
+        # The initialization of FFmpeg extension will fail if supported FFmpeg
+        # libraries are not found in the system.
+        # Since the rest of the torchaudio works without it, we do not report the
+        # error here.
+        # The error will be raised when user code attempts to use these features.
+        _LG.debug("Failed to initialize ffmpeg bindings", exc_info=True)
+
+
+fail_if_no_kaldi = (
+    no_op
+    if _IS_KALDI_AVAILABLE
+    else fail_with_message(
+        "requires kaldi extension, but TorchAudio is not compiled with it. Please build TorchAudio with kaldi support."
+    )
+)
+fail_if_no_sox = (
+    no_op
+    if _SOX_INITIALIZED
+    else fail_with_message(
+        "requires sox extension, but TorchAudio is not compiled with it. Please build TorchAudio with libsox support."
+    )
+)

torchaudio/_extension.py → torchaudio/_extension/utils.py

-import logging
 import os
-import sys
 from pathlib import Path
 
 import torch
 
+import torchaudio
 from torchaudio._internal.module_utils import is_module_available
 
-_LIB_DIR = Path(__file__).parent / "lib"
-
-_LG = logging.getLogger(__name__)
+_LIB_DIR = Path(__file__).parent.parent / "lib"
 
 
 def _get_lib_path(lib: str):
@@ -57,15 +54,19 @@ def _load_lib(lib: str) -> bool:
     return True
 
 
-_IS_TORCHAUDIO_EXT_AVAILABLE = is_module_available("torchaudio.lib._torchaudio")
-_FFMPEG_INITIALIZED = False
+def _init_sox():
+    _load_lib("libtorchaudio_sox")
+    import torchaudio.lib._torchaudio_sox  # noqa
 
+    torch.ops.torchaudio.sox_utils_set_verbosity(0)
 
-def _init_ffmpeg():
-    global _FFMPEG_INITIALIZED
-    if _FFMPEG_INITIALIZED:
-        return
+    import atexit
+
+    torch.ops.torchaudio.sox_effects_initialize_sox_effects()
+    atexit.register(torch.ops.torchaudio.sox_effects_shutdown_sox_effects)
 
+
+def _init_ffmpeg():
     if not is_module_available("torchaudio.lib._torchaudio_ffmpeg"):
         raise RuntimeError(
             "torchaudio is not compiled with FFmpeg integration. Please set USE_FFMPEG=1 when compiling torchaudio."
@@ -82,55 +83,23 @@ def _init_ffmpeg():
     if torch.ops.torchaudio.ffmpeg_get_log_level() > 8:
         torch.ops.torchaudio.ffmpeg_set_log_level(8)
 
-    _FFMPEG_INITIALIZED = True
 
-
-def _init_extension():
+def _init_dll_path():
     # On Windows Python-3.8+ has `os.add_dll_directory` call,
     # which is called to configure dll search path.
     # To find cuda related dlls we need to make sure the
     # conda environment/bin path is configured Please take a look:
     # https://stackoverflow.com/questions/59330863/cant-import-dll-module-in-python
     # Please note: if some path can't be added using add_dll_directory we simply ignore this path
-    if os.name == "nt" and sys.version_info >= (3, 8) and sys.version_info < (3, 9):
-        env_path = os.environ["PATH"]
-        path_arr = env_path.split(";")
-        for path in path_arr:
-            if os.path.exists(path):
-                try:
-                    os.add_dll_directory(path)
-                except Exception:
-                    pass
-
-    if _IS_TORCHAUDIO_EXT_AVAILABLE:
-        try:
-            _load_lib("libtorchaudio")
-            import torchaudio.lib._torchaudio  # noqa
-        except Exception:
-            _LG.debug("Failed to initialize libtorchaudio", exc_info=True)
-
-    if is_module_available("torchaudio.lib._torchaudio_sox"):
-        try:
-            _load_lib("libtorchaudio_sox")
-            import torchaudio.lib._torchaudio_sox  # noqa
-        except Exception:
-            _LG.debug("Failed to initialize libsox bindings", exc_info=True)
-
-    # If the FFmpeg integration is not properly initialized, then detailed error
-    # will be raised when client code attempts to import the dedicated feature.
-    if is_module_available("torchaudio.lib._torchaudio_ffmpeg"):
-        try:
-            _init_ffmpeg()
-        except Exception:
-            _LG.debug("Failed to initialize ffmpeg bindings", exc_info=True)
+    for path in os.environ.get("PATH", "").split(";"):
+        if os.path.exists(path):
+            try:
+                os.add_dll_directory(path)
+            except Exception:
+                pass
 
 
 def _check_cuda_version():
-    if not _IS_TORCHAUDIO_EXT_AVAILABLE:
-        return None
-
-    import torchaudio.lib._torchaudio
-
     version = torchaudio.lib._torchaudio.cuda_version()
     if version is not None and torch.version.cuda is not None:
         version_str = str(version)
@@ -144,7 +113,3 @@ def _check_cuda_version():
                 "Please install the TorchAudio version that matches your PyTorch version."
             )
     return version
-
-
-_init_extension()
-_check_cuda_version()

torchaudio/_internal/module_utils.py

@@ -64,87 +64,19 @@ def deprecated(direction: str, version: Optional[str] = None):
     return decorator
 
 
-def is_kaldi_available():
-    try:
-        import torchaudio.lib._torchaudio
+def fail_with_message(message):
+    """Generate decorator to give users message about missing TorchAudio extension."""
 
-        return torchaudio.lib._torchaudio.is_kaldi_available()
-    except Exception:
-        return False
-
-
-def requires_kaldi():
-    if is_kaldi_available():
-
-        def decorator(func):
-            return func
-
-    else:
-
-        def decorator(func):
-            @wraps(func)
-            def wrapped(*args, **kwargs):
-                raise RuntimeError(f"{func.__module__}.{func.__name__} requires kaldi")
-
-            return wrapped
-
-    return decorator
-
-
-def _check_soundfile_importable():
-    if not is_module_available("soundfile"):
-        return False
-    try:
-        import soundfile  # noqa: F401
-
-        return True
-    except Exception:
-        warnings.warn("Failed to import soundfile. 'soundfile' backend is not available.")
-        return False
-
-
-_is_soundfile_importable = _check_soundfile_importable()
-
-
-def is_soundfile_available():
-    return _is_soundfile_importable
-
-
-def requires_soundfile():
-    if is_soundfile_available():
-
-        def decorator(func):
-            return func
-
-    else:
-
-        def decorator(func):
-            @wraps(func)
-            def wrapped(*args, **kwargs):
-                raise RuntimeError(f"{func.__module__}.{func.__name__} requires soundfile")
+    def decorator(func):
+        @wraps(func)
+        def wrapped(*args, **kwargs):
+            raise RuntimeError(f"{func.__module__}.{func.__name__} {message}")
 
-            return wrapped
+        return wrapped
 
     return decorator
 
 
-def is_sox_available():
-    return is_module_available("torchaudio.lib._torchaudio_sox")
-
-
-def requires_sox():
-    if is_sox_available():
-
-        def decorator(func):
-            return func
-
-    else:
-
-        def decorator(func):
-            @wraps(func)
-            def wrapped(*args, **kwargs):
-                raise RuntimeError(f"{func.__module__}.{func.__name__} requires sox")
-
-            return wrapped
-
-    return decorator
+def no_op(func):
+    """Op-op decorator. Used in place of fail_with_message when a functionality that requires extension works fine."""
+    return func

torchaudio/backend/soundfile_backend.py

@@ -8,8 +8,21 @@ from torchaudio._internal import module_utils as _mod_utils
 from .common import AudioMetaData
 
 
-if _mod_utils.is_soundfile_available():
-    import soundfile
+# TODO: import soundfile only when it is used.
+if _mod_utils.is_module_available("soundfile"):
+    try:
+        import soundfile
+
+        _requires_soundfile = _mod_utils.no_op
+    except Exception:
+        _requires_soundfile = _mod_utils.fail_with_message(
+            "requires soundfile, but we failed to import it. Please check the installation of soundfile."
+        )
+else:
+    _requires_soundfile = _mod_utils.fail_with_message(
+        "requires soundfile, but it is not installed. Please install soundfile."
+    )
+
 
 # Mapping from soundfile subtype to number of bits per sample.
 # This is mostly heuristical and the value is set to 0 when it is irrelevant
@@ -81,7 +94,7 @@ def _get_encoding(format: str, subtype: str):
     return _SUBTYPE_TO_ENCODING.get(subtype, "UNKNOWN")
 
 
-@_mod_utils.requires_soundfile()
+@_requires_soundfile
 def info(filepath: str, format: Optional[str] = None) -> AudioMetaData:
     """Get signal information of an audio file.
 
@@ -120,7 +133,7 @@ _SUBTYPE2DTYPE = {
 }
 
 
-@_mod_utils.requires_soundfile()
+@_requires_soundfile
 def load(
     filepath: str,
     frame_offset: int = 0,
@@ -299,7 +312,7 @@ def _get_subtype(dtype: torch.dtype, format: str, encoding: str, bits_per_sample
     raise ValueError(f"Unsupported format: {format}")
 
 
-@_mod_utils.requires_soundfile()
+@_requires_soundfile
 def save(
     filepath: str,
     src: torch.Tensor,

torchaudio/backend/sox_io_backend.py

@@ -3,7 +3,6 @@ from typing import Optional, Tuple
 
 import torch
 import torchaudio
-from torchaudio._internal import module_utils as _mod_utils
 from torchaudio.utils.sox_utils import get_buffer_size
 
 from .common import AudioMetaData
@@ -48,7 +47,7 @@ else:
     _fallback_load_fileobj = _fail_load_fileobj
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def info(
     filepath: str,
     format: Optional[str] = None,
@@ -106,7 +105,7 @@ def info(
     return _fallback_info(filepath, format)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def load(
     filepath: str,
     frame_offset: int = 0,
@@ -246,7 +245,7 @@ def load(
     return _fallback_load(filepath, frame_offset, num_frames, normalize, channels_first, format)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def save(
     filepath: str,
     src: torch.Tensor,

torchaudio/backend/utils.py

@@ -23,7 +23,7 @@ def list_audio_backends() -> List[str]:
     backends = []
     if _mod_utils.is_module_available("soundfile"):
         backends.append("soundfile")
-    if _mod_utils.is_sox_available():
+    if torchaudio._extension._SOX_INITIALIZED:
         backends.append("sox_io")
     return backends
 

torchaudio/functional/functional.py

@@ -9,7 +9,6 @@ from typing import List, Optional, Tuple, Union
 import torch
 import torchaudio
 from torch import Tensor
-from torchaudio._internal import module_utils as _mod_utils
 
 from .filtering import highpass_biquad, treble_biquad
 
@@ -1265,7 +1264,7 @@ def spectral_centroid(
     return (freqs * specgram).sum(dim=freq_dim) / specgram.sum(dim=freq_dim)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def apply_codec(
     waveform: Tensor,
     sample_rate: int,
@@ -1309,7 +1308,7 @@ def apply_codec(
     return augmented
 
 
-@_mod_utils.requires_kaldi()
+@torchaudio._extension.fail_if_no_kaldi
 def compute_kaldi_pitch(
     waveform: torch.Tensor,
     sample_rate: float,

torchaudio/io/__init__.py

@@ -18,7 +18,8 @@ _LAZILY_IMPORTED = _STREAM_READER + _STREAM_WRITER
 
 def __getattr__(name: str):
     if name in _LAZILY_IMPORTED:
-        torchaudio._extension._init_ffmpeg()
+        if not torchaudio._extension._FFMPEG_INITIALIZED:
+            torchaudio._extension._init_ffmpeg()
 
         if name in _STREAM_READER:
             from . import _stream_reader

torchaudio/sox_effects/__init__.py

-from torchaudio._internal import module_utils as _mod_utils
-
 from .sox_effects import apply_effects_file, apply_effects_tensor, effect_names, init_sox_effects, shutdown_sox_effects
 
 
-if _mod_utils.is_sox_available():
-    import atexit
-
-    init_sox_effects()
-    atexit.register(shutdown_sox_effects)
-
 __all__ = [
     "init_sox_effects",
     "shutdown_sox_effects",

torchaudio/sox_effects/sox_effects.py

@@ -3,11 +3,9 @@ from typing import List, Optional, Tuple
 
 import torch
 import torchaudio
-from torchaudio._internal import module_utils as _mod_utils
 from torchaudio.utils.sox_utils import list_effects
 
 
-@_mod_utils.requires_sox()
 def init_sox_effects():
     """Initialize resources required to use sox effects.
 
@@ -19,10 +17,9 @@ def init_sox_effects():
     Once :func:`shutdown_sox_effects` is called, you can no longer use SoX effects and initializing
     again will result in error.
     """
-    torch.ops.torchaudio.sox_effects_initialize_sox_effects()
+    pass
 
 
-@_mod_utils.requires_sox()
 def shutdown_sox_effects():
     """Clean up resources required to use sox effects.
 
@@ -33,10 +30,10 @@ def shutdown_sox_effects():
     Once :py:func:`shutdown_sox_effects` is called, you can no longer use SoX effects and
     initializing again will result in error.
     """
-    torch.ops.torchaudio.sox_effects_shutdown_sox_effects()
+    pass
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def effect_names() -> List[str]:
     """Gets list of valid sox effect names
 
@@ -50,7 +47,7 @@ def effect_names() -> List[str]:
     return list(list_effects().keys())
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def apply_effects_tensor(
     tensor: torch.Tensor,
     sample_rate: int,
@@ -155,7 +152,7 @@ def apply_effects_tensor(
     return torch.ops.torchaudio.sox_effects_apply_effects_tensor(tensor, sample_rate, effects, channels_first)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def apply_effects_file(
     path: str,
     effects: List[List[str]],

torchaudio/utils/__init__.py

-from torchaudio._internal import module_utils as _mod_utils
-
 from . import ffmpeg_utils, sox_utils
 from .download import download_asset
 
-if _mod_utils.is_sox_available():
-    sox_utils.set_verbosity(0)
 
 __all__ = [
     "download_asset",

torchaudio/utils/sox_utils.py

@@ -5,10 +5,10 @@
 from typing import Dict, List
 
 import torch
-from torchaudio._internal import module_utils as _mod_utils
+import torchaudio
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def set_seed(seed: int):
     """Set libsox's PRNG
 
@@ -21,7 +21,7 @@ def set_seed(seed: int):
     torch.ops.torchaudio.sox_utils_set_seed(seed)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def set_verbosity(verbosity: int):
     """Set libsox's verbosity
 
@@ -39,7 +39,7 @@ def set_verbosity(verbosity: int):
     torch.ops.torchaudio.sox_utils_set_verbosity(verbosity)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def set_buffer_size(buffer_size: int):
     """Set buffer size for sox effect chain
 
@@ -52,7 +52,7 @@ def set_buffer_size(buffer_size: int):
     torch.ops.torchaudio.sox_utils_set_buffer_size(buffer_size)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def set_use_threads(use_threads: bool):
     """Set multithread option for sox effect chain
 
@@ -66,7 +66,7 @@ def set_use_threads(use_threads: bool):
     torch.ops.torchaudio.sox_utils_set_use_threads(use_threads)
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def list_effects() -> Dict[str, str]:
     """List the available sox effect names
 
@@ -76,7 +76,7 @@ def list_effects() -> Dict[str, str]:
     return dict(torch.ops.torchaudio.sox_utils_list_effects())
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def list_read_formats() -> List[str]:
     """List the supported audio formats for read
 
@@ -86,7 +86,7 @@ def list_read_formats() -> List[str]:
     return torch.ops.torchaudio.sox_utils_list_read_formats()
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def list_write_formats() -> List[str]:
     """List the supported audio formats for write
 
@@ -96,7 +96,7 @@ def list_write_formats() -> List[str]:
     return torch.ops.torchaudio.sox_utils_list_write_formats()
 
 
-@_mod_utils.requires_sox()
+@torchaudio._extension.fail_if_no_sox
 def get_buffer_size() -> int:
     """Get buffer size for sox effect chain