Skip to content

Commit

Permalink
Improving typing and typing-related performance in rendezvous.py
Browse files Browse the repository at this point in the history
### **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](python/mypy#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.](76d78bf)

[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.](3cc5844)
Pull Request resolved: #75959
Approved by: https://github.com/kumpera, https://github.com/mrshenli
  • Loading branch information
ishaan-mehta authored and pytorchmergebot committed Apr 24, 2022
1 parent f980c3c commit d65ab9a
Showing 1 changed file with 10 additions and 12 deletions.
22 changes: 10 additions & 12 deletions torch/distributed/rendezvous.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -237,11 +234,12 @@ def _get_env_or_raise(env_var: str) -> str:
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"])
Expand Down

0 comments on commit d65ab9a

Please sign in to comment.