|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import socket |
|
|
from string import Template |
|
|
from typing import List, Any |
|
|
|
|
|
|
|
|
def get_env_variable_or_raise(env_name: str) -> str: |
|
|
r""" |
|
|
Tries to retrieve environment variable. Raises ``ValueError`` |
|
|
if no environment variable found. |
|
|
|
|
|
Args: |
|
|
env_name (str): Name of the env variable |
|
|
""" |
|
|
value = os.environ.get(env_name, None) |
|
|
if value is None: |
|
|
msg = f"Environment variable {env_name} expected, but not set" |
|
|
raise ValueError(msg) |
|
|
return value |
|
|
|
|
|
|
|
|
def get_socket_with_port() -> socket.socket: |
|
|
addrs = socket.getaddrinfo( |
|
|
host="localhost", port=None, family=socket.AF_UNSPEC, type=socket.SOCK_STREAM |
|
|
) |
|
|
for addr in addrs: |
|
|
family, type, proto, _, _ = addr |
|
|
s = socket.socket(family, type, proto) |
|
|
try: |
|
|
s.bind(("localhost", 0)) |
|
|
s.listen(0) |
|
|
return s |
|
|
except OSError as e: |
|
|
s.close() |
|
|
raise RuntimeError("Failed to create a socket") |
|
|
|
|
|
|
|
|
class macros: |
|
|
""" |
|
|
Defines simple macros for caffe2.distributed.launch cmd args substitution |
|
|
""" |
|
|
|
|
|
local_rank = "${local_rank}" |
|
|
|
|
|
@staticmethod |
|
|
def substitute(args: List[Any], local_rank: str) -> List[str]: |
|
|
args_sub = [] |
|
|
for arg in args: |
|
|
if isinstance(arg, str): |
|
|
sub = Template(arg).safe_substitute(local_rank=local_rank) |
|
|
args_sub.append(sub) |
|
|
else: |
|
|
args_sub.append(arg) |
|
|
return args_sub |
|
|
|
