pytorch

1
import functools
2
import importlib.util
3

4
import torch
5

6

7
def _check_module_exists(name: str) -> bool:
8
    r"""Returns if a top-level module with :attr:`name` exists *without**
9
    importing it. This is generally safer than try-catch block around a
10
    `import X`. It avoids third party libraries breaking assumptions of some of
11
    our tests, e.g., setting multiprocessing start method when imported
12
    (see librosa/#747, torchvision/#544).
13
    """
14
    try:
15
        spec = importlib.util.find_spec(name)
16
        return spec is not None
17
    except ImportError:
18
        return False
19

20

21
@functools.lru_cache
22
def dill_available():
23
    return (
24
        _check_module_exists("dill")
25
        # dill fails to import under torchdeploy
26
        and not torch._running_with_deploy()
27
    )
28

29

30
@functools.lru_cache
31
def import_dill():
32
    if not dill_available():
33
        return None
34

35
    import dill
36

37
    # XXX: By default, dill writes the Pickler dispatch table to inject its
38
    # own logic there. This globally affects the behavior of the standard library
39
    # pickler for any user who transitively depends on this module!
40
    # Undo this extension to avoid altering the behavior of the pickler globally.
41
    dill.extend(use_dill=False)
42
    return dill
43