frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +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

9
torch/distributed/rendezvous.py
View File

@ -9,15 +9,16 @@ except ImportError as e:
import numbers
import os
import sys
from collections.abc import Iterator
from datetime import timedelta
from typing import Callable, Dict, Iterator, Optional
from typing import Callable, Optional
from torch.distributed import FileStore, Store, TCPStore
from .constants import default_pg_timeout
_rendezvous_handlers: Dict[str, Callable[..., Iterator[tuple[Store, int, int]]]] = {}
_rendezvous_handlers: dict[str, Callable[..., Iterator[tuple[Store, int, int]]]] = {}
__all__ = ["register_rendezvous_handler", "rendezvous"]
@ -54,14 +55,14 @@ def register_rendezvous_handler(scheme, handler):
# Query will have format "rank=0&world_size=1" and is
# converted into {"rank": 0, "world_size": 1}
def _query_to_dict(query: str) -> Dict[str, str]:
def _query_to_dict(query: str) -> dict[str, str]:
return {
pair[0]: pair[1]
for pair in (pair.split("=") for pair in filter(None, query.split("&")))
}
def _get_use_libuv_from_query_dict(query_dict: Dict[str, str]) -> bool:
def _get_use_libuv_from_query_dict(query_dict: dict[str, str]) -> bool:
# libuv is the default backend for TCPStore. To enable the non-libuv backend,
# user can explicitly specify ``use_libuv=0`` in the URL parameter.
return query_dict.get("use_libuv", os.environ.get("USE_LIBUV", "1")) == "1"