mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
5033a1ca0d
1. We want to take option 3 as discussed in https://github.com/pytorch/pytorch/issues/135712, so every time when we retry, we create a new TCPStore server first so that we don't need to append attempt count as prefix and avoid eventually TCPStore sync failure. (This is only for the TCPStore sharing enabled case) 2. We start a new server bound to an ephemeral port (i.e. 0) so it gets assigned to a free port. We then pass that downstream (trainer or c10d). By doing so, TCPStore is managed by the elastic agent rather than having a race condition on binding to a specific port in the trainer. 3. Then the port be broadcasted for dynamic_rendezvous. Only one more question, what do we do about the store created from (_create_tcp_store) torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py, are we ok with creating a duplicate TCPStore server? Pull Request resolved: https://github.com/pytorch/pytorch/pull/135957 Approved by: https://github.com/d4l3k, https://github.com/c-p-i-o
287 lines
9.7 KiB
Python
287 lines
9.7 KiB
Python
# mypy: allow-untyped-defs
|
|
try:
|
|
from urllib.parse import urlparse, urlunparse
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"urllib cannot be found, urlparse from python2 is no longer supported."
|
|
) from e
|
|
|
|
import numbers
|
|
import os
|
|
import sys
|
|
from datetime import timedelta
|
|
from typing import Callable, Dict, Iterator, Optional, Tuple
|
|
|
|
from torch.distributed import FileStore, Store, TCPStore
|
|
|
|
from .constants import default_pg_timeout
|
|
|
|
|
|
_rendezvous_handlers: Dict[str, Callable[..., Iterator[Tuple[Store, int, int]]]] = {}
|
|
|
|
__all__ = ["register_rendezvous_handler", "rendezvous"]
|
|
|
|
|
|
def register_rendezvous_handler(scheme, handler):
|
|
"""
|
|
Register 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.
|
|
|
|
Args:
|
|
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(f"Rendezvous handler for {scheme}:// already registered")
|
|
_rendezvous_handlers[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]:
|
|
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:
|
|
# 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"
|
|
|
|
|
|
def _rendezvous_helper(url: str, rank: int, world_size_opt: Optional[int], **kwargs):
|
|
result = urlparse(url)
|
|
if world_size_opt is None:
|
|
world_size = -1
|
|
if result.scheme == "env":
|
|
rank = int(os.environ.get("RANK", rank))
|
|
# If the world_size env variable is not present then it is a dynamic group
|
|
world_size = int(os.environ.get("WORLD_SIZE", world_size))
|
|
else:
|
|
world_size = world_size_opt
|
|
if rank != -1 or world_size != -1 or world_size_opt is None:
|
|
query_dict = _query_to_dict(result.query)
|
|
assert (
|
|
"rank" not in query_dict and "world_size" not in query_dict
|
|
), f"The url: {url} has node-specific arguments(rank, world_size) already."
|
|
if rank != -1:
|
|
query_dict["rank"] = str(rank)
|
|
if world_size != -1 or world_size_opt is None:
|
|
query_dict["world_size"] = str(world_size)
|
|
result = result._replace(
|
|
query=f"{'&'.join([f'{k}={v}' for k, v in query_dict.items()])}"
|
|
)
|
|
url = urlunparse(result)
|
|
|
|
if result.scheme not in _rendezvous_handlers:
|
|
raise RuntimeError(f"No rendezvous handler for {result.scheme}://")
|
|
return _rendezvous_handlers[result.scheme](url, **kwargs)
|
|
|
|
|
|
def rendezvous(url: str, rank: int = -1, world_size: int = -1, **kwargs):
|
|
if not isinstance(url, (str, bytes)):
|
|
raise RuntimeError(f"`url` must be a string. {type(url)}: {url}")
|
|
|
|
if not isinstance(rank, numbers.Integral):
|
|
raise RuntimeError(f"`rank` must be an integer. {rank}")
|
|
|
|
if not isinstance(world_size, numbers.Integral):
|
|
raise RuntimeError(f"`world_size` must be an integer. {world_size}")
|
|
|
|
return _rendezvous_helper(url, rank, world_size, **kwargs)
|
|
|
|
|
|
def _create_store_from_options(backend_options, rank):
|
|
store, _, _ = next(_rendezvous_helper(backend_options.init_method, rank, None))
|
|
return store
|
|
|
|
|
|
def _rendezvous_error(msg):
|
|
return ValueError("Error initializing torch.distributed using " + msg)
|
|
|
|
|
|
def _file_rendezvous_handler(url: str, **kwargs):
|
|
def _error(msg):
|
|
return _rendezvous_error("file:// rendezvous: " + msg)
|
|
|
|
result = urlparse(url)
|
|
path = result.path
|
|
if sys.platform == "win32":
|
|
import urllib.request
|
|
|
|
full_path = result.netloc + result.path
|
|
path = urllib.request.url2pathname(full_path)
|
|
if path:
|
|
# Normalizing an empty string produces ".", which is not expected.
|
|
path = os.path.normpath(path)
|
|
|
|
if not path:
|
|
raise _error("path missing")
|
|
query_dict = _query_to_dict(result.query)
|
|
if "rank" not in query_dict:
|
|
raise _error("rank parameter missing")
|
|
if "world_size" not in query_dict:
|
|
raise _error("world size parameter missing")
|
|
|
|
rank = int(query_dict["rank"])
|
|
world_size = int(query_dict["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 _torchelastic_use_agent_store() -> bool:
|
|
return os.environ.get("TORCHELASTIC_USE_AGENT_STORE", None) == str(True)
|
|
|
|
|
|
def _create_c10d_store(
|
|
hostname, port, rank, world_size, timeout, use_libuv=True
|
|
) -> Store:
|
|
"""
|
|
Smartly creates a c10d Store object on ``rank`` based on whether we need to re-use agent store.
|
|
|
|
The TCPStore server is assumed to be hosted
|
|
on ``hostname:port``.
|
|
|
|
By default, the TCPStore server uses the asynchronous implementation
|
|
``LibUVStoreDaemon`` which utilizes libuv.
|
|
|
|
If ``torchelastic_use_agent_store()`` is ``True``, then it is assumed that
|
|
the agent leader (node rank 0) hosts the TCPStore server (for which the
|
|
endpoint is specified by the given ``hostname:port``). Hence
|
|
ALL ranks will create and return a TCPStore client (e.g. ``start_daemon=False``).
|
|
|
|
If ``torchelastic_use_agent_store()`` is ``False``, then rank 0 will host
|
|
the TCPStore (with multi-tenancy) and it is assumed that rank 0's hostname
|
|
and port are correctly passed via ``hostname`` and ``port``. All
|
|
non-zero ranks will create and return a TCPStore client.
|
|
"""
|
|
# check if port is uint16_t
|
|
if not 0 <= port < 2**16:
|
|
raise ValueError(f"port must have value from 0 to 65535 but was {port}.")
|
|
|
|
if _torchelastic_use_agent_store():
|
|
# We create a new TCPStore for every retry so no need to add prefix for each attempt.
|
|
return TCPStore(
|
|
host_name=hostname,
|
|
port=port,
|
|
world_size=world_size,
|
|
is_master=False,
|
|
timeout=timeout,
|
|
)
|
|
else:
|
|
start_daemon = rank == 0
|
|
return TCPStore(
|
|
host_name=hostname,
|
|
port=port,
|
|
world_size=world_size,
|
|
is_master=start_daemon,
|
|
timeout=timeout,
|
|
multi_tenant=True,
|
|
use_libuv=use_libuv,
|
|
)
|
|
|
|
|
|
def _tcp_rendezvous_handler(
|
|
url: str, timeout: timedelta = 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 = _query_to_dict(result.query)
|
|
if "rank" not in query_dict:
|
|
raise _error("rank parameter missing")
|
|
if "world_size" not in query_dict:
|
|
raise _error("world size parameter missing")
|
|
|
|
rank = int(query_dict["rank"])
|
|
world_size = int(query_dict["world_size"])
|
|
use_libuv = _get_use_libuv_from_query_dict(query_dict)
|
|
|
|
assert result.hostname is not None
|
|
|
|
store = _create_c10d_store(
|
|
result.hostname, result.port, rank, world_size, timeout, use_libuv
|
|
)
|
|
|
|
yield (store, rank, world_size)
|
|
|
|
# If this configuration is invalidated, there is nothing we can do about it
|
|
raise RuntimeError("Unable to perform re-rendezvous using tcp:// method")
|
|
|
|
|
|
def _env_rendezvous_handler(
|
|
url: str, timeout: timedelta = default_pg_timeout, **kwargs
|
|
):
|
|
def _error(msg):
|
|
return _rendezvous_error("env:// rendezvous: " + msg)
|
|
|
|
def _env_error(var):
|
|
return _error(f"environment variable {var} expected, but not set")
|
|
|
|
def _get_env_or_raise(env_var: str) -> str:
|
|
env_val = os.environ.get(env_var, None)
|
|
if not env_val:
|
|
raise _env_error(env_var)
|
|
else:
|
|
return env_val
|
|
|
|
result = urlparse(url)
|
|
query_dict = _query_to_dict(result.query)
|
|
|
|
rank: int
|
|
world_size: int
|
|
master_port: int
|
|
master_addr: str
|
|
|
|
if "rank" in query_dict:
|
|
rank = int(query_dict["rank"])
|
|
else:
|
|
rank = int(_get_env_or_raise("RANK"))
|
|
|
|
if "world_size" in query_dict:
|
|
world_size = int(query_dict["world_size"])
|
|
else:
|
|
world_size = int(_get_env_or_raise("WORLD_SIZE"))
|
|
|
|
master_addr = _get_env_or_raise("MASTER_ADDR")
|
|
master_port = int(_get_env_or_raise("MASTER_PORT"))
|
|
use_libuv = _get_use_libuv_from_query_dict(query_dict)
|
|
|
|
store = _create_c10d_store(
|
|
master_addr, master_port, rank, world_size, timeout, use_libuv
|
|
)
|
|
|
|
yield (store, rank, world_size)
|
|
|
|
# If this configuration is invalidated, there is nothing we can do about it
|
|
raise RuntimeError("Unable to perform re-rendezvous using env:// method")
|
|
|
|
|
|
register_rendezvous_handler("tcp", _tcp_rendezvous_handler)
|
|
register_rendezvous_handler("env", _env_rendezvous_handler)
|
|
register_rendezvous_handler("file", _file_rendezvous_handler)
|