Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| 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}" | |
| 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 | |