mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
4753100a3b
Summary: Generally wildcard imports are bad for the reasons described here: https://www.flake8rules.com/rules/F403.html This PR replaces wildcard imports with an explicit list of imported items where possible, and adds a `# noqa: F403` comment in the other cases (mostly re-exports in `__init__.py` files). This is a prerequisite for https://github.com/pytorch/pytorch/issues/55816, because currently [`tools/codegen/dest/register_dispatch_key.py` simply fails if you sort its imports](https://github.com/pytorch/pytorch/actions/runs/742505908). Pull Request resolved: https://github.com/pytorch/pytorch/pull/55838 Test Plan: CI. You can also run `flake8` locally. Reviewed By: jbschlosser Differential Revision: D27724232 Pulled By: samestep fbshipit-source-id: 269fb09cb4168f8a51fd65bfaacc6cda7fb87c34
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
|
|
import torch
|
|
import sys
|
|
import os
|
|
from enum import Enum
|
|
|
|
|
|
def is_available() -> bool:
|
|
"""
|
|
Returns ``True`` if the distributed package is available. Otherwise,
|
|
``torch.distributed`` does not expose any other APIs. Currently,
|
|
``torch.distributed`` is available on Linux, MacOS and Windows. Set
|
|
``USE_DISTRIBUTED=1`` to enable it when building PyTorch from source.
|
|
Currently, the default value is ``USE_DISTRIBUTED=1`` for Linux and Windows,
|
|
``USE_DISTRIBUTED=0`` for MacOS.
|
|
"""
|
|
return hasattr(torch._C, "_c10d_init")
|
|
|
|
|
|
if is_available() and not torch._C._c10d_init():
|
|
raise RuntimeError("Failed to initialize torch.distributed")
|
|
|
|
|
|
if is_available():
|
|
from torch._C._distributed_c10d import (
|
|
Store,
|
|
FileStore,
|
|
TCPStore,
|
|
ProcessGroup,
|
|
Reducer,
|
|
Logger,
|
|
BuiltinCommHookType,
|
|
GradBucket,
|
|
_DEFAULT_FIRST_BUCKET_BYTES,
|
|
_register_comm_hook,
|
|
_register_builtin_comm_hook,
|
|
_broadcast_coalesced,
|
|
_compute_bucket_assignment_by_size,
|
|
_verify_model_across_ranks,
|
|
_test_python_store,
|
|
_DistributedDebugLevel,
|
|
_get_debug_mode
|
|
)
|
|
if sys.platform != 'win32':
|
|
from torch._C._distributed_c10d import (
|
|
HashStore,
|
|
_round_robin_process_groups,
|
|
)
|
|
|
|
from .distributed_c10d import * # noqa: F403
|
|
# Variables prefixed with underscore are not auto imported
|
|
# See the comment in `distributed_c10d.py` above `_backend` on why we expose
|
|
# this.
|
|
|
|
from .distributed_c10d import _backend
|