Refactor the library loading mechanism (#2038)

Summary:
(This is a part of refactor series, followed up by https://github.com/pytorch/audio/issues/2039 and https://github.com/pytorch/audio/issues/2040.
The goal is to make it easy to add a new library artifact alongside with `libtorchudio`, as in https://github.com/pytorch/audio/pull/2048/commits/4ced990849e60f6d19e87ae22819b04d1726648e https://github.com/pytorch/audio/issues/2048 .)

We plan to add prototype/beta third party library integrations,
which could be unstable. (segfault, missing dynamic library dependencies etc...)

If we add such integrations into the existing libtorchaudio,
in the worst case, it will prevent users from just `import torchaudio`.

Instead, we would like to separate the prototype/beta integrations
into separate libraries, so that such issues would not impact all users but
users who attempt to use these prototytpe/beta features.

Say, a prototype feature `foo` is added in `torchaudio.prototype.foo`.
The following initialization procedure will achieve the above mechanism.

1. Place the library file `libtorchaudio_foo` in `torchaudio/lib`.
2. In `torchaudio.prototype.foo.__init__.py`, load the `libtorchaudio_foo`.

Note:
The approach will be slightly different for fbcode, because of how buck deploys
C++ libraries and standardized environment, but the code change here is still
applicable.

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

Reviewed By: carolineechen, nateanl

Differential Revision: D32682900

Pulled By: mthrok

fbshipit-source-id: 0f402a92a366fba8c2894a0fe01f47f8cdd51376

torchaudio/_extension.py

@@ -5,14 +5,17 @@ from pathlib import Path
 import torch
 from torchaudio._internal import module_utils as _mod_utils  # noqa: F401
 
+_LIB_DIR = Path(__file__).parent / 'lib'
 
-def _init_extension():
-    if not _mod_utils.is_module_available('torchaudio._torchaudio'):
-        warnings.warn('torchaudio C++ extension is not available.')
-        return
 
+def _get_lib_path(lib: str):
     suffix = 'pyd' if os.name == 'nt' else 'so'
-    path = Path(__file__).parent / 'lib' / f'libtorchaudio.{suffix}'
+    path = _LIB_DIR / f'{lib}.{suffix}'
+    return path
+
+
+def _load_lib(lib: str):
+    path = _get_lib_path(lib)
     # In case `torchaudio` is deployed with `pex` format, this file does not exist.
     # In this case, we expect that `libtorchaudio` is available somewhere
     # in the search path of dynamic loading mechanism, and importing `_torchaudio`,
@@ -20,7 +23,16 @@ def _init_extension():
     if path.exists():
         torch.ops.load_library(path)
         torch.classes.load_library(path)
+
+
+def _init_extension():
+    if not _mod_utils.is_module_available('torchaudio._torchaudio'):
+        warnings.warn('torchaudio C++ extension is not available.')
+        return
+
+    _load_lib('libtorchaudio')
     # This import is for initializing the methods registered via PyBind11
+    # This has to happen after the base library is loaded
     from torchaudio import _torchaudio  # noqa