From 82f551ded9b12a351f181a93ccd30134b1d79826 Mon Sep 17 00:00:00 2001 From: Moto Hira Date: Mon, 30 Aug 2021 05:05:31 -0700 Subject: [PATCH] Simplify the initialization process Differential Revision: D30633316 fbshipit-source-id: 91102728e5a74f93a5fb7b85bda33b038307f0b0 --- torchaudio/__init__.py | 36 ++++++++++++++++++++++++++++++- torchaudio/extension/__init__.py | 12 ----------- torchaudio/extension/extension.py | 23 -------------------- 3 files changed, 35 insertions(+), 36 deletions(-) delete mode 100644 torchaudio/extension/__init__.py delete mode 100644 torchaudio/extension/extension.py diff --git a/torchaudio/__init__.py b/torchaudio/__init__.py index 76aa052400..bfb5034780 100644 --- a/torchaudio/__init__.py +++ b/torchaudio/__init__.py @@ -1,5 +1,39 @@ -from . import extension # noqa: F401 from torchaudio._internal import module_utils as _mod_utils # noqa: F401 + +if _mod_utils.is_module_available('torchaudio._torchaudio'): + # Note this import has two purposes + # 1. Make _torchaudio accessible by the other modules (regular import) + # 2. Register torchaudio's custom ops bound via TorchScript + # + # For 2, normally function calls `torch.ops.load_library` and `torch.classes.load_library` + # are used. However, in our cases, this is inconvenient and unnecessary. + # + # - Why inconvenient? + # When torchaudio is deployed with `pex` format, all the files are deployed as a single zip + # file, and the extension module is not present as a file with full path. Therefore it is not + # possible to pass the path to library to `torch.[ops|classes].load_library` functions. + # + # - Why unnecessary? + # When torchaudio extension module (C++ module) is available, it is assumed that + # the extension contains both TorchScript-based binding and PyBind11-based binding.* + # Under this assumption, simply performing `from torchaudio import _torchaudio` will load the + # library which contains TorchScript-based binding as well, and the functions/classes bound + # via TorchScript become accessible under `torch.ops` and `torch.classes`. + # + # *Note that this holds true even when these two bindings are split into two library files and + # the library that contains PyBind11-based binding (`_torchaudio.so` in the following diagram) + # depends on the other one (`libtorchaudio.so`), because when the process tries to load + # `_torchaudio.so` it detects undefined symbols from `libtorchaudio.so` and will automatically + # loads `libtorchaudio.so`. (given that the library is found in a search path) + # + # [libtorchaudio.so] <- [_torchaudio.so] + # + # + from torchaudio import _torchaudio # noqa +else: + import warnings + warnings.warn('torchaudio C++ extension is not available.') + from torchaudio import ( compliance, datasets, diff --git a/torchaudio/extension/__init__.py b/torchaudio/extension/__init__.py deleted file mode 100644 index 1a600f45c5..0000000000 --- a/torchaudio/extension/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from .extension import ( - _init_extension, -) - -try: - from . import fb # noqa -except Exception: - pass - -_init_extension() - -del _init_extension diff --git a/torchaudio/extension/extension.py b/torchaudio/extension/extension.py deleted file mode 100644 index d97d932f2b..0000000000 --- a/torchaudio/extension/extension.py +++ /dev/null @@ -1,23 +0,0 @@ -import warnings - -import torch -from torchaudio._internal import module_utils as _mod_utils - - -def _init_extension(): - if _mod_utils.is_module_available('torchaudio._torchaudio'): - # Note this import has two purposes - # 1. to extract the path of the extension module so that - # we can initialize the script module with the path. - # 2. so that torchaudio._torchaudio is accessible in other modules. - # Look at sox_io_backend which uses `torchaudio._torchaudio.XXX`, - # assuming that the module `_torchaudio` is accessible. - import torchaudio._torchaudio - _init_script_module(torchaudio._torchaudio.__file__) - else: - warnings.warn('torchaudio C++ extension is not available.') - - -def _init_script_module(path): - torch.classes.load_library(path) - torch.ops.load_library(path)