mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
f8fafdc7a6
This reverts commit ab26d4fbeb5bc4b4e6ef1c37fbec9fab6e5a9edd. Reverted https://github.com/pytorch/pytorch/pull/158288 on behalf of https://github.com/ZainRizvi due to Reverting as per offline discussion to fix internal breaks. @PaliC will reland this as a codev diff. Instructions here: https://fburl.com/fixing-ghfirst-reverts ([comment](https://github.com/pytorch/pytorch/pull/158288#issuecomment-3119037960))
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import functools
|
|
import importlib.util
|
|
from types import ModuleType
|
|
from typing import Optional
|
|
|
|
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() -> bool:
|
|
return (
|
|
_check_module_exists("dill")
|
|
# dill fails to import under torchdeploy
|
|
and not torch._running_with_deploy()
|
|
)
|
|
|
|
|
|
@functools.lru_cache
|
|
def import_dill() -> Optional[ModuleType]:
|
|
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
|