mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
8db9dfa2d7
See #127836 for details. Pull Request resolved: https://github.com/pytorch/pytorch/pull/127846 Approved by: https://github.com/ezyang ghstack dependencies: #127842, #127843, #127844, #127845
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
# 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
|