1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
|
try:
from urllib.parse import urlparse, urlunparse
except ImportError:
from urlparse import urlparse, urlunparse
import torch._six as six
import numbers
import os
import sys
from . import FileStore
from .constants import default_pg_timeout
if sys.platform != 'win32':
from . import TCPStore
_rendezvous_handlers = {}
def register_rendezvous_handler(scheme, handler):
"""Registers a new rendezvous handler.
Before we can run collective algorithms, participating processes
need to find each other and exchange information to be able to
communicate. We call this process rendezvous.
The outcome of the rendezvous process is a triplet containing a
shared key/value store, the rank of the process, and the total
number of participating processes.
If none of the bundled rendezvous methods apply to your execution
environment you can opt to register your own rendezvous handler.
Pick a unique name and use the URL scheme to identify it when
calling the `rendezvous()` function.
Arguments:
scheme (str): URL scheme to identify your rendezvous handler.
handler (function): Handler that is invoked when the
`rendezvous()` function is called with a URL that uses
the corresponding scheme. It must be a generator function
that yields the triplet.
"""
global _rendezvous_handlers
if scheme in _rendezvous_handlers:
raise RuntimeError(
"Rendezvous handler for {}:// already registered".format(scheme)
)
_rendezvous_handlers[scheme] = handler
def rendezvous(url, rank=-1, world_size=-1, **kwargs):
if not isinstance(url, six.string_classes):
raise RuntimeError("`url` must be a string. {}: {}".format(type(url), url))
if not isinstance(rank, numbers.Integral):
raise RuntimeError("`rank` must be an integer. {}".format(rank))
if not isinstance(world_size, numbers.Integral):
raise RuntimeError("`world_size` must be an integer. {}".format(world_size))
# Append node-specific arguments.
result = urlparse(url)
if rank != -1 or world_size != -1:
query_dict = dict(
pair.split("=") for pair in filter(None, result.query.split("&"))
)
assert (
"rank" not in query_dict and "world_size" not in query_dict
), "The url: {url} has node-specific arguments(rank, world_size) already.".format(
url=url
)
if rank != -1:
query_dict["rank"] = rank
if world_size != -1:
query_dict["world_size"] = world_size
result = result._replace(
query="{}".format("&".join(["{}={}".format(k, v) for k, v in query_dict.items()]))
)
url = urlunparse(result)
if result.scheme not in _rendezvous_handlers:
raise RuntimeError("No rendezvous handler for {}://".format(result.scheme))
return _rendezvous_handlers[result.scheme](url, **kwargs)
def _rendezvous_error(msg):
return ValueError("Error initializing torch.distributed using " + msg)
def _file_rendezvous_handler(url, **kwargs):
def _error(msg):
return _rendezvous_error("file:// rendezvous: " + msg)
result = urlparse(url)
path = result.path
if sys.platform == 'win32':
import urllib.request
path = urllib.request.url2pathname(result.path)
if not path:
raise _error("path missing")
query = dict(pair.split("=") for pair in filter(None, result.query.split("&")))
if "rank" not in query:
raise _error("rank parameter missing")
if "world_size" not in query:
raise _error("world size parameter missing")
rank = int(query["rank"])
world_size = int(query["world_size"])
store = FileStore(path, world_size)
yield (store, rank, world_size)
# If this configuration is invalidated, there is nothing we can do about it
raise RuntimeError("Unable to perform rerendezvous using file:// method")
def _tcp_rendezvous_handler(url, timeout=default_pg_timeout, **kwargs):
def _error(msg):
return _rendezvous_error("tcp:// rendezvous: " + msg)
result = urlparse(url)
if not result.port:
raise _error("port number missing")
query = dict(pair.split("=") for pair in filter(None, result.query.split("&")))
if "rank" not in query:
raise _error("rank parameter missing")
if "world_size" not in query:
raise _error("world size parameter missing")
rank = int(query["rank"])
world_size = int(query["world_size"])
start_daemon = rank == 0
store = TCPStore(result.hostname, result.port, world_size, start_daemon, timeout)
yield (store, rank, world_size)
# If this configuration is invalidated, there is nothing we can do about it
raise RuntimeError("Unable to perform rerendezvous using tcp:// method")
def _env_rendezvous_handler(url, timeout=default_pg_timeout, **kwargs):
def _error(msg):
return _rendezvous_error("env:// rendezvous: " + msg)
def _env_error(var):
return _error("environment variable %s expected, but not set" % var)
result = urlparse(url)
query = dict(pair.split("=") for pair in filter(None, result.query.split("&")))
if "rank" in query:
rank = int(query["rank"])
else:
rank = os.environ.get("RANK", None)
if rank is None:
raise _env_error("RANK")
if "world_size" in query:
world_size = int(query["world_size"])
else:
world_size = os.environ.get("WORLD_SIZE", None)
if world_size is None:
raise _env_error("WORLD_SIZE")
master_addr = os.environ.get("MASTER_ADDR", None)
if master_addr is None:
raise _env_error("MASTER_ADDR")
master_port = os.environ.get("MASTER_PORT", None)
if master_port is None:
raise _env_error("MASTER_PORT")
# Converting before creating the store
rank = int(rank)
world_size = int(world_size)
master_port = int(master_port)
# Now start the TCP store daemon on the rank 0
start_daemon = rank == 0
store = TCPStore(master_addr, master_port, world_size, start_daemon, timeout)
yield (store, rank, world_size)
# If this configuration is invalidated, there is nothing we can do about it
raise RuntimeError("Unable to perform rerendezvous using env:// method")
if sys.platform != 'win32':
register_rendezvous_handler("tcp", _tcp_rendezvous_handler)
register_rendezvous_handler("env", _env_rendezvous_handler)
register_rendezvous_handler("file", _file_rendezvous_handler)
|
