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:
Ishaan
2022-04-24 21:49:51 +00:00
committed by PyTorch MergeBot
parent f980c3c193
commit d65ab9a689

View File

@ -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"])