mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Improving typing and typing-related performance in rendezvous.py
### **SUMMARY:** It is unnecessary to perform `n + 1` calls to `cast` (one cast for each parameter name-value pair and one cast for the filter generator itself) in a dictionary comprehension in an effort to avoid mypy errors. Previously, the `cast` to `Tuple[str, str]` was necessary to prevent mypy from complaining that we are trying to create a dictionary out of lists as opposed to tuples (see the mypy issue [here](https://github.com/python/mypy/issues/7509)). However, a `cast` is both adding unnecessary overhead due to the function call and should generally only be used when a linter is unable to properly infer the type of a variable, not to "lie" to it about the type. We can avoid this by instead using a generator within the dictionary comprehension and then indexing into it twice to produce tuples of size 2; mypy recognizes this as a valid dictionary initialization. The above change is much more performant than the previous version of the code. Timing the two versions of the dictionary construction yielded the following results: ``` >python -m timeit -s "from typing import cast, Dict, Tuple, Iterable" -n 100000 -p "dict(cast(Tuple[str, str], pair.split('=')) for pair in cast(Iterable[str], filter(None, 'rank=3&world_size=5'.split('&'))))" 100000 loops, best of 5: 2.66 usec per loop >python -m timeit -n 100000 -p "dict((pair[0], pair[1]) for pair in (pair.split('=') for pair in filter(None, 'rank=3&world_size=5'.split('&'))))" 100000 loops, best of 5: 1.09 usec per loop ``` The `cast` to `Iterable[str]` is similarly a "lie" that is not necessary. It is best to be as transparent as possible to the linter rather than using `cast` to eliminate errors. This actually does not even produce any mypy errors when removed in isolation from the other changes. Further, it is good practice to type hint the return value of a function rather than specifying the type of the return value inside the function. Thus, the unnecessary intermediate variable `query_dict` inside `_query_to_dict` was removed, and the type hint of the return value was moved to the function declaration. The type of the argument `query` is specified as `str`. ### **EDITS (additional commits):** [The sole type hint for `query_dict` (in `_env_rendezvous_handler`) was removed to match all other functions in the file.](76d78bfc9c
) [Incorrect typing is fixed for _env_rendezvous_handler typing so that `rank`, `world_size`, `master_port`, and `master_addr` are specified to be `int`, `int`, `int`, and `str`, respectively.](3cc5844264
) Pull Request resolved: https://github.com/pytorch/pytorch/pull/75959 Approved by: https://github.com/kumpera, https://github.com/mrshenli
This commit is contained in:
@ -9,7 +9,7 @@ import numbers
|
||||
import os
|
||||
import sys
|
||||
from datetime import timedelta
|
||||
from typing import cast, Dict, Iterable, Optional, Tuple, Union
|
||||
from typing import Dict
|
||||
|
||||
import torch._six as six
|
||||
from torch.distributed import FileStore, PrefixStore, Store, TCPStore
|
||||
@ -52,11 +52,8 @@ 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):
|
||||
query_dict: Dict[str, str] = dict(
|
||||
cast(Tuple[str, str], pair.split("=")) for pair in cast(Iterable[str], filter(None, query.split("&")))
|
||||
)
|
||||
return query_dict
|
||||
def _query_to_dict(query: str) -> Dict[str, str]:
|
||||
return dict((pair[0], pair[1]) for pair in (pair.split("=") for pair in filter(None, query.split("&"))))
|
||||
|
||||
def rendezvous(url: str, rank: int = -1, world_size: int = -1, **kwargs):
|
||||
if not isinstance(url, six.string_classes):
|
||||
@ -107,8 +104,8 @@ def _create_store_from_options(backend_options, rank):
|
||||
query_dict = _query_to_dict(result.query)
|
||||
# if rank is -1 then intentionally exclude rank for the query, error will be thrown later
|
||||
if rank != -1:
|
||||
query_dict["rank"] = rank
|
||||
query_dict["world_size"] = world_size
|
||||
query_dict["rank"] = str(rank)
|
||||
query_dict["world_size"] = str(world_size)
|
||||
|
||||
result = result._replace(
|
||||
query="{}".format(
|
||||
@ -237,11 +234,12 @@ def _env_rendezvous_handler(
|
||||
return env_val
|
||||
|
||||
result = urlparse(url)
|
||||
query_dict: Dict[str, Union[int, str]] = _query_to_dict(result.query)
|
||||
query_dict = _query_to_dict(result.query)
|
||||
|
||||
rank: Optional[Union[str, int]]
|
||||
world_size: Optional[Union[str, int]]
|
||||
master_port: Optional[Union[str, int]]
|
||||
rank: int
|
||||
world_size: int
|
||||
master_port: int
|
||||
master_addr: str
|
||||
|
||||
if "rank" in query_dict:
|
||||
rank = int(query_dict["rank"])
|
||||
|
Reference in New Issue
Block a user