# Buggy code
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}".format(  # expect a lot of environment variables
            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)


# Traceback (ModuleNotFoundError)
File "/usr/lib/python3.7/runpy.py", line 183, in _run_module_as_main
mod_name, mod_spec, code = _get_module_details(mod_name, _Error)
File "/usr/lib/python3.7/runpy.py", line 109, in _get_module_details
__import__(pkg_name)
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/__init__.py", line 18, in <module>
from .runner import run
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/runner.py", line 20, in <module>
import pyspark
ModuleNotFoundError: No module named 'pyspark'
--------------------------------------------------------------------------
ORTE was unable to reliably start one or more daemons.
This usually is caused by:
...
--------------------------------------------------------------------------
Exception in thread Thread-3:
Traceback (most recent call last):
File "/usr/lib/python3.7/threading.py", line 926, in _bootstrap_inner
self.run()
File "/usr/lib/python3.7/threading.py", line 870, in run
self._target(*self._args, **self._kwargs)
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/runner.py", line 100, in run_spark
result = procs.mapPartitionsWithIndex(_make_mapper(driver.addresses(), settings, use_gloo)).collect()
File "/home/weichen.xu/spark-3.0.0-preview2-bin-hadoop2.7/python/pyspark/rdd.py", line 889, in collect
sock_info = self.ctx._jvm.PythonRDD.collectAndServe(self._jrdd.rdd())
File "/home/weichen.xu/spark-3.0.0-preview2-bin-hadoop2.7/python/lib/py4j-0.10.8.1-src.zip/py4j/java_gateway.py", line 1286, in __call__
answer, self.gateway_client, self.target_id, self.name)
File "/home/weichen.xu/spark-3.0.0-preview2-bin-hadoop2.7/python/pyspark/sql/utils.py", line 98, in deco
return f(*a, **kw)
File "/home/weichen.xu/spark-3.0.0-preview2-bin-hadoop2.7/python/lib/py4j-0.10.8.1-src.zip/py4j/protocol.py", line 328, in get_return_value
format(target_id, ".", name), value)
py4j.protocol.Py4JJavaError: An error occurred while calling z:org.apache.spark.api.python.PythonRDD.collectAndServe.
: org.apache.spark.SparkException: Job 0 cancelled part of cancelled job group horovod.spark.run.0
at org.apache.spark.scheduler.DAGScheduler.failJobAndIndependentStages(DAGScheduler.scala:1989)
at org.apache.spark.scheduler.DAGScheduler.handleJobCancellation(DAGScheduler.scala:1924)
at org.apache.spark.scheduler.DAGScheduler.$anonfun$handleJobGroupCancelled$4(DAGScheduler.scala:937)
at scala.runtime.java8.JFunction1$mcVI$sp.apply(JFunction1$mcVI$sp.java:23)
at scala.collection.mutable.HashSet.foreach(HashSet.scala:79)
at org.apache.spark.scheduler.DAGScheduler.handleJobGroupCancelled(DAGScheduler.scala:936)
at org.apache.spark.scheduler.DAGSchedulerEventProcessLoop.doOnReceive(DAGScheduler.scala:2175)
at org.apache.spark.scheduler.DAGSchedulerEventProcessLoop.onReceive(DAGScheduler.scala:2155)
at org.apache.spark.scheduler.DAGSchedulerEventProcessLoop.onReceive(DAGScheduler.scala:2144)
at org.apache.spark.util.EventLoop$$anon$1.run(EventLoop.scala:49)
at org.apache.spark.scheduler.DAGScheduler.runJob(DAGScheduler.scala:758)
at org.apache.spark.SparkContext.runJob(SparkContext.scala:2116)
at org.apache.spark.SparkContext.runJob(SparkContext.scala:2137)
at org.apache.spark.SparkContext.runJob(SparkContext.scala:2156)
at org.apache.spark.SparkContext.runJob(SparkContext.scala:2181)
at org.apache.spark.rdd.RDD.$anonfun$collect$1(RDD.scala:1004)
at org.apache.spark.rdd.RDDOperationScope$.withScope(RDDOperationScope.scala:151)
at org.apache.spark.rdd.RDDOperationScope$.withScope(RDDOperationScope.scala:112)
at org.apache.spark.rdd.RDD.withScope(RDD.scala:388)
at org.apache.spark.rdd.RDD.collect(RDD.scala:1003)
at org.apache.spark.api.python.PythonRDD$.collectAndServe(PythonRDD.scala:168)
at org.apache.spark.api.python.PythonRDD.collectAndServe(PythonRDD.scala)
at sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
at sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
at sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
at java.lang.reflect.Method.invoke(Method.java:498)
at py4j.reflection.MethodInvoker.invoke(MethodInvoker.java:244)
at py4j.reflection.ReflectionEngine.invoke(ReflectionEngine.java:357)
at py4j.Gateway.invoke(Gateway.java:282)
at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132)
at py4j.commands.CallCommand.execute(CallCommand.java:79)
at py4j.GatewayConnection.run(GatewayConnection.java:238)
at java.lang.Thread.run(Thread.java:748)


Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/runner.py", line 227, in run
_launch_job(use_mpi, use_gloo, settings, driver, env, stdout, stderr)
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/runner.py", line 123, in _launch_job
settings.verbose)
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/run/runner.py", line 686, in run_controller
mpi_run()
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/runner.py", line 121, in <lambda>
use_mpi, lambda: mpi_run(settings, nics, driver, env, stdout, stderr),
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/spark/mpi_run.py", line 54, in mpi_run
hr_mpi_run(settings, nics, env, command, stdout=stdout, stderr=stderr)
File "/home/weichen.xu/.local/lib/python3.7/site-packages/horovod/run/mpi_run.py", line 201, in mpi_run
raise RuntimeError("mpirun failed with exit code {exit_code}".format(exit_code=exit_code))
RuntimeError: mpirun failed with exit code 1
20/06/18 05:40:55 WARN PythonRunner: Incomplete task 0.0 in stage 0 (TID 0) interrupted: Attempting to kill Python Worker
20/06/18 05:40:55 WARN PythonRunner: Incomplete task 1.0 in stage 0 (TID 1) interrupted: Attempting to kill Python Worker
20/06/18 05:40:56 WARN TaskSetManager: Lost task 1.0 in stage 0.0 (TID 1, ip-10-20-4-87.us-west-2.compute.internal, executor driver): TaskKilled (Stage cancelled)
20/06/18 05:40:56 WARN TaskSetManager: Lost task 0.0 in stage 0.0 (TID 0, ip-10-20-4-87.us-west-2.compute.internal, executor driver): TaskKilled (Stage cancelled)

# Fixed code
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}".format(  # expect a lot of environment variables
            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 and PYTHONPATH in env to run mpirun,
    # env for mpirun is different to env encoded in mpirun_command
    for var in ["PATH", "PYTHONPATH"]:
        if var not in env and var in os.environ:
            # copy env so we do not leak env modifications
            env = copy.copy(env)
            # copy var over from os.environ
            env[var] = os.environ[var]

    # 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)