def rsh(driver_addresses, key, host_hash, command, env, local_rank, verbose,
stdout=None, stderr=None, prefix_output_with_timestamp=False,
background=True, events=None):
"""
Method to run a command remotely given a host hash, local rank and driver addresses.
This method connects to the SparkDriverService running on the Spark driver,
retrieves all information required to connect to the task with given local rank
of that host hash and invoke the command there.
The method returns immediately after launching the command if background is True (default).
When background is set to False, this method waits for command termination and returns
command's result. If there is an exception while waiting for the result (i.e. connection reset)
it returns -1.
:param driver_addresses: driver's addresses
:param key: used for encryption of parameters passed across the hosts
:param host_hash: host hash to connect to
:param command: command and arguments to invoke
:param env: environment to use
:param local_rank: local rank on the host of task to run the command in
:param verbose: verbosity level
:param stdout: Task stdout is redirected to this stream.
:param stderr: Task stderr is redirected to this stream.
:param prefix_output_with_timestamp: shows timestamp in stdout/stderr forwarding on the driver if True
:param background: run command in background if True, returns command result otherwise
:param events: events to abort the command, only if background is True
:return exit code if background is False
"""
if ':' in host_hash:
raise Exception('Illegal host hash provided. Are you using Open MPI 4.0.0+?')
driver_client = driver_service.SparkDriverClient(driver_addresses, key, verbose=verbose)
task_indices = driver_client.task_host_hash_indices(host_hash)
task_index = task_indices[local_rank]
task_addresses = driver_client.all_task_addresses(task_index)
task_client = task_service.SparkTaskClient(task_index, task_addresses, key, verbose=verbose)
task_client.stream_command_output(stdout, stderr)
task_client.run_command(command, env,
capture_stdout=stdout is not None,
capture_stderr=stderr is not None,
prefix_output_with_timestamp=prefix_output_with_timestamp)
if not background:
events = events or []
stop = threading.Event()
for event in events:
on_event(event, task_client.abort_command, stop=stop)
try:
exit_code = task_client.wait_for_command_exit_code()
logging.debug('rsh exit code %s for host %s slot %s', exit_code, host_hash, local_rank)
return exit_code
except:
traceback.print_exc()
return -1
finally:
stop.set()
def _exec_middleman(command, env, exit_event, stdout, stderr, rw):
stdout_r, stdout_w = stdout
stderr_r, stderr_w = stderr
r, w = rw
# Close unused file descriptors to enforce PIPE behavior.
stdout_r.close()
stderr_r.close()
w.close()
os.setsid()
executor_shell = subprocess.Popen(command,
shell=True,
env=env,
stdout=stdout_w,
stderr=stderr_w)
# we don't bother stopping the on_event thread, this process sys.exits soon
# so the on_event thread has to be a deamon thread
on_event(exit_event,
terminate_executor_shell_and_children,
args=(executor_shell.pid, ),
daemon=True)
def kill_executor_children_if_parent_dies():
# This read blocks until the pipe is closed on the other side
# due to parent process termination (for any reason, including -9).
os.read(r.fileno(), 1)
terminate_executor_shell_and_children(executor_shell.pid)
in_thread(kill_executor_children_if_parent_dies)
exit_code = executor_shell.wait()
if exit_code < 0:
# See: https://www.gnu.org/software/bash/manual/html_node/Exit-Status.html
exit_code = 128 + abs(exit_code)
sys.exit(exit_code)
def test_on_event(self):
# a happy run without args and stop event
event = threading.Event()
fn = mock.Mock()
thread = on_event(event, fn)
fn.assert_not_called()
event.set()
thread.join(1.0)
self.assertFalse(thread.is_alive())
fn.assert_called_once()
# a happy run with args but without stop event
event = threading.Event()
fn = mock.Mock()
thread = on_event(event, fn, ('a', 1))
fn.assert_not_called()
event.set()
thread.join(1.0)
self.assertFalse(thread.is_alive())
fn.assert_called_once()
fn.assert_called_once_with('a', 1)
# a happy run with stop event but unused
event = threading.Event()
stop = threading.Event()
fn = mock.Mock()
thread = on_event(event, fn, stop=stop, check_stop_interval_s=0.01)
fn.assert_not_called()
event.set()
thread.join(1.0)
self.assertFalse(thread.is_alive())
fn.assert_called_once()
stop.set()
time.sleep(0.1)
fn.assert_called_once()
# stop the thread before we set the event
event = threading.Event()
stop = threading.Event()
fn = mock.Mock()
thread = on_event(event, fn, stop=stop, check_stop_interval_s=0.01)
fn.assert_not_called()
stop.set()
thread.join(1.0)
self.assertFalse(thread.is_alive())
fn.assert_not_called()
event.set()
time.sleep(0.1)
fn.assert_not_called()
# test with exception
def exception():
raise Exception("Test Exception")
event = threading.Event()
fn = mock.Mock(side_effect=exception)
thread = on_event(event, fn)
fn.assert_not_called()
event.set()
thread.join(1.0)
self.assertFalse(thread.is_alive())
fn.assert_called_once()
# test with exception but silent
event = threading.Event()
fn = mock.Mock(side_effect=exception)
thread = on_event(event, fn)
fn.assert_not_called()
event.set()
thread.join(1.0)
self.assertFalse(thread.is_alive())
fn.assert_called_once()
# test None event
event = None
fn = mock.Mock()
with pytest.raises(ValueError, match="^Event must not be None$"):
on_event(event, fn)
fn.assert_not_called()
# test non-tuple args
event = threading.Event()
fn = mock.Mock()
with pytest.raises(
ValueError,
match="^args must be a tuple, not <(class|type) 'int'>, "
"for a single argument use \\(arg,\\)$"):
on_event(event, fn, args=1)
fn.assert_not_called()
# test None stop and non-daemon
event = threading.Event()
fn = mock.Mock()
with pytest.raises(
ValueError,
match="^Stop event must be given for non-daemon event thread$"
):
on_event(event, fn, stop=None, daemon=False)
fn.assert_not_called()
def execute(command,
env=None,
stdout=None,
stderr=None,
index=None,
events=None,
prefix_output_with_timestamp=False):
ctx = multiprocessing.get_context('spawn')
# When this event is set, signal to middleman to terminate its children and exit.
exit_event = _create_event(ctx)
# Make a pipe for the subprocess stdout/stderr.
(stdout_r, stdout_w) = ctx.Pipe()
(stderr_r, stderr_w) = ctx.Pipe()
# This Pipe is how we ensure that the executed process is properly terminated (not orphaned) if
# the parent process is hard killed (-9). If the parent (this process) is killed for any reason,
# this Pipe will be closed, which can be detected by the middleman. When the middleman sees the
# closed Pipe, it will issue a SIGTERM to the subprocess executing the command. The assumption
# here is that users will be inclined to hard kill this process, not the middleman.
(r, w) = ctx.Pipe()
middleman = ctx.Process(target=_exec_middleman,
args=(command, env, exit_event,
(stdout_r, stdout_w), (stderr_r,
stderr_w), (r, w)))
middleman.start()
# Close unused file descriptors to enforce PIPE behavior.
r.close()
stdout_w.close()
stderr_w.close()
# Redirect command stdout & stderr to provided streams or sys.stdout/sys.stderr.
# This is useful for Jupyter Notebook that uses custom sys.stdout/sys.stderr or
# for redirecting to a file on disk.
if stdout is None:
stdout = sys.stdout
if stderr is None:
stderr = sys.stderr
stdout_fwd = in_thread(target=forward_stream,
args=(stdout_r, stdout, 'stdout', index,
prefix_output_with_timestamp))
stderr_fwd = in_thread(target=forward_stream,
args=(stderr_r, stderr, 'stderr', index,
prefix_output_with_timestamp))
# TODO: Currently this requires explicitly declaration of the events and signal handler to set
# the event (gloo_run.py:_launch_jobs()). Need to figure out a generalized way to hide this behind
# interfaces.
stop = threading.Event()
events = events or []
for event in events:
on_event(event, exit_event.set, stop=stop, silent=True)
try:
middleman.join()
except:
# interrupted, send middleman TERM signal which will terminate children
exit_event.set()
while True:
try:
middleman.join()
break
except:
# interrupted, wait for middleman to finish
pass
finally:
stop.set()
stdout_fwd.join()
stderr_fwd.join()
return middleman.exitcode