def slot_info_to_command(slot_info): """ Given a slot_info, creates a command used by gloo to launch a single job. :param slot_info: host and slot to execute the run command on :return: """ host_name = slot_info.hostname horovod_rendez_env = ( 'HOROVOD_HOSTNAME={hostname} ' 'HOROVOD_RANK={rank} ' 'HOROVOD_SIZE={size} ' 'HOROVOD_LOCAL_RANK={local_rank} ' 'HOROVOD_LOCAL_SIZE={local_size} ' 'HOROVOD_CROSS_RANK={cross_rank} ' 'HOROVOD_CROSS_SIZE={cross_size} ' .format(hostname=host_name, rank=slot_info.rank, size=slot_info.size, local_rank=slot_info.local_rank, local_size=slot_info.local_size, cross_rank=slot_info.cross_rank, cross_size=slot_info.cross_size)) return '{horovod_env} {env} {run_command}' .format( horovod_env=horovod_rendez_env, env=' '.join(['%s=%s' % (key, quote(value)) for key, value in env.items() if env_util.is_exportable(key)]), run_command=run_command)
def mpi_run(settings, common_intfs, env): mpi_impl_flags = _get_mpi_implementation_flags() if mpi_impl_flags is None: raise Exception( 'horovodrun convenience script does not find an installed OpenMPI.\n\n' 'Choose one of:\n' '1. Install Open MPI 4.0.0+ or IBM Spectrum MPI and re-install Horovod ' '(use --no-cache-dir pip option).\n' '2. Run distributed ' 'training script using the standard way provided by your' ' MPI distribution (usually mpirun, srun, or jsrun).\n' '3. Use built-in gloo option (horovodrun --gloo ...).') ssh_port_arg = '-mca plm_rsh_args \"-p {ssh_port}\"'.format( ssh_port=settings.ssh_port) if settings.ssh_port else '' # if user does not specify any hosts, mpirun by default uses local host. # There is no need to specify localhost. hosts_arg = '-H {hosts}'.format(hosts=settings.hosts) tcp_intf_arg = '-mca btl_tcp_if_include {common_intfs}'.format( common_intfs=','.join(common_intfs)) if common_intfs else '' nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={common_intfs}'.format( common_intfs=','.join(common_intfs)) if common_intfs else '' # On large cluster runs (e.g. Summit), we need extra settings to work around OpenMPI issues if settings.num_hosts >= 64: mpi_impl_flags.append('-mca plm_rsh_no_tree_spawn true') mpi_impl_flags.append('-mca plm_rsh_num_concurrent {}'.format( settings.num_proc)) # Pass all the env variables to the mpirun command. mpirun_command = ( 'mpirun --allow-run-as-root --tag-output ' '-np {num_proc} {hosts_arg} ' '-bind-to none -map-by slot ' '{mpi_args} ' '{ssh_port_arg} ' '{tcp_intf_arg} ' '{nccl_socket_intf_arg} ' '{output_filename_arg} ' '{env} {extra_mpi_args} {command}' # expect a lot of environment variables .format(num_proc=settings.num_proc, hosts_arg=hosts_arg, mpi_args=' '.join(mpi_impl_flags), tcp_intf_arg=tcp_intf_arg, nccl_socket_intf_arg=nccl_socket_intf_arg, ssh_port_arg=ssh_port_arg, output_filename_arg='--output-filename ' + settings.output_filename if settings.output_filename else '', env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)), extra_mpi_args=settings.extra_mpi_args if settings.extra_mpi_args else '', command=' '.join(quote(par) for par in settings.command))) if settings.verbose >= 2: print(mpirun_command) # Execute the mpirun command. os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)
def mpi_run(settings, common_intfs): if not _is_open_mpi_installed(): raise Exception( 'horovodrun convenience script does not find an installed OpenMPI.\n\n' 'Choose one of:\n' '1. Install Open MPI 4.0.0+ and re-install Horovod ' '(use --no-cache-dir pip option).\n' '2. Run distributed ' 'training script using the standard way provided by your' ' MPI distribution (usually mpirun, srun, or jsrun).\n' '3. Use built-in gloo option (horovodrun --gloo ...).') # Pass all the env variables to the mpirun command. env = os.environ.copy() ssh_port_arg = '-mca plm_rsh_args \"-p {ssh_port}\"'.format( ssh_port=settings.ssh_port) if settings.ssh_port else '' # if user does not specify any hosts, mpirun by default uses local host. # There is no need to specify localhost. hosts_arg = '-H {hosts}'.format(hosts=settings.hosts) tcp_intf_arg = '-mca btl_tcp_if_include {common_intfs}'.format( common_intfs=','.join(common_intfs)) if common_intfs else '' nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={common_intfs}'.format( common_intfs=','.join(common_intfs)) if common_intfs else '' mpirun_command = ( 'mpirun --allow-run-as-root --tag-output ' '-np {num_proc} {hosts_arg} ' '-bind-to none -map-by slot ' '-mca pml ob1 -mca btl ^openib ' '{ssh_port_arg} ' '{tcp_intf_arg} ' '{nccl_socket_intf_arg} ' '{env} {command}' # expect a lot of environment variables .format(num_proc=settings.num_proc, hosts_arg=hosts_arg, tcp_intf_arg=tcp_intf_arg, nccl_socket_intf_arg=nccl_socket_intf_arg, ssh_port_arg=ssh_port_arg, env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)), command=' '.join(quote(par) for par in settings.command)) ) if settings.verbose >= 2: print(mpirun_command) # Execute the mpirun command. os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)
def alloc_info_to_command(alloc_info): """ Given an alloc_info, creates a command used by gloo to launch a single job. :param alloc_info: host and slot to execute the run command on :return: """ # generate env for rendezvous horovod_rendez_env = 'HOROVOD_RANK={rank} HOROVOD_SIZE={size} ' \ 'HOROVOD_LOCAL_RANK={local_rank} HOROVOD_LOCAL_SIZE={local_size} ' \ 'HOROVOD_CROSS_RANK={cross_rank} HOROVOD_CROSS_SIZE={cross_size} ' \ .format(rank=alloc_info.rank, size=alloc_info.size, local_rank=alloc_info.local_rank, local_size=alloc_info.local_size, cross_rank=alloc_info.cross_rank, cross_size=alloc_info.cross_size) return '{horovod_env} {env} {run_command}' .format( horovod_env=horovod_rendez_env, env=' '.join(['%s=%s' % (key, quote(value)) for key, value in env.items() if env_util.is_exportable(key)]), run_command=run_command)
def run(fn, args=(), kwargs={}, num_proc=None, start_timeout=None, env=None, stdout=None, stderr=None, verbose=1): """ Runs Horovod in Spark. Runs `num_proc` processes executing `fn` using the same amount of Spark tasks. Args: fn: Function to run. args: Arguments to pass to `fn`. kwargs: Keyword arguments to pass to `fn`. num_proc: Number of Horovod processes. Defaults to `spark.default.parallelism`. start_timeout: Timeout for Spark tasks to spawn, register and start running the code, in seconds. If not set, falls back to `HOROVOD_SPARK_START_TIMEOUT` environment variable value. If it is not set as well, defaults to 600 seconds. env: Environment dictionary to use in Horovod run. Defaults to `os.environ`. stdout: Horovod stdout is redirected to this stream. Defaults to sys.stdout. stderr: Horovod stderr is redirected to this stream. Defaults to sys.stderr. verbose: Debug output verbosity (0-2). Defaults to 1. Returns: List of results returned by running `fn` on each rank. """ if start_timeout is None: # Lookup default timeout from the environment variable. start_timeout = int(os.getenv('HOROVOD_SPARK_START_TIMEOUT', '600')) tmout = timeout.Timeout(start_timeout, message='Timed out waiting for {activity}. Please check that you have ' 'enough resources to run all Horovod processes. Each Horovod ' 'process runs in a Spark task. You may need to increase the ' 'start_timeout parameter to a larger value if your Spark resources ' 'are allocated on-demand.') settings = hvd_settings.Settings(verbose=verbose, key=secret.make_secret_key(), timeout=tmout) spark_context = pyspark.SparkContext._active_spark_context if spark_context is None: raise Exception('Could not find an active SparkContext, are you ' 'running in a PySpark session?') if num_proc is None: num_proc = spark_context.defaultParallelism if settings.verbose >= 1: print('Running %d processes (inferred from spark.default.parallelism)...' % num_proc) else: if settings.verbose >= 1: print('Running %d processes...' % num_proc) settings.num_proc = num_proc result_queue = queue.Queue(1) spark_job_group = 'horovod.spark.run.%d' % job_id.next_job_id() driver = driver_service.SparkDriverService(settings.num_proc, fn, args, kwargs, settings.key) spark_thread = _make_spark_thread(spark_context, spark_job_group, driver, result_queue, settings) try: driver.wait_for_initial_registration(settings.timeout) if settings.verbose >= 2: print('Initial Spark task registration is complete.') task_clients = [ task_service.SparkTaskClient(index, driver.task_addresses_for_driver(index), settings.key, settings.verbose) for index in range(settings.num_proc)] for task_client in task_clients: task_client.notify_initial_registration_complete() driver.wait_for_task_to_task_address_updates(settings.timeout) if settings.verbose >= 2: print('Spark task-to-task address registration is complete.') # Determine a set of common interfaces for task-to-task communication. common_intfs = set(driver.task_addresses_for_tasks(0).keys()) for index in range(1, settings.num_proc): common_intfs.intersection_update(driver.task_addresses_for_tasks(index).keys()) if not common_intfs: raise Exception('Unable to find a set of common task-to-task communication interfaces: %s' % [(index, driver.task_addresses_for_tasks(index)) for index in range(settings.num_proc)]) # Determine the index grouping based on host hashes. # Barrel shift until index 0 is in the first host. host_hashes = list(driver.task_host_hash_indices().keys()) host_hashes.sort() while 0 not in driver.task_host_hash_indices()[host_hashes[0]]: host_hashes = host_hashes[1:] + host_hashes[:1] ranks_to_indices = [] for host_hash in host_hashes: ranks_to_indices += driver.task_host_hash_indices()[host_hash] driver.set_ranks_to_indices(ranks_to_indices) if env is None: env = os.environ.copy() # Pass secret key through the environment variables. env[secret.HOROVOD_SECRET_KEY] = codec.dumps_base64(settings.key) mpirun_command = ( 'mpirun --allow-run-as-root --tag-output ' '-np {num_proc} -H {hosts} ' '-bind-to none -map-by slot ' '-mca pml ob1 -mca btl ^openib -mca btl_tcp_if_include {common_intfs} ' '-x NCCL_DEBUG=INFO -x NCCL_SOCKET_IFNAME={common_intfs} ' '{env} ' # expect a lot of environment variables '-mca plm_rsh_agent "{python} -m horovod.spark.driver.mpirun_rsh {encoded_driver_addresses} {settings}" ' '{python} -m horovod.spark.task.mpirun_exec_fn {encoded_driver_addresses} {settings}' .format(num_proc=settings.num_proc, hosts=','.join('%s:%d' % (host_hash, len(driver.task_host_hash_indices()[host_hash])) for host_hash in host_hashes), common_intfs=','.join(common_intfs), env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)), python=sys.executable, encoded_driver_addresses=codec.dumps_base64(driver.addresses()), settings=codec.dumps_base64(settings))) if settings.verbose >= 2: print('+ %s' % mpirun_command) exit_code = safe_shell_exec.execute(mpirun_command, env, stdout, stderr) if exit_code != 0: raise Exception('mpirun exited with code %d, see the error above.' % exit_code) except: # Terminate Spark job. spark_context.cancelJobGroup(spark_job_group) # Re-raise exception. raise finally: spark_thread.join() driver.shutdown() # Make sure Spark Job did not fail. driver.check_for_spark_job_failure() # If there's no exception, execution results are in this queue. results = result_queue.get_nowait() return [results[index] for index in ranks_to_indices]
def mpi_run(settings, common_intfs, env, command, stdout=None, stderr=None, run_func=safe_shell_exec.execute): """ Runs mpi_run. Args: settings: Settings for running MPI. Note: settings.num_proc and settings.hosts must not be None. common_intfs: Interfaces to include by MPI. env: Environment dictionary to use for running MPI. command: Command and arguments to run as a list of string. stdout: Stdout of the mpi process. Only used when settings.run_func_mode is True. stderr: Stderr of the mpi process. Only used when settings.run_func_mode is True. run_func: Run function to use. Must have arguments 'command' and 'env'. Only used when settings.run_func_mode is True. Defaults to safe_shell_exec.execute. """ mpi_impl_flags = _get_mpi_implementation_flags() if mpi_impl_flags is None: raise Exception( 'horovodrun convenience script does not find an installed MPI.\n\n' 'Choose one of:\n' '1. Install Open MPI 4.0.0+ or IBM Spectrum MPI or MPICH and re-install Horovod ' '(use --no-cache-dir pip option).\n' '2. Run distributed ' 'training script using the standard way provided by your' ' MPI distribution (usually mpirun, srun, or jsrun).\n' '3. Use built-in gloo option (horovodrun --gloo ...).') ssh_port_arg = '-mca plm_rsh_args \"-p {ssh_port}\"'.format( ssh_port=settings.ssh_port) if settings.ssh_port else '' # if user does not specify any hosts, mpirun by default uses local host. # There is no need to specify localhost. hosts_arg = '-H {hosts}'.format(hosts=settings.hosts) tcp_intf_arg = '-mca btl_tcp_if_include {common_intfs}'.format( common_intfs=','.join(common_intfs)) if common_intfs else '' nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={common_intfs}'.format( common_intfs=','.join(common_intfs)) if common_intfs else '' # On large cluster runs (e.g. Summit), we need extra settings to work around OpenMPI issues if settings.num_hosts and settings.num_hosts >= _LARGE_CLUSTER_THRESHOLD: mpi_impl_flags.append('-mca plm_rsh_no_tree_spawn true') mpi_impl_flags.append('-mca plm_rsh_num_concurrent {}'.format( settings.num_proc)) # Pass all the env variables to the mpirun command. mpirun_command = ( 'mpirun --allow-run-as-root --tag-output ' '-np {num_proc} {hosts_arg} ' '-bind-to none -map-by slot ' '{mpi_args} ' '{ssh_port_arg} ' '{tcp_intf_arg} ' '{nccl_socket_intf_arg} ' '{output_filename_arg} ' '{env} {extra_mpi_args} {command}' # expect a lot of environment variables .format(num_proc=settings.num_proc, hosts_arg=hosts_arg, mpi_args=' '.join(mpi_impl_flags), tcp_intf_arg=tcp_intf_arg, nccl_socket_intf_arg=nccl_socket_intf_arg, ssh_port_arg=ssh_port_arg, output_filename_arg='--output-filename ' + settings.output_filename if settings.output_filename else '', env=' '.join('-x %s' % key for key in sorted(env.keys()) if env_util.is_exportable(key)), extra_mpi_args=settings.extra_mpi_args if settings.extra_mpi_args else '', command=' '.join(quote(par) for par in command))) if settings.verbose >= 2: print(mpirun_command) # Execute the mpirun command. if settings.run_func_mode: exit_code = run_func(command=mpirun_command, env=env, stdout=stdout, stderr=stderr) if exit_code != 0: raise RuntimeError( "mpirun failed with exit code {exit_code}".format( exit_code=exit_code)) else: os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)
def mpi_run(settings, nics, env, command, stdout=None, stderr=None): """ Runs mpi_run. Args: settings: Settings for running MPI. Note: settings.num_proc and settings.hosts must not be None. nics: Interfaces to include by MPI. env: Environment dictionary to use for running command. command: Command and arguments to run as a list of string. stdout: Stdout of the mpi process. Only used when settings.run_func_mode is True. stderr: Stderr of the mpi process. Only used when settings.run_func_mode is True. """ if env is not None and not isinstance(env, dict): raise Exception( 'env argument must be a dict, not {type}: {env}'.format( type=type(env), env=env)) mpi_impl_flags, impl_binding_args = _get_mpi_implementation_flags( settings.tcp_flag, env=env) if mpi_impl_flags is None: raise Exception(_MPI_NOT_FOUND_ERROR_MSG) ssh_port_arg = '-mca plm_rsh_args \"-p {ssh_port}\"'.format( ssh_port=settings.ssh_port) if settings.ssh_port else '' # if user does not specify any hosts, mpirun by default uses local host. # There is no need to specify localhost. hosts_arg = '-H {hosts}'.format(hosts=settings.hosts) tcp_intf_arg = '-mca btl_tcp_if_include {nics}'.format( nics=','.join(nics)) if nics else '' nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={nics}'.format( nics=','.join(nics)) if nics else '' # On large cluster runs (e.g. Summit), we need extra settings to work around OpenMPI issues if settings.num_hosts and settings.num_hosts >= _LARGE_CLUSTER_THRESHOLD: mpi_impl_flags.append('-mca plm_rsh_no_tree_spawn true') mpi_impl_flags.append('-mca plm_rsh_num_concurrent {}'.format( settings.num_hosts)) binding_args = settings.binding_args if settings.binding_args else ' '.join( impl_binding_args) # Pass all the env variables to the mpirun command. mpirun_command = ( 'mpirun --allow-run-as-root --tag-output ' '-np {num_proc} {hosts_arg} ' '{binding_args} ' '{mpi_args} ' '{ssh_port_arg} ' '{tcp_intf_arg} ' '{nccl_socket_intf_arg} ' '{output_filename_arg} ' '{env} {extra_mpi_args} {command}' # expect a lot of environment variables .format(num_proc=settings.num_proc, hosts_arg=hosts_arg, binding_args=binding_args, mpi_args=' '.join(mpi_impl_flags), tcp_intf_arg=tcp_intf_arg, nccl_socket_intf_arg=nccl_socket_intf_arg, ssh_port_arg=ssh_port_arg, output_filename_arg='--output-filename ' + settings.output_filename if settings.output_filename else '', env=' '.join('-x %s' % key for key in sorted(env.keys()) if env_util.is_exportable(key)), extra_mpi_args=settings.extra_mpi_args if settings.extra_mpi_args else '', command=' '.join(quote(par) for par in command))) if settings.verbose >= 2: print(mpirun_command) # we need the driver's PATH in env to run mpirun, # env for mpirun is different to env encoded in mpirun_command if 'PATH' not in env and 'PATH' in os.environ: env = copy.copy(env) # copy env so we do not leak env modifications env['PATH'] = os.environ['PATH'] # Execute the mpirun command. if settings.run_func_mode: exit_code = safe_shell_exec.execute(mpirun_command, env=env, stdout=stdout, stderr=stderr) if exit_code != 0: raise RuntimeError( "mpirun failed with exit code {exit_code}".format( exit_code=exit_code)) else: os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)
def run(): args = parse_args() if args.version: print(horovod.__version__) exit(0) if args.host: all_host_names = [ x for x in [y.split(':')[0] for y in args.host.split(',')] ] else: all_host_names = [] # This cache stores the results of checks performed by horovodrun # during the initialization step. It can be disabled by setting # --disable-cache flag. fn_cache = None if not args.disable_cache: params = '' if args.np: params += str(args.np) + ' ' if args.host: params += str(args.host) + ' ' if args.ssh_port: params += str(args.ssh_port) parameters_hash = hashlib.md5(params.encode('utf-8')).hexdigest() fn_cache = cache.Cache(CACHE_FOLDER, CACHE_STALENESS_THRESHOLD_MINUTES, parameters_hash) # horovodrun has to finish all the checks before this timeout runs out. if args.start_timeout: start_timeout = args.start_timeout else: # Lookup default timeout from the environment variable. start_timeout = int(os.getenv('HOROVOD_START_TIMEOUT', '600')) tmout = timeout.Timeout(start_timeout) key = secret.make_secret_key() remote_host_names = [] if args.host: if args.verbose: print("Filtering local host names.") remote_host_names = network.filter_local_addresses(all_host_names) if len(remote_host_names) > 0: if args.verbose: print("Checking ssh on all remote hosts.") # Check if we can ssh into all remote hosts successfully. _check_all_hosts_ssh_successful(remote_host_names, args.ssh_port, fn_cache=fn_cache) if args.verbose: print("SSH was successful into all the remote hosts.") hosts_arg = "-H {hosts}".format(hosts=args.host) else: # if user does not specify any hosts, mpirun by default uses local host. # There is no need to specify localhost. hosts_arg = "" if args.host and len(remote_host_names) > 0: if args.verbose: print("Testing interfaces on all the hosts.") local_host_names = set(all_host_names) - set(remote_host_names) # Find the set of common, routed interfaces on all the hosts (remote # and local) and specify it in the args to be used by NCCL. It is # expected that the following function will find at least one interface # otherwise, it will raise an exception. common_intfs = _driver_fn(key, all_host_names, local_host_names, tmout, args.ssh_port, args.verbose, fn_cache=fn_cache) tcp_intf_arg = "-mca btl_tcp_if_include {common_intfs}".format( common_intfs=','.join(common_intfs)) nccl_socket_intf_arg = "-x NCCL_SOCKET_IFNAME={common_intfs}".format( common_intfs=','.join(common_intfs)) if args.verbose: print("Interfaces on all the hosts were successfully checked.") else: # If all the given hosts are local, no need to specify the interfaces # because MPI does not use network for local execution. tcp_intf_arg = "" nccl_socket_intf_arg = "" # Pass all the env variables to the mpirun command. env = os.environ.copy() # Pass secret key through the environment variables. env[secret.HOROVOD_SECRET_KEY] = codec.dumps_base64(key) if not _is_open_mpi_installed(): raise Exception( 'horovodrun convenience script currently only supports ' 'Open MPI.\n\n' 'Choose one of:\n' '1. Install Open MPI 4.0.0+ and re-install Horovod ' '(use --no-cache-dir pip option).\n' '2. Run distributed ' 'training script using the standard way provided by your' ' MPI distribution (usually mpirun, srun, or jsrun).') if args.ssh_port: ssh_port_arg = "-mca plm_rsh_args \"-p {ssh_port}\"".format( ssh_port=args.ssh_port) else: ssh_port_arg = "" mpirun_command = ( 'mpirun --allow-run-as-root --tag-output ' '-np {num_proc} {hosts_arg} ' '-bind-to none -map-by slot ' '-mca pml ob1 -mca btl ^openib ' '{ssh_port_arg} ' '{tcp_intf_arg} ' '-x NCCL_DEBUG=INFO ' '{nccl_socket_intf_arg} ' '{env} {command}' # expect a lot of environment variables .format(num_proc=args.np, hosts_arg=hosts_arg, tcp_intf_arg=tcp_intf_arg, nccl_socket_intf_arg=nccl_socket_intf_arg, ssh_port_arg=ssh_port_arg, env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)), command=' '.join(quote(par) for par in args.command))) if args.verbose: print(mpirun_command) # Execute the mpirun command. os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)
def _launch_jobs(settings, env, host_alloc_plan, remote_host_names, _run_command): """ executes the jobs defined by run command on hosts. :param hosts_alloc: list of dict indicating the allocating info. For example, [{'Hostname':'worker-0', 'Rank': 0, 'Local_rank': 0, 'Cross_rank':0, 'Size':2, 'Local_size':1, 'Cross_size':2}, {'Hostname':'worker-1', 'Rank': 1, 'Local_rank': 0, 'Cross_rank':1, 'Size':2, 'Local_size':1, 'Cross_size':2} ] :type hosts_alloc: list(dict) :param remote_host_names: names that are resolved to one of the addresses of remote hosts interfaces. :type remote_host_names: set :param _run_command: command to execute :type _run_command: string :return: :rtype: """ def _exec_command(command, index, event): if settings.verbose: print(command) # Redirect output if requested stdout = stderr = None stdout_file = stderr_file = None if settings.output_filename: padded_rank = _pad_rank(index, settings.num_proc) output_dir_rank = os.path.join( settings.output_filename, 'rank.{rank}'.format(rank=padded_rank)) if not os.path.exists(output_dir_rank): os.mkdir(output_dir_rank) stdout_file = open(os.path.join(output_dir_rank, 'stdout'), 'w') stderr_file = open(os.path.join(output_dir_rank, 'stderr'), 'w') stdout = MultiFile([sys.stdout, stdout_file]) stderr = MultiFile([sys.stderr, stderr_file]) try: exit_code = safe_shell_exec.execute(command, index=index, event=event, stdout=stdout, stderr=stderr) if exit_code != 0: print('Process {idx} exit with status code {ec}.'.format( idx=index, ec=exit_code)) except Exception as e: print('Exception happened during safe_shell_exec, exception ' 'message: {message}'.format(message=e)) finally: if stdout_file: stdout_file.close() if stderr_file: stderr_file.close() return 0 ssh_port_arg = '-p {ssh_port}'.format( ssh_port=settings.ssh_port) if settings.ssh_port else '' # Create a event for communication between threads event = threading.Event() def set_event_on_sigterm(signum, frame): event.set() signal.signal(signal.SIGINT, set_event_on_sigterm) signal.signal(signal.SIGTERM, set_event_on_sigterm) args_list = [] for alloc_info in host_alloc_plan: # generate env for rendezvous horovod_rendez_env = 'HOROVOD_RANK={rank} HOROVOD_SIZE={size} ' \ 'HOROVOD_LOCAL_RANK={local_rank} HOROVOD_LOCAL_SIZE={local_size} ' \ 'HOROVOD_CROSS_RANK={cross_rank} HOROVOD_CROSS_SIZE={cross_size} ' \ .format(rank=alloc_info.rank, size=alloc_info.size, local_rank=alloc_info.local_rank, local_size=alloc_info.local_size, cross_rank=alloc_info.cross_rank, cross_size=alloc_info.cross_size) host_name = alloc_info.hostname # TODO: Workaround for over-buffered outputs. Investigate how mpirun avoids this problem. env['PYTHONUNBUFFERED'] = '1' local_command = '{horovod_env} {env} {run_command}'.format( horovod_env=horovod_rendez_env, env=' '.join([ '%s=%s' % (key, quote(value)) for key, value in env.items() if env_util.is_exportable(key) ]), run_command=_run_command) if host_name not in remote_host_names: command = local_command else: command = 'ssh -o StrictHostKeyChecking=no {host} {ssh_port_arg} ' \ '{local_command}'.format( host=host_name, ssh_port_arg=ssh_port_arg, local_command=quote('cd {pwd} >& /dev/null ; {local_command}' .format(pwd=os.getcwd(), local_command=local_command)) ) args_list.append([command, alloc_info.rank, event]) # Make the output directory if it does not exist if settings.output_filename: _mkdir_p(settings.output_filename) # Each thread will use ssh command to launch the job on each remote host. If an # error occurs in one thread, entire process will be terminated. Otherwise, # threads will keep running and ssh session. In case, the main thread receives # a SIGINT, the event will be set and the spawned threads will kill their # corresponding middleman processes and thus the jobs will be killed as # well. threads.execute_function_multithreaded(_exec_command, args_list, block_until_all_done=True)