We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Running horovod.spark.run with env=os.environ fails
Example:

    horovod.spark.run(fn, num_proc=2, env=os.environ)

That `env` is an object, not a dictionary. It cannot be pickled:

```
Traceback (most recent call last):
  File "horovod/run/common/util/tiny_shell_exec.py", line 32, in execute
    exit_code = safe_shell_exec.execute(command, env=env, stdout=output, stderr=output)
  File "horovod/run/common/util/safe_shell_exec.py", line 183, in execute
    middleman.start()
  File "multiprocessing/process.py", line 105, in start
    self._popen = self._Popen(self)
  File "multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object '_createenviron.<locals>.encode'
```

It works with

    horovod.spark.run(fn, num_proc=2, env=os.environ.copy())

The `run` function needs to copy `env` itself first.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `horovod/run/mpi_run.py`
Content:
```
1 # Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 
16 import copy
17 import os
18 import sys
19 
20 from shlex import quote
21 
22 from horovod.run.common.util import env as env_util, safe_shell_exec, tiny_shell_exec
23 
24 # MPI implementations
25 _OMPI_IMPL = 'OpenMPI'
26 _SMPI_IMPL = 'SpectrumMPI'
27 _MPICH_IMPL = 'MPICH'
28 _UNKNOWN_IMPL = 'Unknown'
29 _MISSING_IMPL = 'Missing'
30 
31 # Open MPI Flags
32 _OMPI_FLAGS = ['-mca pml ob1', '-mca btl ^openib']
33 # Spectrum MPI Flags
34 _SMPI_FLAGS = []
35 _SMPI_FLAGS_TCP = ['-tcp']
36 # MPICH Flags
37 _MPICH_FLAGS = []
38 
39 # Threshold for large cluster MPI issues:
40 _LARGE_CLUSTER_THRESHOLD = 64
41 # No process binding args
42 _NO_BINDING_ARGS = ['-bind-to none', '-map-by slot']
43 # Process socket binding args
44 _SOCKET_BINDING_ARGS = ['-bind-to socket', '-map-by socket', '-rank-by core']
45 
46 # MPI not found error message
47 _MPI_NOT_FOUND_ERROR_MSG= ('horovod does not find an installed MPI.\n\n'
48                            'Choose one of:\n'
49                            '1. Install Open MPI 4.0.0+ or IBM Spectrum MPI or MPICH and re-install Horovod '
50                            '(use --no-cache-dir pip option).\n'
51                            '2. Run distributed '
52                            'training script using the standard way provided by your'
53                            ' MPI distribution (usually mpirun, srun, or jsrun).\n'
54                            '3. Use built-in gloo option (horovodrun --gloo ...).')
55 
56 
57 def mpi_available(env=None):
58     return _get_mpi_implementation(env) not in {_UNKNOWN_IMPL, _MISSING_IMPL}
59 
60 
61 def is_open_mpi(env=None):
62     return _get_mpi_implementation(env) == _OMPI_IMPL
63 
64 
65 def is_spectrum_mpi(env=None):
66     return _get_mpi_implementation(env) == _SMPI_IMPL
67 
68 
69 def is_mpich(env=None):
70     return _get_mpi_implementation(env) == _MPICH_IMPL
71 
72 
73 def _get_mpi_implementation(env=None):
74     """
75     Detects the available MPI implementation by invoking `mpirun --version`.
76     This command is executed by the given execute function, which takes the
77     command as the only argument and returns (output, exit code). Output
78     represents the stdout and stderr as a string.
79 
80     Returns one of:
81     - _OMPI_IMPL, _SMPI_IMPL or _MPICH_IMPL for known implementations
82     - _UNKNOWN_IMPL for any unknown implementation
83     - _MISSING_IMPL if `mpirun --version` could not be executed.
84 
85     :param env: environment variable to use to run mpirun
86     :return: string representing identified implementation
87     """
88     command = 'mpirun --version'
89     res = tiny_shell_exec.execute(command, env)
90     if res is None:
91         return _MISSING_IMPL
92     (output, exit_code) = res
93 
94     if exit_code == 0:
95         if 'Open MPI' in output or 'OpenRTE' in output:
96             return _OMPI_IMPL
97         elif 'IBM Spectrum MPI' in output:
98             return _SMPI_IMPL
99         elif 'MPICH' in output:
100             return _MPICH_IMPL
101 
102         print('Unknown MPI implementation given in output of mpirun --version:', file=sys.stderr)
103         print(output, file=sys.stderr)
104         return _UNKNOWN_IMPL
105     else:
106         print('Was unable to run {command}:'.format(command=command), file=sys.stderr)
107         print(output, file=sys.stderr)
108         return _MISSING_IMPL
109 
110 
111 def _get_mpi_implementation_flags(tcp_flag, env=None):
112     if is_open_mpi(env):
113         return list(_OMPI_FLAGS), list(_NO_BINDING_ARGS)
114     elif is_spectrum_mpi(env):
115         return list(_SMPI_FLAGS) if not tcp_flag else list(_SMPI_FLAGS_TCP), list(_SOCKET_BINDING_ARGS)
116     elif is_mpich(env):
117         return list(_MPICH_FLAGS), list(_NO_BINDING_ARGS)
118     else:
119         return None, None
120 
121 
122 def mpi_run(settings, nics, env, command, stdout=None, stderr=None):
123     """
124     Runs mpi_run.
125 
126     Args:
127         settings: Settings for running MPI.
128                   Note: settings.num_proc and settings.hosts must not be None.
129         nics: Interfaces to include by MPI.
130         env: Environment dictionary to use for running command.
131         command: Command and arguments to run as a list of string.
132         stdout: Stdout of the mpi process.
133                 Only used when settings.run_func_mode is True.
134         stderr: Stderr of the mpi process.
135                 Only used when settings.run_func_mode is True.
136     """
137     mpi_impl_flags, impl_binding_args = _get_mpi_implementation_flags(settings.tcp_flag, env=env)
138     if mpi_impl_flags is None:
139         raise Exception(_MPI_NOT_FOUND_ERROR_MSG)
140 
141     ssh_port_arg = '-mca plm_rsh_args \"-p {ssh_port}\"'.format(
142             ssh_port=settings.ssh_port) if settings.ssh_port else ''
143 
144     # if user does not specify any hosts, mpirun by default uses local host.
145     # There is no need to specify localhost.
146     hosts_arg = '-H {hosts}'.format(hosts=settings.hosts)
147 
148     tcp_intf_arg = '-mca btl_tcp_if_include {nics}'.format(
149         nics=','.join(nics)) if nics else ''
150     nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={nics}'.format(
151         nics=','.join(nics)) if nics else ''
152 
153     # On large cluster runs (e.g. Summit), we need extra settings to work around OpenMPI issues
154     if settings.num_hosts and settings.num_hosts >= _LARGE_CLUSTER_THRESHOLD:
155         mpi_impl_flags.append('-mca plm_rsh_no_tree_spawn true')
156         mpi_impl_flags.append('-mca plm_rsh_num_concurrent {}'.format(settings.num_hosts))
157 
158     binding_args = settings.binding_args if settings.binding_args else ' '.join(impl_binding_args)
159 
160     # Pass all the env variables to the mpirun command.
161     mpirun_command = (
162         'mpirun --allow-run-as-root --tag-output '
163         '-np {num_proc} {hosts_arg} '
164         '{binding_args} '
165         '{mpi_args} '
166         '{ssh_port_arg} '
167         '{tcp_intf_arg} '
168         '{nccl_socket_intf_arg} '
169         '{output_filename_arg} '
170         '{env} {extra_mpi_args} {command}'  # expect a lot of environment variables
171         .format(num_proc=settings.num_proc,
172                 hosts_arg=hosts_arg,
173                 binding_args=binding_args,
174                 mpi_args=' '.join(mpi_impl_flags),
175                 tcp_intf_arg=tcp_intf_arg,
176                 nccl_socket_intf_arg=nccl_socket_intf_arg,
177                 ssh_port_arg=ssh_port_arg,
178                 output_filename_arg='--output-filename ' + settings.output_filename
179                                     if settings.output_filename else '',
180                 env=' '.join('-x %s' % key for key in sorted(env.keys())
181                              if env_util.is_exportable(key)),
182 
183                 extra_mpi_args=settings.extra_mpi_args if settings.extra_mpi_args else '',
184                 command=' '.join(quote(par) for par in command))
185     )
186 
187     if settings.verbose >= 2:
188         print(mpirun_command)
189 
190     # we need the driver's PATH in env to run mpirun,
191     # env for mpirun is different to env encoded in mpirun_command
192     if 'PATH' not in env and 'PATH' in os.environ:
193         env = copy.copy(env)  # copy env so we do not leak env modifications
194         env['PATH'] = os.environ['PATH']
195 
196     # Execute the mpirun command.
197     if settings.run_func_mode:
198         exit_code = safe_shell_exec.execute(mpirun_command, env=env, stdout=stdout, stderr=stderr)
199         if exit_code != 0:
200             raise RuntimeError("mpirun failed with exit code {exit_code}".format(exit_code=exit_code))
201     else:
202         os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)
203
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/horovod/run/mpi_run.py b/horovod/run/mpi_run.py
--- a/horovod/run/mpi_run.py
+++ b/horovod/run/mpi_run.py
@@ -134,6 +134,10 @@
         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)

{"golden_diff": "diff --git a/horovod/run/mpi_run.py b/horovod/run/mpi_run.py\n--- a/horovod/run/mpi_run.py\n+++ b/horovod/run/mpi_run.py\n@@ -134,6 +134,10 @@\n         stderr: Stderr of the mpi process.\n                 Only used when settings.run_func_mode is True.\n     \"\"\"\n+    if env is not None and not isinstance(env, dict):\n+        raise Exception('env argument must be a dict, not {type}: {env}'\n+                        .format(type=type(env), env=env))\n+\n     mpi_impl_flags, impl_binding_args = _get_mpi_implementation_flags(settings.tcp_flag, env=env)\n     if mpi_impl_flags is None:\n         raise Exception(_MPI_NOT_FOUND_ERROR_MSG)\n", "issue": "Running horovod.spark.run with env=os.environ fails\nExample:\r\n\r\n    horovod.spark.run(fn, num_proc=2, env=os.environ)\r\n\r\nThat `env` is an object, not a dictionary. It cannot be pickled:\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"horovod/run/common/util/tiny_shell_exec.py\", line 32, in execute\r\n    exit_code = safe_shell_exec.execute(command, env=env, stdout=output, stderr=output)\r\n  File \"horovod/run/common/util/safe_shell_exec.py\", line 183, in execute\r\n    middleman.start()\r\n  File \"multiprocessing/process.py\", line 105, in start\r\n    self._popen = self._Popen(self)\r\n  File \"multiprocessing/context.py\", line 284, in _Popen\r\n    return Popen(process_obj)\r\n  File \"multiprocessing/popen_spawn_posix.py\", line 32, in __init__\r\n    super().__init__(process_obj)\r\n  File \"multiprocessing/popen_fork.py\", line 19, in __init__\r\n    self._launch(process_obj)\r\n  File \"multiprocessing/popen_spawn_posix.py\", line 47, in _launch\r\n    reduction.dump(process_obj, fp)\r\n  File \"multiprocessing/reduction.py\", line 60, in dump\r\n    ForkingPickler(file, protocol).dump(obj)\r\nAttributeError: Can't pickle local object '_createenviron.<locals>.encode'\r\n```\r\n\r\nIt works with\r\n\r\n    horovod.spark.run(fn, num_proc=2, env=os.environ.copy())\r\n\r\nThe `run` function needs to copy `env` itself first.\n", "before_files": [{"content": "# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nimport copy\nimport os\nimport sys\n\nfrom shlex import quote\n\nfrom horovod.run.common.util import env as env_util, safe_shell_exec, tiny_shell_exec\n\n# MPI implementations\n_OMPI_IMPL = 'OpenMPI'\n_SMPI_IMPL = 'SpectrumMPI'\n_MPICH_IMPL = 'MPICH'\n_UNKNOWN_IMPL = 'Unknown'\n_MISSING_IMPL = 'Missing'\n\n# Open MPI Flags\n_OMPI_FLAGS = ['-mca pml ob1', '-mca btl ^openib']\n# Spectrum MPI Flags\n_SMPI_FLAGS = []\n_SMPI_FLAGS_TCP = ['-tcp']\n# MPICH Flags\n_MPICH_FLAGS = []\n\n# Threshold for large cluster MPI issues:\n_LARGE_CLUSTER_THRESHOLD = 64\n# No process binding args\n_NO_BINDING_ARGS = ['-bind-to none', '-map-by slot']\n# Process socket binding args\n_SOCKET_BINDING_ARGS = ['-bind-to socket', '-map-by socket', '-rank-by core']\n\n# MPI not found error message\n_MPI_NOT_FOUND_ERROR_MSG= ('horovod does not find an installed MPI.\\n\\n'\n                           'Choose one of:\\n'\n                           '1. Install Open MPI 4.0.0+ or IBM Spectrum MPI or MPICH and re-install Horovod '\n                           '(use --no-cache-dir pip option).\\n'\n                           '2. Run distributed '\n                           'training script using the standard way provided by your'\n                           ' MPI distribution (usually mpirun, srun, or jsrun).\\n'\n                           '3. Use built-in gloo option (horovodrun --gloo ...).')\n\n\ndef mpi_available(env=None):\n    return _get_mpi_implementation(env) not in {_UNKNOWN_IMPL, _MISSING_IMPL}\n\n\ndef is_open_mpi(env=None):\n    return _get_mpi_implementation(env) == _OMPI_IMPL\n\n\ndef is_spectrum_mpi(env=None):\n    return _get_mpi_implementation(env) == _SMPI_IMPL\n\n\ndef is_mpich(env=None):\n    return _get_mpi_implementation(env) == _MPICH_IMPL\n\n\ndef _get_mpi_implementation(env=None):\n    \"\"\"\n    Detects the available MPI implementation by invoking `mpirun --version`.\n    This command is executed by the given execute function, which takes the\n    command as the only argument and returns (output, exit code). Output\n    represents the stdout and stderr as a string.\n\n    Returns one of:\n    - _OMPI_IMPL, _SMPI_IMPL or _MPICH_IMPL for known implementations\n    - _UNKNOWN_IMPL for any unknown implementation\n    - _MISSING_IMPL if `mpirun --version` could not be executed.\n\n    :param env: environment variable to use to run mpirun\n    :return: string representing identified implementation\n    \"\"\"\n    command = 'mpirun --version'\n    res = tiny_shell_exec.execute(command, env)\n    if res is None:\n        return _MISSING_IMPL\n    (output, exit_code) = res\n\n    if exit_code == 0:\n        if 'Open MPI' in output or 'OpenRTE' in output:\n            return _OMPI_IMPL\n        elif 'IBM Spectrum MPI' in output:\n            return _SMPI_IMPL\n        elif 'MPICH' in output:\n            return _MPICH_IMPL\n\n        print('Unknown MPI implementation given in output of mpirun --version:', file=sys.stderr)\n        print(output, file=sys.stderr)\n        return _UNKNOWN_IMPL\n    else:\n        print('Was unable to run {command}:'.format(command=command), file=sys.stderr)\n        print(output, file=sys.stderr)\n        return _MISSING_IMPL\n\n\ndef _get_mpi_implementation_flags(tcp_flag, env=None):\n    if is_open_mpi(env):\n        return list(_OMPI_FLAGS), list(_NO_BINDING_ARGS)\n    elif is_spectrum_mpi(env):\n        return list(_SMPI_FLAGS) if not tcp_flag else list(_SMPI_FLAGS_TCP), list(_SOCKET_BINDING_ARGS)\n    elif is_mpich(env):\n        return list(_MPICH_FLAGS), list(_NO_BINDING_ARGS)\n    else:\n        return None, None\n\n\ndef mpi_run(settings, nics, env, command, stdout=None, stderr=None):\n    \"\"\"\n    Runs mpi_run.\n\n    Args:\n        settings: Settings for running MPI.\n                  Note: settings.num_proc and settings.hosts must not be None.\n        nics: Interfaces to include by MPI.\n        env: Environment dictionary to use for running command.\n        command: Command and arguments to run as a list of string.\n        stdout: Stdout of the mpi process.\n                Only used when settings.run_func_mode is True.\n        stderr: Stderr of the mpi process.\n                Only used when settings.run_func_mode is True.\n    \"\"\"\n    mpi_impl_flags, impl_binding_args = _get_mpi_implementation_flags(settings.tcp_flag, env=env)\n    if mpi_impl_flags is None:\n        raise Exception(_MPI_NOT_FOUND_ERROR_MSG)\n\n    ssh_port_arg = '-mca plm_rsh_args \\\"-p {ssh_port}\\\"'.format(\n            ssh_port=settings.ssh_port) if settings.ssh_port else ''\n\n    # if user does not specify any hosts, mpirun by default uses local host.\n    # There is no need to specify localhost.\n    hosts_arg = '-H {hosts}'.format(hosts=settings.hosts)\n\n    tcp_intf_arg = '-mca btl_tcp_if_include {nics}'.format(\n        nics=','.join(nics)) if nics else ''\n    nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={nics}'.format(\n        nics=','.join(nics)) if nics else ''\n\n    # On large cluster runs (e.g. Summit), we need extra settings to work around OpenMPI issues\n    if settings.num_hosts and settings.num_hosts >= _LARGE_CLUSTER_THRESHOLD:\n        mpi_impl_flags.append('-mca plm_rsh_no_tree_spawn true')\n        mpi_impl_flags.append('-mca plm_rsh_num_concurrent {}'.format(settings.num_hosts))\n\n    binding_args = settings.binding_args if settings.binding_args else ' '.join(impl_binding_args)\n\n    # Pass all the env variables to the mpirun command.\n    mpirun_command = (\n        'mpirun --allow-run-as-root --tag-output '\n        '-np {num_proc} {hosts_arg} '\n        '{binding_args} '\n        '{mpi_args} '\n        '{ssh_port_arg} '\n        '{tcp_intf_arg} '\n        '{nccl_socket_intf_arg} '\n        '{output_filename_arg} '\n        '{env} {extra_mpi_args} {command}'  # expect a lot of environment variables\n        .format(num_proc=settings.num_proc,\n                hosts_arg=hosts_arg,\n                binding_args=binding_args,\n                mpi_args=' '.join(mpi_impl_flags),\n                tcp_intf_arg=tcp_intf_arg,\n                nccl_socket_intf_arg=nccl_socket_intf_arg,\n                ssh_port_arg=ssh_port_arg,\n                output_filename_arg='--output-filename ' + settings.output_filename\n                                    if settings.output_filename else '',\n                env=' '.join('-x %s' % key for key in sorted(env.keys())\n                             if env_util.is_exportable(key)),\n\n                extra_mpi_args=settings.extra_mpi_args if settings.extra_mpi_args else '',\n                command=' '.join(quote(par) for par in command))\n    )\n\n    if settings.verbose >= 2:\n        print(mpirun_command)\n\n    # we need the driver's PATH in env to run mpirun,\n    # env for mpirun is different to env encoded in mpirun_command\n    if 'PATH' not in env and 'PATH' in os.environ:\n        env = copy.copy(env)  # copy env so we do not leak env modifications\n        env['PATH'] = os.environ['PATH']\n\n    # Execute the mpirun command.\n    if settings.run_func_mode:\n        exit_code = safe_shell_exec.execute(mpirun_command, env=env, stdout=stdout, stderr=stderr)\n        if exit_code != 0:\n            raise RuntimeError(\"mpirun failed with exit code {exit_code}\".format(exit_code=exit_code))\n    else:\n        os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)\n", "path": "horovod/run/mpi_run.py"}], "after_files": [{"content": "# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nimport copy\nimport os\nimport sys\n\nfrom shlex import quote\n\nfrom horovod.run.common.util import env as env_util, safe_shell_exec, tiny_shell_exec\n\n# MPI implementations\n_OMPI_IMPL = 'OpenMPI'\n_SMPI_IMPL = 'SpectrumMPI'\n_MPICH_IMPL = 'MPICH'\n_UNKNOWN_IMPL = 'Unknown'\n_MISSING_IMPL = 'Missing'\n\n# Open MPI Flags\n_OMPI_FLAGS = ['-mca pml ob1', '-mca btl ^openib']\n# Spectrum MPI Flags\n_SMPI_FLAGS = []\n_SMPI_FLAGS_TCP = ['-tcp']\n# MPICH Flags\n_MPICH_FLAGS = []\n\n# Threshold for large cluster MPI issues:\n_LARGE_CLUSTER_THRESHOLD = 64\n# No process binding args\n_NO_BINDING_ARGS = ['-bind-to none', '-map-by slot']\n# Process socket binding args\n_SOCKET_BINDING_ARGS = ['-bind-to socket', '-map-by socket', '-rank-by core']\n\n# MPI not found error message\n_MPI_NOT_FOUND_ERROR_MSG= ('horovod does not find an installed MPI.\\n\\n'\n                           'Choose one of:\\n'\n                           '1. Install Open MPI 4.0.0+ or IBM Spectrum MPI or MPICH and re-install Horovod '\n                           '(use --no-cache-dir pip option).\\n'\n                           '2. Run distributed '\n                           'training script using the standard way provided by your'\n                           ' MPI distribution (usually mpirun, srun, or jsrun).\\n'\n                           '3. Use built-in gloo option (horovodrun --gloo ...).')\n\n\ndef mpi_available(env=None):\n    return _get_mpi_implementation(env) not in {_UNKNOWN_IMPL, _MISSING_IMPL}\n\n\ndef is_open_mpi(env=None):\n    return _get_mpi_implementation(env) == _OMPI_IMPL\n\n\ndef is_spectrum_mpi(env=None):\n    return _get_mpi_implementation(env) == _SMPI_IMPL\n\n\ndef is_mpich(env=None):\n    return _get_mpi_implementation(env) == _MPICH_IMPL\n\n\ndef _get_mpi_implementation(env=None):\n    \"\"\"\n    Detects the available MPI implementation by invoking `mpirun --version`.\n    This command is executed by the given execute function, which takes the\n    command as the only argument and returns (output, exit code). Output\n    represents the stdout and stderr as a string.\n\n    Returns one of:\n    - _OMPI_IMPL, _SMPI_IMPL or _MPICH_IMPL for known implementations\n    - _UNKNOWN_IMPL for any unknown implementation\n    - _MISSING_IMPL if `mpirun --version` could not be executed.\n\n    :param env: environment variable to use to run mpirun\n    :return: string representing identified implementation\n    \"\"\"\n    command = 'mpirun --version'\n    res = tiny_shell_exec.execute(command, env)\n    if res is None:\n        return _MISSING_IMPL\n    (output, exit_code) = res\n\n    if exit_code == 0:\n        if 'Open MPI' in output or 'OpenRTE' in output:\n            return _OMPI_IMPL\n        elif 'IBM Spectrum MPI' in output:\n            return _SMPI_IMPL\n        elif 'MPICH' in output:\n            return _MPICH_IMPL\n\n        print('Unknown MPI implementation given in output of mpirun --version:', file=sys.stderr)\n        print(output, file=sys.stderr)\n        return _UNKNOWN_IMPL\n    else:\n        print('Was unable to run {command}:'.format(command=command), file=sys.stderr)\n        print(output, file=sys.stderr)\n        return _MISSING_IMPL\n\n\ndef _get_mpi_implementation_flags(tcp_flag, env=None):\n    if is_open_mpi(env):\n        return list(_OMPI_FLAGS), list(_NO_BINDING_ARGS)\n    elif is_spectrum_mpi(env):\n        return list(_SMPI_FLAGS) if not tcp_flag else list(_SMPI_FLAGS_TCP), list(_SOCKET_BINDING_ARGS)\n    elif is_mpich(env):\n        return list(_MPICH_FLAGS), list(_NO_BINDING_ARGS)\n    else:\n        return None, None\n\n\ndef mpi_run(settings, nics, env, command, stdout=None, stderr=None):\n    \"\"\"\n    Runs mpi_run.\n\n    Args:\n        settings: Settings for running MPI.\n                  Note: settings.num_proc and settings.hosts must not be None.\n        nics: Interfaces to include by MPI.\n        env: Environment dictionary to use for running command.\n        command: Command and arguments to run as a list of string.\n        stdout: Stdout of the mpi process.\n                Only used when settings.run_func_mode is True.\n        stderr: Stderr of the mpi process.\n                Only used when settings.run_func_mode is True.\n    \"\"\"\n    if env is not None and not isinstance(env, dict):\n        raise Exception('env argument must be a dict, not {type}: {env}'\n                        .format(type=type(env), env=env))\n\n    mpi_impl_flags, impl_binding_args = _get_mpi_implementation_flags(settings.tcp_flag, env=env)\n    if mpi_impl_flags is None:\n        raise Exception(_MPI_NOT_FOUND_ERROR_MSG)\n\n    ssh_port_arg = '-mca plm_rsh_args \\\"-p {ssh_port}\\\"'.format(\n            ssh_port=settings.ssh_port) if settings.ssh_port else ''\n\n    # if user does not specify any hosts, mpirun by default uses local host.\n    # There is no need to specify localhost.\n    hosts_arg = '-H {hosts}'.format(hosts=settings.hosts)\n\n    tcp_intf_arg = '-mca btl_tcp_if_include {nics}'.format(\n        nics=','.join(nics)) if nics else ''\n    nccl_socket_intf_arg = '-x NCCL_SOCKET_IFNAME={nics}'.format(\n        nics=','.join(nics)) if nics else ''\n\n    # On large cluster runs (e.g. Summit), we need extra settings to work around OpenMPI issues\n    if settings.num_hosts and settings.num_hosts >= _LARGE_CLUSTER_THRESHOLD:\n        mpi_impl_flags.append('-mca plm_rsh_no_tree_spawn true')\n        mpi_impl_flags.append('-mca plm_rsh_num_concurrent {}'.format(settings.num_hosts))\n\n    binding_args = settings.binding_args if settings.binding_args else ' '.join(impl_binding_args)\n\n    # Pass all the env variables to the mpirun command.\n    mpirun_command = (\n        'mpirun --allow-run-as-root --tag-output '\n        '-np {num_proc} {hosts_arg} '\n        '{binding_args} '\n        '{mpi_args} '\n        '{ssh_port_arg} '\n        '{tcp_intf_arg} '\n        '{nccl_socket_intf_arg} '\n        '{output_filename_arg} '\n        '{env} {extra_mpi_args} {command}'  # expect a lot of environment variables\n        .format(num_proc=settings.num_proc,\n                hosts_arg=hosts_arg,\n                binding_args=binding_args,\n                mpi_args=' '.join(mpi_impl_flags),\n                tcp_intf_arg=tcp_intf_arg,\n                nccl_socket_intf_arg=nccl_socket_intf_arg,\n                ssh_port_arg=ssh_port_arg,\n                output_filename_arg='--output-filename ' + settings.output_filename\n                                    if settings.output_filename else '',\n                env=' '.join('-x %s' % key for key in sorted(env.keys())\n                             if env_util.is_exportable(key)),\n\n                extra_mpi_args=settings.extra_mpi_args if settings.extra_mpi_args else '',\n                command=' '.join(quote(par) for par in command))\n    )\n\n    if settings.verbose >= 2:\n        print(mpirun_command)\n\n    # we need the driver's PATH in env to run mpirun,\n    # env for mpirun is different to env encoded in mpirun_command\n    if 'PATH' not in env and 'PATH' in os.environ:\n        env = copy.copy(env)  # copy env so we do not leak env modifications\n        env['PATH'] = os.environ['PATH']\n\n    # Execute the mpirun command.\n    if settings.run_func_mode:\n        exit_code = safe_shell_exec.execute(mpirun_command, env=env, stdout=stdout, stderr=stderr)\n        if exit_code != 0:\n            raise RuntimeError(\"mpirun failed with exit code {exit_code}\".format(exit_code=exit_code))\n    else:\n        os.execve('/bin/sh', ['/bin/sh', '-c', mpirun_command], env)\n", "path": "horovod/run/mpi_run.py"}]}