_import_utils.py 1.2 KB

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