def _task_fn(index, driver_addresses, settings):
task = task_service.SparkTaskService(index, settings.key)
try:
driver_client = driver_service.SparkDriverClient(driver_addresses, settings.key, settings.verbose)
driver_client.register_task(index, task.addresses(), host_hash.host_hash())
task.wait_for_initial_registration(settings.timeout)
# Tasks ping each other in a circular fashion to determine interfaces reachable within
# the cluster.
next_task_index = (index + 1) % settings.num_proc
next_task_addresses = driver_client.all_task_addresses(next_task_index)
# We request interface matching to weed out all the NAT'ed interfaces.
next_task_client = \
task_service.SparkTaskClient(next_task_index, next_task_addresses,
settings.key, settings.verbose,
match_intf=True)
driver_client.register_task_to_task_addresses(next_task_index, next_task_client.addresses())
task_indices_on_this_host = driver_client.task_host_hash_indices(
host_hash.host_hash())
if task_indices_on_this_host[0] == index:
# Task with first index will execute orted that will run mpirun_exec_fn for all tasks.
task.wait_for_command_start(settings.timeout)
task.wait_for_command_termination()
else:
# The rest of tasks need to wait for the first task to finish.
first_task_addresses = driver_client.all_task_addresses(task_indices_on_this_host[0])
first_task_client = \
task_service.SparkTaskClient(task_indices_on_this_host[0],
first_task_addresses, settings.key,
settings.verbose)
first_task_client.wait_for_command_termination()
return task.fn_result()
finally:
task.shutdown()
def rsh(driver_addresses, settings, host_hash, command, env, local_rank):
"""
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.
:param driver_addresses: driver's addresses
:param settings: settings
: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
"""
if ':' in host_hash:
raise Exception('Illegal host hash provided. Are you using Open MPI 4.0.0+?')
key = codec.loads_base64(env[secret.HOROVOD_SECRET_KEY])
driver_client = driver_service.SparkDriverClient(driver_addresses, key,
verbose=settings.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=settings.verbose)
task_client.run_command(command, env)
def task_exec(driver_addresses, settings, rank_env, local_rank_env):
# Die if parent process terminates
in_thread(target=_parent_process_monitor, args=(os.getppid(), ))
key = codec.loads_base64(os.environ[secret.HOROVOD_SECRET_KEY])
rank = int(os.environ[rank_env])
local_rank = int(os.environ[local_rank_env])
driver_client = driver_service.SparkDriverClient(driver_addresses,
key,
verbose=settings.verbose)
# tell driver about local rank and rank
# in elastic mode the driver already knows this mapping
# for simplicity we keep code paths the same for elastic and static mode
host_hash = os.environ['HOROVOD_HOSTNAME']
task_index = driver_client.set_local_rank_to_rank(host_hash, local_rank,
rank)
# gather available resources from task service
task_addresses = driver_client.all_task_addresses(task_index)
task_client = task_service.SparkTaskClient(task_index,
task_addresses,
key,
verbose=settings.verbose)
task_info.set_resources(task_client.resources())
fn, args, kwargs = driver_client.code()
result = fn(*args, **kwargs)
task_client.register_code_result(result)
def main(driver_addresses, settings, host_hash, command):
"""
Method to run `orted` remotely given a host hash and driver addresses.
This method connects to the SparkDriverService running on the Spark driver,
retrieves all information required to connect to the task with the lowest task index
of that host hash and invoke the command there.
All other tasks with the same host hash are expected to no-op (see `horovod.spark._task_fn`)
and wait for the first task to terminate.
:param driver_addresses: driver's addresses
:param settings: settings
:param host_hash: host hash to connect to
:param command: command and arguments to invoke
"""
if ':' in host_hash:
raise Exception(
'Illegal host hash provided. Are you using Open MPI 4.0.0+?')
key = codec.loads_base64(os.environ[secret.HOROVOD_SECRET_KEY])
driver_client = driver_service.SparkDriverClient(driver_addresses,
key,
verbose=settings.verbose)
task_indices = driver_client.task_host_hash_indices(host_hash)
# Since tasks with the same host hash have shared memory, we will run only
# one ORTED process on the first task.
first_task_index = task_indices[0]
task_addresses = driver_client.all_task_addresses(first_task_index)
task_client = task_service.SparkTaskClient(first_task_index,
task_addresses,
key,
verbose=settings.verbose)
task_client.run_command(command, os.environ)
def _task_fn(index, driver_addresses, key, settings, use_gloo):
# deserialized on Spark workers, settings do not contain the key, so it is given here explicitly
# Spark RPC communicates the key and supports encryption
# for convenience, we put it back into settings
settings.key = key
task = task_service.SparkTaskService(index, settings.key, settings.nics,
settings.verbose)
try:
driver_client = driver_service.SparkDriverClient(
driver_addresses, settings.key, settings.verbose)
driver_client.register_task(index, task.addresses(),
host_hash.host_hash())
task.wait_for_initial_registration(settings.timeout)
task_indices_on_this_host = driver_client.task_host_hash_indices(
host_hash.host_hash())
# With Gloo all tasks wait for the command
# With MPI task with first index executes orted which will run mpirun_exec_fn for all tasks.
if use_gloo or task_indices_on_this_host[0] == index:
task.wait_for_command_start(settings.timeout)
task.wait_for_command_termination()
else:
# The rest of tasks need to wait for the first task to finish.
first_task_addresses = driver_client.all_task_addresses(
task_indices_on_this_host[0])
first_task_client = \
task_service.SparkTaskClient(task_indices_on_this_host[0],
first_task_addresses, settings.key,
settings.verbose)
first_task_client.wait_for_command_termination()
return task.fn_result()
finally:
task.shutdown()
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 rsh(driver_addresses, key, settings, host_hash, command, env, local_rank,
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 settings: settings
: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 background: run command in background if True, returns command result otherwise
:param events: events to abort the command, only if background is True
"""
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=settings.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=settings.verbose)
task_client.run_command(command, env)
if not background:
stop = None
events = events or []
for event in events:
stop = threading.Event()
on_event(event, task_client.abort_command, stop=stop)
try:
return task_client.wait_for_command_exit_code()
except:
traceback.print_exc()
return -1
finally:
if stop is not None:
stop.set()
def main(driver_addresses, settings, host_hash, command):
if ':' in host_hash:
raise Exception('Illegal host hash provided. Are you using Open MPI 4.0.0+?')
key = codec.loads_base64(os.environ[secret.HOROVOD_SECRET_KEY])
driver_client = driver_service.SparkDriverClient(driver_addresses, key,
verbose=settings.verbose)
task_indices = driver_client.task_host_hash_indices(host_hash)
# Since tasks with the same host hash have shared memory, we will run only
# one ORTED process on the first task.
first_task_index = task_indices[0]
task_addresses = driver_client.all_task_addresses(first_task_index)
task_client = task_service.SparkTaskClient(first_task_index, task_addresses,
key, verbose=settings.verbose)
task_client.run_command(command, os.environ)
def main(driver_addresses):
# Die if parent process terminates
bg = threading.Thread(target=parent_process_monitor, args=(os.getppid(), ))
bg.daemon = True
bg.start()
key = codec.loads_base64(os.environ[secret.HOROVOD_SECRET_KEY])
rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
driver_client = driver_service.SparkDriverClient(driver_addresses, key)
task_index = driver_client.task_index_by_rank(rank)
task_addresses = driver_client.all_task_addresses(task_index)
task_client = task_service.SparkTaskClient(task_index, task_addresses, key)
fn, args, kwargs = driver_client.code()
result = fn(*args, **kwargs)
task_client.register_code_result(result)
def _task_fn(index, driver_addresses, key, settings, use_gloo):
# deserialized on Spark workers, settings do not contain the key, so it is given here explicitly
# Spark RPC communicates the key and supports encryption
# for convenience, we put it back into settings
settings.key = key
task = task_service.SparkTaskService(index, settings.key, settings.nics,
settings.verbose)
try:
driver_client = driver_service.SparkDriverClient(
driver_addresses, settings.key, settings.verbose)
driver_client.register_task(index, task.addresses(),
host_hash.host_hash())
task.wait_for_initial_registration(settings.timeout)
task_indices_on_this_host = driver_client.task_host_hash_indices(
host_hash.host_hash())
# With Gloo all tasks wait for the command
# With MPI task with first index executes orted which will run mpirun_exec_fn for all tasks.
minimum_lifetime_after_start = None
if use_gloo or task_indices_on_this_host[0] == index:
task.wait_for_command_start(settings.timeout)
minimum_lifetime_after_start = timeout.Timeout(
MINIMUM_COMMAND_LIFETIME_S, message='Just measuring runtime')
task.wait_for_command_termination()
else:
# The rest of tasks need to wait for the first task to finish.
first_task_addresses = driver_client.all_task_addresses(
task_indices_on_this_host[0])
first_task_client = \
task_service.SparkTaskClient(task_indices_on_this_host[0],
first_task_addresses, settings.key,
settings.verbose)
first_task_client.wait_for_command_termination()
# command terminated, make sure this task service does not shutdown too quickly after
# the client started the command as it needs some time to connect again
# to wait for the result after starting the command (see horovod.spark.driver.rsh).
if minimum_lifetime_after_start is not None:
time.sleep(minimum_lifetime_after_start.remaining())
return task.fn_result()
finally:
# this has to block on running requests (wait_for_command_exit_code)
# so they can finish serving the exit code
# shutdown does block with network.BasicService._server._block_on_close = True
task.shutdown()
def task_exec(driver_addresses, settings, rank_env):
# Die if parent process terminates
in_thread(target=_parent_process_monitor, args=(os.getppid(), ))
key = codec.loads_base64(os.environ[secret.HOROVOD_SECRET_KEY])
rank = int(os.environ[rank_env])
driver_client = driver_service.SparkDriverClient(driver_addresses,
key,
verbose=settings.verbose)
task_index = driver_client.task_index_by_rank(rank)
task_addresses = driver_client.all_task_addresses(task_index)
task_client = task_service.SparkTaskClient(task_index,
task_addresses,
key,
verbose=settings.verbose)
task_info.set_resources(task_client.resources())
fn, args, kwargs = driver_client.code()
result = fn(*args, **kwargs)
task_client.register_code_result(result)
def _task_fn(index, driver_addresses, key, settings, use_gloo, is_elastic):
# deserialized on Spark workers, settings do not contain the key, so it is given here explicitly
# Spark RPC communicates the key and supports encryption
# for convenience, we put it back into settings
settings.key = key
# to simplify things, each task is an individual host in Elastic Horovod on Spark
# further, each attempt (instance) of a task is an individual host in Elastic Horovod on Spark
# hides availability of shared memory among executors on the same Spark node
hosthash = host_hash(
salt='{}-{}'.format(index, time.time()) if is_elastic else None)
# provide host hash to mpirun_exec_fn.py via task service
# gloo_exec_fn.py will get this env var set in request env as well
os.environ['HOROVOD_HOSTNAME'] = hosthash
task = task_service.SparkTaskService(
index, settings.key, settings.nics,
MINIMUM_COMMAND_LIFETIME_S if is_elastic or use_gloo else None,
settings.verbose)
try:
driver_client = driver_service.SparkDriverClient(
driver_addresses, settings.key, settings.verbose)
driver_client.register_task(index, task.addresses(), hosthash)
if not is_elastic:
task.wait_for_initial_registration(settings.start_timeout)
task_indices_on_this_host = driver_client.task_host_hash_indices(
hosthash)
local_rank_zero_index = task_indices_on_this_host[0]
else:
local_rank_zero_index = None
# In elastic all tasks wait for task shutdown signal from driver.
# With Gloo all tasks wait for the command to start and terminate.
# With MPI task with first index executes orted which will run mpirun_exec_fn for all tasks.
if is_elastic:
# either terminate on task shutdown or command termination
shutdown_thread = in_thread(driver_client.wait_for_task_shutdown)
while shutdown_thread.is_alive():
# Once the command started we wait for its termination
if task.check_for_command_start(
WAIT_FOR_COMMAND_START_DELAY_SECONDS):
task.wait_for_command_termination()
if task.command_exit_code() != 0:
raise Exception(
'Command failed, making Spark task fail to restart the task'
)
break
# While no command started, we can shutdown any time
shutdown_thread.join(WAIT_FOR_SHUTDOWN_DELAY_SECONDS)
elif use_gloo or index == local_rank_zero_index:
# Either Gloo or first task with MPI.
task.wait_for_command_start(settings.start_timeout)
task.wait_for_command_termination()
else:
# The other tasks with MPI need to wait for the first task to finish.
first_task_addresses = driver_client.all_task_addresses(
local_rank_zero_index)
first_task_client = \
task_service.SparkTaskClient(local_rank_zero_index,
first_task_addresses, settings.key,
settings.verbose)
first_task_client.wait_for_command_termination()
return task.fn_result()
finally:
# we must not call into shutdown too quickly, task clients run a command
# and want to wait on the result, we have told task service not to return
# from wait_for_command_termination too quickly, so we are safe here to shutdown
# clients have had enough time to connect to the service already
#
# the shutdown has to block on running requests (wait_for_command_exit_code)
# so they can finish serving the exit code
# shutdown does block with network.BasicService._server._block_on_close = True
task.shutdown()
