wehub-resource-sync
447 lines
15 KiB
Python
447 lines
15 KiB
Python
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from typing import IO, Any, List, Optional
|
|
|
|
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
|
|
|
CONN_REFUSED_PATIENCE = 30 # how long to wait for sshd to run
|
|
|
|
_redirect_output = False # Whether to log command output to a temporary file
|
|
_allow_interactive = True # whether to pass on stdin to running commands.
|
|
|
|
|
|
def is_output_redirected():
|
|
return _redirect_output
|
|
|
|
|
|
def set_output_redirected(val: bool):
|
|
"""Choose between logging to a temporary file and to `sys.stdout`.
|
|
|
|
The default is to log to a file.
|
|
|
|
Args:
|
|
val: If true, subprocess output will be redirected to
|
|
a temporary file.
|
|
"""
|
|
global _redirect_output
|
|
_redirect_output = val
|
|
|
|
|
|
def does_allow_interactive():
|
|
return _allow_interactive
|
|
|
|
|
|
def set_allow_interactive(val: bool):
|
|
"""Choose whether to pass on stdin to running commands.
|
|
|
|
The default is to pipe stdin and close it immediately.
|
|
|
|
Args:
|
|
val: If true, stdin will be passed to commands.
|
|
"""
|
|
global _allow_interactive
|
|
_allow_interactive = val
|
|
|
|
|
|
class ProcessRunnerError(Exception):
|
|
def __init__(self, msg, msg_type, code=None, command=None, special_case=None):
|
|
super(ProcessRunnerError, self).__init__(
|
|
"{} (discovered={}): type={}, code={}, command={}".format(
|
|
msg, special_case, msg_type, code, command
|
|
)
|
|
)
|
|
|
|
self.msg_type = msg_type
|
|
self.code = code
|
|
self.command = command
|
|
|
|
self.special_case = special_case
|
|
|
|
|
|
_ssh_output_regexes = {
|
|
"known_host_update": re.compile(
|
|
r"\s*Warning: Permanently added '.+' \(.+\) " r"to the list of known hosts.\s*"
|
|
),
|
|
"connection_closed": re.compile(r"\s*Shared connection to .+ closed.\s*"),
|
|
"timeout": re.compile(
|
|
r"\s*ssh: connect to host .+ port .+: " r"Operation timed out\s*"
|
|
),
|
|
"conn_refused": re.compile(
|
|
r"\s*ssh: connect to host .+ port .+: Connection refused\s*"
|
|
)
|
|
# todo: check for other connection failures for better error messages?
|
|
}
|
|
|
|
|
|
def _read_subprocess_stream(
|
|
f: IO[str], output_file: Optional[IO[str]], is_stdout: bool = False
|
|
) -> Optional[str]:
|
|
"""Read and process a subprocess output stream.
|
|
|
|
The goal is to find error messages and respond to them in a clever way.
|
|
Currently just used for SSH messages (CONN_REFUSED, TIMEOUT, etc.), so
|
|
the user does not get confused by these.
|
|
|
|
Ran in a thread each for both `stdout` and `stderr` to
|
|
allow for cross-platform asynchronous IO.
|
|
|
|
Note: `select`-based IO is another option, but Windows has
|
|
no support for `select`ing pipes, and Linux support varies somewhat.
|
|
Spefically, Older *nix systems might also have quirks in how they
|
|
handle `select` on pipes.
|
|
|
|
Args:
|
|
f: File object for the stream.
|
|
output_file: File object to which filtered output is written.
|
|
is_stdout:
|
|
When `is_stdout` is `False`, the stream is assumed to
|
|
be `stderr`. Different error message detectors are used,
|
|
and the output is displayed to the user unless it matches
|
|
a special case (e.g. SSH timeout), in which case this is
|
|
left up to the caller.
|
|
|
|
Returns:
|
|
The detected special case name (e.g. "ssh_timeout") or None.
|
|
"""
|
|
|
|
detected_special_case = None
|
|
while True:
|
|
# ! Readline here is crucial.
|
|
# ! Normal `read()` will block until EOF instead of until
|
|
# something is available.
|
|
line = f.readline()
|
|
|
|
if line is None or line == "":
|
|
# EOF
|
|
break
|
|
|
|
if line[-1] == "\n":
|
|
line = line[:-1]
|
|
|
|
if not is_stdout:
|
|
if _ssh_output_regexes["connection_closed"].fullmatch(line) is not None:
|
|
# Do not log "connection closed" messages which SSH
|
|
# puts in stderr for no reason.
|
|
#
|
|
# They are never errors since the connection will
|
|
# close no matter whether the command succeeds or not.
|
|
continue
|
|
|
|
if _ssh_output_regexes["timeout"].fullmatch(line) is not None:
|
|
# Timeout is not really an error but rather a special
|
|
# condition. It should be handled by the caller, since
|
|
# network conditions/nodes in the early stages of boot
|
|
# are expected to sometimes cause connection timeouts.
|
|
if detected_special_case is not None:
|
|
raise ValueError(
|
|
"Bug: ssh_timeout conflicts with another "
|
|
"special codition: " + detected_special_case
|
|
)
|
|
|
|
detected_special_case = "ssh_timeout"
|
|
continue
|
|
|
|
if _ssh_output_regexes["conn_refused"].fullmatch(line) is not None:
|
|
# Connection refused is not really an error but
|
|
# rather a special condition. It should be handled by
|
|
# the caller, since network conditions/nodes in the
|
|
# early stages of boot are expected to sometimes cause
|
|
# CONN_REFUSED.
|
|
if detected_special_case is not None:
|
|
raise ValueError(
|
|
"Bug: ssh_conn_refused conflicts with another "
|
|
"special codition: " + detected_special_case
|
|
)
|
|
|
|
detected_special_case = "ssh_conn_refused"
|
|
continue
|
|
|
|
if _ssh_output_regexes["known_host_update"].fullmatch(line) is not None:
|
|
# Since we ignore SSH host control anyway
|
|
# (-o UserKnownHostsFile=/dev/null),
|
|
# we should silence the host control warnings.
|
|
continue
|
|
|
|
cli_logger.error(line)
|
|
|
|
if output_file is not None and output_file != subprocess.DEVNULL:
|
|
output_file.write(line + "\n")
|
|
|
|
return detected_special_case
|
|
|
|
|
|
def _run_and_process_output(
|
|
cmd: List[str],
|
|
stdout_file: Optional[IO[str]],
|
|
process_runner: Any = subprocess,
|
|
stderr_file: Optional[IO[str]] = None,
|
|
use_login_shells: bool = False,
|
|
):
|
|
"""Run a command and process its output for special cases.
|
|
|
|
Calls a standard 'check_call' if process_runner is not subprocess.
|
|
|
|
Specifically, run all command output through regex to detect
|
|
error conditions and filter out non-error messages that went to stderr
|
|
anyway (SSH writes ALL of its "system" messages to stderr even if they
|
|
are not actually errors).
|
|
|
|
Args:
|
|
cmd: Command to run.
|
|
stdout_file: File to redirect stdout to.
|
|
process_runner: Used for command execution. Assumed to have
|
|
'check_call' and 'check_output' inplemented.
|
|
stderr_file: File to redirect stderr to.
|
|
use_login_shells: Whether to disable special output processing because
|
|
an interactive login shell is in use.
|
|
|
|
Returns:
|
|
The return code of the executed process.
|
|
|
|
Implementation notes:
|
|
1. `use_login_shells` disables special processing
|
|
If we run interactive apps, output processing will likely get
|
|
overwhelmed with the interactive output elements.
|
|
Thus, we disable output processing for login shells. This makes
|
|
the logging experience considerably worse, but it only degrades
|
|
to old-style logging.
|
|
|
|
For example, `pip install` outputs HUNDREDS of progress-bar lines
|
|
when downloading a package, and we have to
|
|
read + regex + write all of them.
|
|
|
|
After all, even just printing output to console can often slow
|
|
down a fast-printing app, and we do more than just print, and
|
|
all that from Python, which is much slower than C regarding
|
|
stream processing.
|
|
|
|
2. `stdin=PIPE` for subprocesses
|
|
Do not inherit stdin as it messes with bash signals
|
|
(ctrl-C for SIGINT) and these commands aren't supposed to
|
|
take input anyway.
|
|
|
|
3. `ThreadPoolExecutor` without the `Pool`
|
|
We use `ThreadPoolExecutor` to create futures from threads.
|
|
Threads are never reused.
|
|
|
|
This approach allows us to have no custom synchronization by
|
|
off-loading the return value and exception passing to the
|
|
standard library (`ThreadPoolExecutor` internals).
|
|
|
|
This instance will be `shutdown()` ASAP so it's fine to
|
|
create one in such a weird place.
|
|
|
|
The code is thus 100% thread-safe as long as the stream readers
|
|
are read-only except for return values and possible exceptions.
|
|
"""
|
|
stdin_overwrite = subprocess.PIPE
|
|
# This already should be validated in a higher place of the stack.
|
|
assert not (
|
|
does_allow_interactive() and is_output_redirected()
|
|
), "Cannot redirect output while in interactive mode."
|
|
if process_runner != subprocess or (
|
|
does_allow_interactive() and not is_output_redirected()
|
|
):
|
|
stdin_overwrite = None
|
|
|
|
# See implementation note #1
|
|
|
|
if use_login_shells or process_runner != subprocess:
|
|
return process_runner.check_call(
|
|
cmd,
|
|
# See implementation note #2
|
|
stdin=stdin_overwrite,
|
|
stdout=stdout_file,
|
|
stderr=stderr_file,
|
|
)
|
|
|
|
with subprocess.Popen(
|
|
cmd,
|
|
# See implementation note #2
|
|
stdin=stdin_overwrite,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
bufsize=1, # line buffering
|
|
universal_newlines=True, # text mode outputs
|
|
) as p:
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
# Closing stdin might be necessary to signal EOF to some
|
|
# apps (they might get stuck waiting for input forever otherwise).
|
|
p.stdin.close()
|
|
|
|
# See implementation note #3
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
stdout_future = executor.submit(
|
|
_read_subprocess_stream, p.stdout, stdout_file, is_stdout=True
|
|
)
|
|
stderr_future = executor.submit(
|
|
_read_subprocess_stream, p.stderr, stderr_file, is_stdout=False
|
|
)
|
|
# Wait for completion.
|
|
executor.shutdown()
|
|
|
|
# Update `p.returncode`
|
|
p.poll()
|
|
|
|
detected_special_case = stdout_future.result()
|
|
if stderr_future.result() is not None:
|
|
if detected_special_case is not None:
|
|
# This might some day need to be changed.
|
|
# We should probably make sure the two special cases
|
|
# are compatible then and that we can handle both by
|
|
# e.g. reporting both to the caller.
|
|
raise ValueError(
|
|
"Bug: found a special case in both stdout and "
|
|
"stderr. This is not valid behavior at the time "
|
|
"of writing this code."
|
|
)
|
|
detected_special_case = stderr_future.result()
|
|
|
|
if p.returncode > 0:
|
|
# Process failed, but not due to a signal, since signals
|
|
# set the exit code to a negative value.
|
|
raise ProcessRunnerError(
|
|
"Command failed",
|
|
"ssh_command_failed",
|
|
code=p.returncode,
|
|
command=cmd,
|
|
special_case=detected_special_case,
|
|
)
|
|
elif p.returncode < 0:
|
|
# Process failed due to a signal, since signals
|
|
# set the exit code to a negative value.
|
|
raise ProcessRunnerError(
|
|
"Command failed",
|
|
"ssh_command_failed",
|
|
code=p.returncode,
|
|
command=cmd,
|
|
special_case="died_to_signal",
|
|
)
|
|
|
|
return p.returncode
|
|
|
|
|
|
def run_cmd_redirected(
|
|
cmd: List[str],
|
|
process_runner: Any = subprocess,
|
|
silent: bool = False,
|
|
use_login_shells: bool = False,
|
|
):
|
|
"""Run a command and optionally redirect output to a file.
|
|
|
|
Args:
|
|
cmd: Command to run.
|
|
process_runner: Process runner used for executing commands.
|
|
silent: If true, the command output will be silenced completely
|
|
(redirected to /dev/null), unless verbose logging
|
|
is enabled. Use this for running utility commands like
|
|
rsync.
|
|
use_login_shells: Whether to disable special output processing because
|
|
an interactive login shell is in use.
|
|
|
|
Returns:
|
|
The return code of the executed process.
|
|
"""
|
|
if silent and cli_logger.verbosity < 1:
|
|
return _run_and_process_output(
|
|
cmd,
|
|
process_runner=process_runner,
|
|
stdout_file=process_runner.DEVNULL,
|
|
stderr_file=process_runner.DEVNULL,
|
|
use_login_shells=use_login_shells,
|
|
)
|
|
|
|
if not is_output_redirected():
|
|
return _run_and_process_output(
|
|
cmd,
|
|
process_runner=process_runner,
|
|
stdout_file=sys.stdout,
|
|
stderr_file=sys.stderr,
|
|
use_login_shells=use_login_shells,
|
|
)
|
|
else:
|
|
tmpfile_path = os.path.join(
|
|
tempfile.gettempdir(), "ray-up-{}-{}.txt".format(cmd[0], time.time())
|
|
)
|
|
with open(
|
|
tmpfile_path,
|
|
mode="w",
|
|
# line buffering
|
|
buffering=1,
|
|
) as tmp:
|
|
cli_logger.verbose("Command stdout is redirected to {}", cf.bold(tmp.name))
|
|
|
|
return _run_and_process_output(
|
|
cmd,
|
|
process_runner=process_runner,
|
|
stdout_file=tmp,
|
|
stderr_file=tmp,
|
|
use_login_shells=use_login_shells,
|
|
)
|
|
|
|
|
|
def handle_ssh_fails(
|
|
e: ProcessRunnerError,
|
|
first_conn_refused_time: Optional[float],
|
|
retry_interval: float,
|
|
) -> Optional[float]:
|
|
"""Handle SSH system failures coming from a subprocess.
|
|
|
|
Args:
|
|
e: The `ProcessRunnerException` to handle.
|
|
first_conn_refused_time:
|
|
The time (as reported by this function) or None,
|
|
indicating the last time a CONN_REFUSED error was caught.
|
|
|
|
After exceeding a patience value, the program will be aborted
|
|
since SSH will likely never recover.
|
|
retry_interval: The interval after which the command will be retried,
|
|
used here just to inform the user.
|
|
|
|
Returns:
|
|
The updated `first_conn_refused_time`, or None if not applicable.
|
|
"""
|
|
if e.msg_type != "ssh_command_failed":
|
|
return
|
|
|
|
if e.special_case == "ssh_conn_refused":
|
|
if (
|
|
first_conn_refused_time is not None
|
|
and time.time() - first_conn_refused_time > CONN_REFUSED_PATIENCE
|
|
):
|
|
cli_logger.error(
|
|
"SSH connection was being refused "
|
|
"for {} seconds. Head node assumed "
|
|
"unreachable.",
|
|
cf.bold(str(CONN_REFUSED_PATIENCE)),
|
|
)
|
|
cli_logger.abort(
|
|
"Check the node's firewall settings "
|
|
"and the cloud network configuration."
|
|
)
|
|
|
|
cli_logger.warning("SSH connection was refused.")
|
|
cli_logger.warning(
|
|
"This might mean that the SSH daemon is "
|
|
"still setting up, or that "
|
|
"the host is inaccessible (e.g. due to "
|
|
"a firewall)."
|
|
)
|
|
|
|
return time.time()
|
|
|
|
if e.special_case in ["ssh_timeout", "ssh_conn_refused"]:
|
|
cli_logger.print(
|
|
"SSH still not available, retrying in {} seconds.",
|
|
cf.bold(str(retry_interval)),
|
|
)
|
|
else:
|
|
raise e
|
|
|
|
return first_conn_refused_time
|