pytorch

1
# mypy: allow-untyped-defs
2
import functools
3
import importlib.util
4

5
import torch
6

7

8
def _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
23
def 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
32
def 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