| 12345678910111213141516171819202122232425262728293031323334353637383940414243 |
- # mypy: allow-untyped-defs
- import functools
- import importlib.util
- import torch
- def _check_module_exists(name: str) -> bool:
- r"""Returns if a top-level module with :attr:`name` exists *without**
- importing it. This is generally safer than try-catch block around a
- `import X`. It avoids third party libraries breaking assumptions of some of
- our tests, e.g., setting multiprocessing start method when imported
- (see librosa/#747, torchvision/#544).
- """
- try:
- spec = importlib.util.find_spec(name)
- return spec is not None
- except ImportError:
- return False
- @functools.lru_cache
- def dill_available():
- return (
- _check_module_exists("dill")
- # dill fails to import under torchdeploy
- and not torch._running_with_deploy()
- )
- @functools.lru_cache
- def import_dill():
- if not dill_available():
- return None
- import dill
- # XXX: By default, dill writes the Pickler dispatch table to inject its
- # own logic there. This globally affects the behavior of the standard library
- # pickler for any user who transitively depends on this module!
- # Undo this extension to avoid altering the behavior of the pickler globally.
- dill.extend(use_dill=False)
- return dill
|