Spaces:
Paused
Paused
File size: 5,051 Bytes
6b6d66f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import importlib
import logging
import os
import types
from pathlib import Path
import torch
_LG = logging.getLogger(__name__)
_LIB_DIR = Path(__file__).parent.parent / "lib"
class _LazyImporter(types.ModuleType):
"""Lazily import module/extension."""
def __init__(self, name, import_func):
super().__init__(name)
self.import_func = import_func
self.module = None
# Note:
# Python caches what was retrieved with `__getattr__`, so this method will not be
# called again for the same item.
def __getattr__(self, item):
self._import_once()
return getattr(self.module, item)
def __repr__(self):
if self.module is None:
return f"<module '{self.__module__}.{self.__class__.__name__}(\"{self.name}\")'>"
return repr(self.module)
def __dir__(self):
self._import_once()
return dir(self.module)
def _import_once(self):
if self.module is None:
self.module = self.import_func()
# Note:
# By attaching the module attributes to self,
# module attributes are directly accessible.
# This allows to avoid calling __getattr__ for every attribute access.
self.__dict__.update(self.module.__dict__)
def is_available(self):
try:
self._import_once()
except Exception:
return False
return True
def _get_lib_path(lib: str):
suffix = "pyd" if os.name == "nt" else "so"
path = _LIB_DIR / f"{lib}.{suffix}"
return path
def _load_lib(lib: str) -> bool:
"""Load extension module
Note:
In case `torio` is deployed with `pex` format, the library file
is not in a standard location.
In this case, we expect that `libtorio` is available somewhere
in the search path of dynamic loading mechanism, so that importing
`_torio` will have library loader find and load `libtorio`.
This is the reason why the function should not raising an error when the library
file is not found.
Returns:
bool:
True if the library file is found AND the library loaded without failure.
False if the library file is not found (like in the case where torio
is deployed with pex format, thus the shared library file is
in a non-standard location.).
If the library file is found but there is an issue loading the library,
(such as missing dependency) then this function raises the exception as-is.
Raises:
Exception:
If the library file is found, but there is an issue loading the library file,
(when underlying `ctype.DLL` throws an exception), this function will pass
the exception as-is, instead of catching it and returning bool.
The expected case is `OSError` thrown by `ctype.DLL` when a dynamic dependency
is not found.
This behavior was chosen because the expected failure case is not recoverable.
If a dependency is missing, then users have to install it.
"""
path = _get_lib_path(lib)
if not path.exists():
return False
torch.ops.load_library(path)
return True
_FFMPEG_VERS = ["6", "5", "4", ""]
def _find_versionsed_ffmpeg_extension(version: str):
ext = f"torio.lib._torio_ffmpeg{version}"
lib = f"libtorio_ffmpeg{version}"
if not importlib.util.find_spec(ext):
raise RuntimeError(f"FFmpeg{version} extension is not available.")
_load_lib(lib)
return importlib.import_module(ext)
def _find_ffmpeg_extension(ffmpeg_vers):
for ffmpeg_ver in ffmpeg_vers:
_LG.debug("Loading FFmpeg%s", ffmpeg_ver)
try:
ext = _find_versionsed_ffmpeg_extension(ffmpeg_ver)
_LG.debug("Successfully loaded FFmpeg%s", ffmpeg_ver)
return ext
except Exception:
_LG.debug("Failed to load FFmpeg%s extension.", ffmpeg_ver, exc_info=True)
continue
raise ImportError(
f"Failed to intialize FFmpeg extension. Tried versions: {ffmpeg_vers}. "
"Enable DEBUG logging to see more details about the error."
)
def _get_ffmpeg_versions():
ffmpeg_vers = _FFMPEG_VERS
# User override
if (ffmpeg_ver := os.environ.get("TORIO_USE_FFMPEG_VERSION")) is not None:
if ffmpeg_ver not in ffmpeg_vers:
raise ValueError(
f"The FFmpeg version '{ffmpeg_ver}' (read from TORIO_USE_FFMPEG_VERSION) "
f"is not one of supported values. Possible values are {ffmpeg_vers}"
)
ffmpeg_vers = [ffmpeg_ver]
return ffmpeg_vers
def _init_ffmpeg():
ffmpeg_vers = _get_ffmpeg_versions()
ext = _find_ffmpeg_extension(ffmpeg_vers)
ext.init()
if ext.get_log_level() > 8:
ext.set_log_level(8)
return ext
|