frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

PEP585 update - torch/distributed (#145164)

Browse Source 
See #145101 for details.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/145164
Approved by: https://github.com/bobrenjc93
This commit is contained in:
Aaron Orenstein
2025-01-20 14:50:01 -08:00
committed by PyTorch MergeBot
parent c6986ca2e1
commit 00ffeca1b1
79 changed files with 805 additions and 860 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/distributed/rpc/api.py
View File

@ -7,7 +7,7 @@ import functools
import inspect
import logging
import threading
from typing import Any, Dict, Generic, Set, TYPE_CHECKING, TypeVar
from typing import Any, Generic, TYPE_CHECKING, TypeVar
import torch
from torch._C._distributed_rpc import (
@ -115,9 +115,9 @@ class AllGatherStates:
# States used by `def _all_gather()`.
# `_ALL_WORKER_NAMES` is initialized on initializing RPC layer.
_ALL_WORKER_NAMES: Set[Any] = set()
_ALL_WORKER_NAMES: set[Any] = set()
_all_gather_dict_lock = threading.RLock()
_all_gather_sequence_id: Dict[str, int] = {}
_all_gather_sequence_id: dict[str, int] = {}
_all_gather_sequence_id_to_states: collections.defaultdict = collections.defaultdict(
AllGatherStates
)