wehub-resource-sync
653 lines
19 KiB
Python
653 lines
19 KiB
Python
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from typing import Any, Iterator, List, Optional, Sequence, Tuple
|
|
|
|
import yaml
|
|
|
|
import ray # noqa: F401
|
|
from ray.autoscaler._private.cli_logger import cli_logger
|
|
from ray.autoscaler._private.providers import _get_node_provider
|
|
from ray.autoscaler.tags import NODE_KIND_HEAD, NODE_KIND_WORKER, TAG_RAY_NODE_KIND
|
|
|
|
# Import psutil after ray so the packaged version is used.
|
|
import psutil
|
|
|
|
MAX_PARALLEL_SSH_WORKERS = 8
|
|
DEFAULT_SSH_USER = "ubuntu"
|
|
DEFAULT_SSH_KEYS = ["~/ray_bootstrap_key.pem", "~/.ssh/ray-autoscaler_2_us-west-2.pem"]
|
|
|
|
|
|
class CommandFailed(RuntimeError):
|
|
pass
|
|
|
|
|
|
class LocalCommandFailed(CommandFailed):
|
|
pass
|
|
|
|
|
|
class RemoteCommandFailed(CommandFailed):
|
|
pass
|
|
|
|
|
|
class GetParameters:
|
|
def __init__(
|
|
self,
|
|
logs: bool = True,
|
|
debug_state: bool = True,
|
|
pip: bool = True,
|
|
processes: bool = True,
|
|
processes_verbose: bool = True,
|
|
processes_list: Optional[List[Tuple[str, bool]]] = None,
|
|
):
|
|
self.logs = logs
|
|
self.debug_state = debug_state
|
|
self.pip = pip
|
|
self.processes = processes
|
|
self.processes_verbose = processes_verbose
|
|
self.processes_list = processes_list
|
|
|
|
|
|
class Node:
|
|
"""Node (as in "machine")"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
ssh_user: str = "ubuntu",
|
|
ssh_key: str = "~/ray_bootstrap_key.pem",
|
|
docker_container: Optional[str] = None,
|
|
is_head: bool = False,
|
|
):
|
|
self.host = host
|
|
self.ssh_user = ssh_user
|
|
self.ssh_key = ssh_key
|
|
self.docker_container = docker_container
|
|
self.is_head = is_head
|
|
|
|
|
|
class Archive:
|
|
"""Archive object to collect and compress files into a single file.
|
|
|
|
Objects of this class can be passed around to different data collection
|
|
functions. These functions can use the :meth:`subdir` method to add
|
|
files to a sub directory of the archive.
|
|
|
|
"""
|
|
|
|
def __init__(self, file: Optional[str] = None):
|
|
self.file = file or tempfile.mkstemp(prefix="ray_logs_", suffix=".tar.gz")[1]
|
|
self.tar = None
|
|
self._lock = threading.Lock()
|
|
|
|
@property
|
|
def is_open(self):
|
|
return bool(self.tar)
|
|
|
|
def open(self):
|
|
self.tar = tarfile.open(self.file, "w:gz")
|
|
|
|
def close(self):
|
|
self.tar.close()
|
|
self.tar = None
|
|
|
|
def __enter__(self):
|
|
self.open()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.close()
|
|
|
|
@contextmanager
|
|
def subdir(self, subdir: str, root: Optional[str] = "/") -> Iterator[Any]:
|
|
"""Open a context to add files to the archive.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
with Archive("file.tar.gz") as archive:
|
|
with archive.subdir("logfiles", root="/tmp/logs") as sd:
|
|
# Will be added as `logfiles/nested/file.txt`
|
|
sd.add("/tmp/logs/nested/file.txt")
|
|
|
|
Args:
|
|
subdir: Subdir to which to add files to. Calling the
|
|
``add(path)`` command will place files into the ``subdir``
|
|
directory of the archive.
|
|
root: Root path. Files without an explicit ``arcname``
|
|
will be named relatively to this path.
|
|
|
|
Yields:
|
|
Any: A context object that can be used to add files to the archive.
|
|
"""
|
|
root = os.path.abspath(root)
|
|
|
|
class _Context:
|
|
@staticmethod
|
|
def add(path: str, arcname: Optional[str] = None):
|
|
path = os.path.abspath(path)
|
|
arcname = arcname or os.path.join(subdir, os.path.relpath(path, root))
|
|
|
|
self._lock.acquire()
|
|
self.tar.add(path, arcname=arcname)
|
|
self._lock.release()
|
|
|
|
yield _Context()
|
|
|
|
|
|
###
|
|
# Functions to gather logs and information on the local node
|
|
###
|
|
|
|
|
|
def get_local_ray_logs(
|
|
archive: Archive,
|
|
exclude: Optional[Sequence[str]] = None,
|
|
session_log_dir: str = "/tmp/ray/session_latest",
|
|
) -> Archive:
|
|
"""Copy local log files into an archive.
|
|
|
|
Args:
|
|
archive: Archive object to add log files to.
|
|
exclude: Sequence of regex patterns. Files that match
|
|
any of these patterns will not be included in the archive.
|
|
session_log_dir: Path to the Ray session files. Defaults to
|
|
``/tmp/ray/session_latest``
|
|
|
|
Returns:
|
|
Open archive object.
|
|
|
|
"""
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
exclude = exclude or []
|
|
|
|
session_log_dir = os.path.join(os.path.expanduser(session_log_dir), "logs")
|
|
|
|
with archive.subdir("logs", root=session_log_dir) as sd:
|
|
for root, dirs, files in os.walk(session_log_dir):
|
|
for file in files:
|
|
file_path = os.path.join(root, file)
|
|
rel_path = os.path.relpath(file_path, start=session_log_dir)
|
|
# Skip file if it matches any pattern in `exclude`
|
|
if any(re.match(pattern, rel_path) for pattern in exclude):
|
|
continue
|
|
sd.add(file_path)
|
|
|
|
return archive
|
|
|
|
|
|
def get_local_debug_state(
|
|
archive: Archive, session_dir: str = "/tmp/ray/session_latest"
|
|
) -> Archive:
|
|
"""Copy local log files into an archive.
|
|
|
|
Args:
|
|
archive: Archive object to add log files to.
|
|
session_dir: Path to the Ray session files. Defaults to
|
|
``/tmp/ray/session_latest``
|
|
|
|
Returns:
|
|
Open archive object.
|
|
|
|
"""
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
session_dir = os.path.expanduser(session_dir)
|
|
debug_state_file = os.path.join(session_dir, "logs/debug_state.txt")
|
|
|
|
if not os.path.exists(debug_state_file):
|
|
raise LocalCommandFailed("No `debug_state.txt` file found.")
|
|
|
|
with archive.subdir("", root=session_dir) as sd:
|
|
sd.add(debug_state_file)
|
|
|
|
return archive
|
|
|
|
|
|
def get_local_pip_packages(archive: Archive):
|
|
"""Get currently installed pip packages and write into an archive.
|
|
|
|
Args:
|
|
archive: Archive object to add meta files to.
|
|
|
|
Returns:
|
|
Open archive object.
|
|
"""
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
try:
|
|
from pip._internal.operations import freeze
|
|
except ImportError: # pip < 10.0
|
|
from pip.operations import freeze
|
|
|
|
with tempfile.NamedTemporaryFile("wt") as fp:
|
|
for line in freeze.freeze():
|
|
fp.writelines([line, "\n"])
|
|
|
|
fp.flush()
|
|
with archive.subdir("") as sd:
|
|
sd.add(fp.name, "pip_packages.txt")
|
|
|
|
return archive
|
|
|
|
|
|
def get_local_ray_processes(
|
|
archive: Archive,
|
|
processes: Optional[List[Tuple[str, bool]]] = None,
|
|
verbose: bool = False,
|
|
):
|
|
"""Get the status of all the relevant ray processes.
|
|
Args:
|
|
archive: Archive object to add process info files to.
|
|
processes: List of processes to get information on. The first
|
|
element of the tuple is a string to filter by, and the second
|
|
element is a boolean indicating if we should filter by command
|
|
name (True) or command line including parameters (False)
|
|
verbose: If True, show entire executable command line.
|
|
If False, show just the first term.
|
|
Returns:
|
|
Open archive object.
|
|
"""
|
|
if not processes:
|
|
# local import to avoid circular dependencies
|
|
from ray.autoscaler._private.constants import RAY_PROCESSES
|
|
|
|
processes = RAY_PROCESSES
|
|
|
|
process_infos = []
|
|
for process in psutil.process_iter(["pid", "name", "cmdline", "status"]):
|
|
try:
|
|
with process.oneshot():
|
|
cmdline = " ".join(process.cmdline())
|
|
process_infos.append(
|
|
(
|
|
{
|
|
"executable": cmdline
|
|
if verbose
|
|
else cmdline.split("--", 1)[0][:-1],
|
|
"name": process.name(),
|
|
"pid": process.pid,
|
|
"status": process.status(),
|
|
},
|
|
process.cmdline(),
|
|
)
|
|
)
|
|
except Exception as exc:
|
|
raise LocalCommandFailed(exc) from exc
|
|
|
|
relevant_processes = {}
|
|
for process_dict, cmdline in process_infos:
|
|
for keyword, filter_by_cmd in processes:
|
|
if filter_by_cmd:
|
|
corpus = process_dict["name"]
|
|
else:
|
|
corpus = subprocess.list2cmdline(cmdline)
|
|
if keyword in corpus and process_dict["pid"] not in relevant_processes:
|
|
relevant_processes[process_dict["pid"]] = process_dict
|
|
|
|
with tempfile.NamedTemporaryFile("wt") as fp:
|
|
for line in relevant_processes.values():
|
|
fp.writelines([yaml.dump(line), "\n"])
|
|
|
|
fp.flush()
|
|
with archive.subdir("meta") as sd:
|
|
sd.add(fp.name, "process_info.txt")
|
|
|
|
return archive
|
|
|
|
|
|
def get_all_local_data(archive: Archive, parameters: GetParameters):
|
|
"""Get all local data.
|
|
|
|
Gets:
|
|
- The Ray logs of the latest session
|
|
- The currently installed pip packages
|
|
|
|
Args:
|
|
archive: Archive object to add meta files to.
|
|
parameters: Parameters (settings) for getting data.
|
|
|
|
Returns:
|
|
Open archive object.
|
|
"""
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
if parameters.logs:
|
|
try:
|
|
get_local_ray_logs(archive=archive)
|
|
except LocalCommandFailed as exc:
|
|
cli_logger.error(exc)
|
|
if parameters.debug_state:
|
|
try:
|
|
get_local_debug_state(archive=archive)
|
|
except LocalCommandFailed as exc:
|
|
cli_logger.error(exc)
|
|
if parameters.pip:
|
|
try:
|
|
get_local_pip_packages(archive=archive)
|
|
except LocalCommandFailed as exc:
|
|
cli_logger.error(exc)
|
|
if parameters.processes:
|
|
try:
|
|
get_local_ray_processes(
|
|
archive=archive,
|
|
processes=parameters.processes_list,
|
|
verbose=parameters.processes_verbose,
|
|
)
|
|
except LocalCommandFailed as exc:
|
|
cli_logger.error(exc)
|
|
|
|
return archive
|
|
|
|
|
|
###
|
|
# Functions to invoke remote scripts and gather data from remote nodes
|
|
###
|
|
|
|
|
|
def _wrap(items: List[str], quotes="'"):
|
|
return f"{quotes}{' '.join(items)}{quotes}"
|
|
|
|
|
|
def create_and_get_archive_from_remote_node(
|
|
remote_node: Node, parameters: GetParameters, script_path: str = "ray"
|
|
) -> Optional[str]:
|
|
"""Create an archive containing logs on a remote node and transfer.
|
|
|
|
This will call ``ray local-dump --stream`` on the remote
|
|
node. The resulting file will be saved locally in a temporary file and
|
|
returned.
|
|
|
|
Args:
|
|
remote_node: Remote node to gather archive from.
|
|
parameters: Parameters (settings) for getting data.
|
|
script_path: Path to this script on the remote node.
|
|
|
|
Returns:
|
|
Path to a temporary file containing the node's collected data.
|
|
|
|
"""
|
|
cmd = [
|
|
"ssh",
|
|
"-o StrictHostKeyChecking=no",
|
|
"-o UserKnownHostsFile=/dev/null",
|
|
"-o LogLevel=ERROR",
|
|
"-i",
|
|
remote_node.ssh_key,
|
|
f"{remote_node.ssh_user}@{remote_node.host}",
|
|
]
|
|
|
|
if remote_node.docker_container:
|
|
cmd += [
|
|
"docker",
|
|
"exec",
|
|
remote_node.docker_container,
|
|
]
|
|
|
|
collect_cmd = [script_path, "local-dump", "--stream"]
|
|
collect_cmd += ["--logs"] if parameters.logs else ["--no-logs"]
|
|
collect_cmd += ["--debug-state"] if parameters.debug_state else ["--no-debug-state"]
|
|
collect_cmd += ["--pip"] if parameters.pip else ["--no-pip"]
|
|
collect_cmd += ["--processes"] if parameters.processes else ["--no-processes"]
|
|
if parameters.processes:
|
|
collect_cmd += (
|
|
["--processes-verbose"]
|
|
if parameters.processes_verbose
|
|
else ["--no-proccesses-verbose"]
|
|
)
|
|
|
|
cmd += ["/bin/bash", "-c", _wrap(collect_cmd, quotes='"')]
|
|
|
|
cat = "node" if not remote_node.is_head else "head"
|
|
|
|
cli_logger.print(f"Collecting data from remote node: {remote_node.host}")
|
|
tmp = tempfile.mkstemp(prefix=f"ray_{cat}_{remote_node.host}_", suffix=".tar.gz")[1]
|
|
with open(tmp, "wb") as fp:
|
|
try:
|
|
subprocess.check_call(cmd, stdout=fp, stderr=sys.stderr)
|
|
except subprocess.CalledProcessError as exc:
|
|
raise RemoteCommandFailed(
|
|
f"Gathering logs from remote node failed: {' '.join(cmd)}"
|
|
) from exc
|
|
|
|
return tmp
|
|
|
|
|
|
def create_and_add_remote_data_to_local_archive(
|
|
archive: Archive, remote_node: Node, parameters: GetParameters
|
|
):
|
|
"""Create and get data from remote node and add to local archive.
|
|
|
|
Args:
|
|
archive: Archive object to add remote data to.
|
|
remote_node: Remote node to gather archive from.
|
|
parameters: Parameters (settings) for getting data.
|
|
|
|
Returns:
|
|
Open archive object.
|
|
"""
|
|
tmp = create_and_get_archive_from_remote_node(remote_node, parameters)
|
|
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
cat = "node" if not remote_node.is_head else "head"
|
|
|
|
with archive.subdir("", root=os.path.dirname(tmp)) as sd:
|
|
sd.add(tmp, arcname=f"ray_{cat}_{remote_node.host}.tar.gz")
|
|
|
|
return archive
|
|
|
|
|
|
def create_and_add_local_data_to_local_archive(
|
|
archive: Archive, parameters: GetParameters
|
|
):
|
|
"""Create and get data from this node and add to archive.
|
|
|
|
Args:
|
|
archive: Archive object to add remote data to.
|
|
parameters: Parameters (settings) for getting data.
|
|
|
|
Returns:
|
|
Open archive object.
|
|
"""
|
|
with Archive() as local_data_archive:
|
|
get_all_local_data(local_data_archive, parameters)
|
|
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
with archive.subdir("", root=os.path.dirname(local_data_archive.file)) as sd:
|
|
sd.add(local_data_archive.file, arcname="local_node.tar.gz")
|
|
|
|
os.remove(local_data_archive.file)
|
|
|
|
return archive
|
|
|
|
|
|
def create_archive_for_remote_nodes(
|
|
archive: Archive, remote_nodes: Sequence[Node], parameters: GetParameters
|
|
):
|
|
"""Create an archive combining data from the remote nodes.
|
|
|
|
This will parallelize calls to get data from remote nodes.
|
|
|
|
Args:
|
|
archive: Archive object to add remote data to.
|
|
remote_nodes: Sequence of remote nodes.
|
|
parameters: Parameters (settings) for getting data.
|
|
|
|
Returns:
|
|
Open archive object.
|
|
|
|
"""
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
with ThreadPoolExecutor(max_workers=MAX_PARALLEL_SSH_WORKERS) as executor:
|
|
for remote_node in remote_nodes:
|
|
executor.submit(
|
|
create_and_add_remote_data_to_local_archive,
|
|
archive=archive,
|
|
remote_node=remote_node,
|
|
parameters=parameters,
|
|
)
|
|
|
|
return archive
|
|
|
|
|
|
def create_archive_for_local_and_remote_nodes(
|
|
archive: Archive, remote_nodes: Sequence[Node], parameters: GetParameters
|
|
):
|
|
"""Create an archive combining data from the local and remote nodes.
|
|
|
|
This will parallelize calls to get data from remote nodes.
|
|
|
|
Args:
|
|
archive: Archive object to add data to.
|
|
remote_nodes: Sequence of remote nodes.
|
|
parameters: Parameters (settings) for getting data.
|
|
|
|
Returns:
|
|
Open archive object.
|
|
|
|
"""
|
|
if not archive.is_open:
|
|
archive.open()
|
|
|
|
try:
|
|
create_and_add_local_data_to_local_archive(archive, parameters)
|
|
except CommandFailed as exc:
|
|
cli_logger.error(exc)
|
|
|
|
create_archive_for_remote_nodes(archive, remote_nodes, parameters)
|
|
|
|
cli_logger.print(
|
|
f"Collected data from local node and {len(remote_nodes)} " f"remote nodes."
|
|
)
|
|
return archive
|
|
|
|
|
|
###
|
|
# Ray cluster info
|
|
###
|
|
def get_info_from_ray_cluster_config(
|
|
cluster_config: str,
|
|
) -> Tuple[List[str], str, str, Optional[str], Optional[str]]:
|
|
"""Get information from Ray cluster config.
|
|
|
|
Return list of host IPs, ssh user, ssh key file, and optional docker
|
|
container.
|
|
|
|
Args:
|
|
cluster_config: Path to ray cluster config.
|
|
|
|
Returns:
|
|
Tuple of list of host IPs, ssh user name, ssh key file path,
|
|
optional docker container name, optional cluster name.
|
|
"""
|
|
from ray.autoscaler._private.commands import _bootstrap_config
|
|
|
|
cli_logger.print(
|
|
f"Retrieving cluster information from ray cluster file: " f"{cluster_config}"
|
|
)
|
|
|
|
cluster_config = os.path.expanduser(cluster_config)
|
|
|
|
config = yaml.safe_load(open(cluster_config).read())
|
|
config = _bootstrap_config(config, no_config_cache=True)
|
|
|
|
provider = _get_node_provider(config["provider"], config["cluster_name"])
|
|
head_nodes = provider.non_terminated_nodes({TAG_RAY_NODE_KIND: NODE_KIND_HEAD})
|
|
worker_nodes = provider.non_terminated_nodes({TAG_RAY_NODE_KIND: NODE_KIND_WORKER})
|
|
|
|
hosts = [provider.external_ip(node) for node in head_nodes + worker_nodes]
|
|
ssh_user = config["auth"]["ssh_user"]
|
|
ssh_key = config["auth"]["ssh_private_key"]
|
|
|
|
docker = None
|
|
docker_config = config.get("docker", None)
|
|
if docker_config:
|
|
docker = docker_config.get("container_name", None)
|
|
|
|
cluster_name = config.get("cluster_name", None)
|
|
|
|
return hosts, ssh_user, ssh_key, docker, cluster_name
|
|
|
|
|
|
def _info_from_params(
|
|
cluster: Optional[str] = None,
|
|
host: Optional[str] = None,
|
|
ssh_user: Optional[str] = None,
|
|
ssh_key: Optional[str] = None,
|
|
docker: Optional[str] = None,
|
|
):
|
|
"""Parse command line arguments.
|
|
|
|
Note: This returns a list of hosts, not a comma separated string!
|
|
"""
|
|
if not host and not cluster:
|
|
bootstrap_config = os.path.expanduser("~/ray_bootstrap_config.yaml")
|
|
if os.path.exists(bootstrap_config):
|
|
cluster = bootstrap_config
|
|
cli_logger.warning(
|
|
f"Detected cluster config file at {cluster}. "
|
|
f"If this is incorrect, specify with "
|
|
f"`ray cluster-dump <config>`"
|
|
)
|
|
elif cluster:
|
|
cluster = os.path.expanduser(cluster)
|
|
|
|
cluster_name = None
|
|
|
|
if cluster:
|
|
h, u, k, d, cluster_name = get_info_from_ray_cluster_config(cluster)
|
|
|
|
ssh_user = ssh_user or u
|
|
ssh_key = ssh_key or k
|
|
docker = docker or d
|
|
hosts = host.split(",") if host else h
|
|
|
|
if not hosts:
|
|
raise LocalCommandFailed(
|
|
f"Invalid cluster file or cluster has no running nodes: " f"{cluster}"
|
|
)
|
|
elif host:
|
|
hosts = host.split(",")
|
|
else:
|
|
raise LocalCommandFailed(
|
|
"You need to either specify a `<cluster_config>` or `--host`."
|
|
)
|
|
|
|
if not ssh_user:
|
|
ssh_user = DEFAULT_SSH_USER
|
|
cli_logger.warning(
|
|
f"Using default SSH user `{ssh_user}`. "
|
|
f"If this is incorrect, specify with `--ssh-user <user>`"
|
|
)
|
|
|
|
if not ssh_key:
|
|
for cand_key in DEFAULT_SSH_KEYS:
|
|
cand_key_file = os.path.expanduser(cand_key)
|
|
if os.path.exists(cand_key_file):
|
|
ssh_key = cand_key_file
|
|
cli_logger.warning(
|
|
f"Auto detected SSH key file: {ssh_key}. "
|
|
f"If this is incorrect, specify with `--ssh-key <key>`"
|
|
)
|
|
break
|
|
|
|
return cluster, hosts, ssh_user, ssh_key, docker, cluster_name
|