1# mypy: allow-untyped-defs 2import functools 3import importlib.util 4 5import torch 6 7 8def _check_module_exists(name: str) -> bool: 9 r"""Returns if a top-level module with :attr:`name` exists *without** 10 importing it. This is generally safer than try-catch block around a 11 `import X`. It avoids third party libraries breaking assumptions of some of 12 our tests, e.g., setting multiprocessing start method when imported 13 (see librosa/#747, torchvision/#544). 14 """ 15 try: 16 spec = importlib.util.find_spec(name) 17 return spec is not None 18 except ImportError: 19 return False 20 21 22@functools.lru_cache 23def dill_available(): 24 return ( 25 _check_module_exists("dill") 26 # dill fails to import under torchdeploy 27 and not torch._running_with_deploy() 28 ) 29 30 31@functools.lru_cache 32def import_dill(): 33 if not dill_available(): 34 return None 35 36 import dill 37 38 # XXX: By default, dill writes the Pickler dispatch table to inject its 39 # own logic there. This globally affects the behavior of the standard library 40 # pickler for any user who transitively depends on this module! 41 # Undo this extension to avoid altering the behavior of the pickler globally. 42 dill.extend(use_dill=False) 43 return dill 44