chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,446 @@
|
||||
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
|
||||
Reference in New Issue
Block a user