We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
TF estimators and horovod.spark
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet)
TensorFlow
2. Framework version:
1.13.1
3. Horovod version:
0.16
4. MPI version:
3.0.0
5. CUDA version:
N/A
6. NCCL version:
N/A
7. Python version:
3.5
8. OS and version:
Amazon Linux AMI release 2018.03

**Your question:**
Is there any guidance on how to use horovod.spark with TF estimators? I've been able to train models using TF-Keras and horovod.spark in the style of the keras_spark_rossman.py example script, but I'm curious about how horovod.spark would work with Estimator's?

Moreover, I think I'm just confused about how horovod.spark actually works - is it okay to instantiate an Estimator instance inside the input_fn for a horovod.spark.run() call, or will that mean each worker has a different starting point for the model?

--- 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/spark/__init__.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 os
17 import pyspark
18 from six.moves import queue
19 import sys
20 import threading
21 
22 from horovod.spark.task import task_service
23 from horovod.run.common.util import codec, env as env_util, safe_shell_exec, \
24     timeout, host_hash, secret
25 from horovod.run.common.util import settings as hvd_settings
26 from horovod.spark.driver import driver_service, job_id
27 
28 
29 def _task_fn(index, driver_addresses, settings):
30     task = task_service.SparkTaskService(index, settings.key)
31     try:
32         driver_client = driver_service.SparkDriverClient(driver_addresses, settings.key, settings.verbose)
33         driver_client.register_task(index, task.addresses(), host_hash.host_hash())
34         task.wait_for_initial_registration(settings.timeout)
35         # Tasks ping each other in a circular fashion to determine interfaces reachable within
36         # the cluster.
37         next_task_index = (index + 1) % settings.num_proc
38         next_task_addresses = driver_client.all_task_addresses(next_task_index)
39         # We request interface matching to weed out all the NAT'ed interfaces.
40         next_task_client = \
41             task_service.SparkTaskClient(next_task_index, next_task_addresses,
42                                          settings.key, settings.verbose,
43                                          match_intf=True)
44         driver_client.register_task_to_task_addresses(next_task_index, next_task_client.addresses())
45         task_indices_on_this_host = driver_client.task_host_hash_indices(
46             host_hash.host_hash())
47         if task_indices_on_this_host[0] == index:
48             # Task with first index will execute orted that will run mpirun_exec_fn for all tasks.
49             task.wait_for_command_start(settings.timeout)
50             task.wait_for_command_termination()
51         else:
52             # The rest of tasks need to wait for the first task to finish.
53             first_task_addresses = driver_client.all_task_addresses(task_indices_on_this_host[0])
54             first_task_client = \
55                 task_service.SparkTaskClient(task_indices_on_this_host[0],
56                                              first_task_addresses, settings.key,
57                                              settings.verbose)
58             first_task_client.wait_for_command_termination()
59         return task.fn_result()
60     finally:
61         task.shutdown()
62 
63 
64 def _make_mapper(driver_addresses, settings):
65     def _mapper(index, _):
66         yield _task_fn(index, driver_addresses, settings)
67     return _mapper
68 
69 
70 def _make_spark_thread(spark_context, spark_job_group, driver, result_queue,
71                        settings):
72     def run_spark():
73         try:
74             spark_context.setJobGroup(spark_job_group,
75                                       "Horovod Spark Run",
76                                       interruptOnCancel=True)
77             procs = spark_context.range(0, numSlices=settings.num_proc)
78             # We assume that folks caring about security will enable Spark RPC
79             # encryption, thus ensuring that key that is passed here remains
80             # secret.
81             result = procs.mapPartitionsWithIndex(_make_mapper(driver.addresses(), settings)).collect()
82             result_queue.put(result)
83         except:
84             driver.notify_spark_job_failed()
85             raise
86 
87     spark_thread = threading.Thread(target=run_spark)
88     spark_thread.start()
89     return spark_thread
90 
91 
92 def run(fn, args=(), kwargs={}, num_proc=None, start_timeout=None, env=None,
93         stdout=None, stderr=None, verbose=1):
94     """
95     Runs Horovod in Spark.  Runs `num_proc` processes executing `fn` using the same amount of Spark tasks.
96 
97     Args:
98         fn: Function to run.
99         args: Arguments to pass to `fn`.
100         kwargs: Keyword arguments to pass to `fn`.
101         num_proc: Number of Horovod processes.  Defaults to `spark.default.parallelism`.
102         start_timeout: Timeout for Spark tasks to spawn, register and start running the code, in seconds.
103                        If not set, falls back to `HOROVOD_SPARK_START_TIMEOUT` environment variable value.
104                        If it is not set as well, defaults to 600 seconds.
105         env: Environment dictionary to use in Horovod run.  Defaults to `os.environ`.
106         stdout: Horovod stdout is redirected to this stream. Defaults to sys.stdout.
107         stderr: Horovod stderr is redirected to this stream. Defaults to sys.stderr.
108         verbose: Debug output verbosity (0-2). Defaults to 1.
109 
110     Returns:
111         List of results returned by running `fn` on each rank.
112     """
113 
114     if start_timeout is None:
115         # Lookup default timeout from the environment variable.
116         start_timeout = int(os.getenv('HOROVOD_SPARK_START_TIMEOUT', '600'))
117 
118     tmout = timeout.Timeout(start_timeout,
119                             message='Timed out waiting for {activity}. Please check that you have '
120                                     'enough resources to run all Horovod processes. Each Horovod '
121                                     'process runs in a Spark task. You may need to increase the '
122                                     'start_timeout parameter to a larger value if your Spark resources '
123                                     'are allocated on-demand.')
124     settings = hvd_settings.Settings(verbose=verbose,
125                                      key=secret.make_secret_key(),
126                                      timeout=tmout)
127 
128     spark_context = pyspark.SparkContext._active_spark_context
129     if spark_context is None:
130         raise Exception('Could not find an active SparkContext, are you '
131                         'running in a PySpark session?')
132 
133     if num_proc is None:
134         num_proc = spark_context.defaultParallelism
135         if settings.verbose >= 1:
136             print('Running %d processes (inferred from spark.default.parallelism)...' % num_proc)
137     else:
138         if settings.verbose >= 1:
139             print('Running %d processes...' % num_proc)
140     settings.num_proc = num_proc
141 
142     result_queue = queue.Queue(1)
143 
144     spark_job_group = 'horovod.spark.run.%d' % job_id.next_job_id()
145     driver = driver_service.SparkDriverService(settings.num_proc, fn, args, kwargs,
146                                                settings.key)
147     spark_thread = _make_spark_thread(spark_context, spark_job_group, driver,
148                                       result_queue, settings)
149     try:
150         driver.wait_for_initial_registration(settings.timeout)
151         if settings.verbose >= 2:
152             print('Initial Spark task registration is complete.')
153         task_clients = [
154             task_service.SparkTaskClient(index,
155                                          driver.task_addresses_for_driver(index),
156                                          settings.key, settings.verbose)
157             for index in range(settings.num_proc)]
158         for task_client in task_clients:
159             task_client.notify_initial_registration_complete()
160         driver.wait_for_task_to_task_address_updates(settings.timeout)
161         if settings.verbose >= 2:
162             print('Spark task-to-task address registration is complete.')
163 
164         # Determine a set of common interfaces for task-to-task communication.
165         common_intfs = set(driver.task_addresses_for_tasks(0).keys())
166         for index in range(1, settings.num_proc):
167             common_intfs.intersection_update(driver.task_addresses_for_tasks(index).keys())
168         if not common_intfs:
169             raise Exception('Unable to find a set of common task-to-task communication interfaces: %s'
170                             % [(index, driver.task_addresses_for_tasks(index)) for index in range(settings.num_proc)])
171 
172         # Determine the index grouping based on host hashes.
173         # Barrel shift until index 0 is in the first host.
174         host_hashes = list(driver.task_host_hash_indices().keys())
175         host_hashes.sort()
176         while 0 not in driver.task_host_hash_indices()[host_hashes[0]]:
177             host_hashes = host_hashes[1:] + host_hashes[:1]
178 
179         ranks_to_indices = []
180         for host_hash in host_hashes:
181             ranks_to_indices += driver.task_host_hash_indices()[host_hash]
182         driver.set_ranks_to_indices(ranks_to_indices)
183 
184         if env is None:
185             env = os.environ.copy()
186 
187         # Pass secret key through the environment variables.
188         env[secret.HOROVOD_SECRET_KEY] = codec.dumps_base64(settings.key)
189 
190         mpirun_command = (
191             'mpirun --allow-run-as-root --tag-output '
192             '-np {num_proc} -H {hosts} '
193             '-bind-to none -map-by slot '
194             '-mca pml ob1 -mca btl ^openib -mca btl_tcp_if_include {common_intfs} '
195             '-x NCCL_DEBUG=INFO -x NCCL_SOCKET_IFNAME={common_intfs} '
196             '{env} '  # expect a lot of environment variables
197             '-mca plm_rsh_agent "{python} -m horovod.spark.driver.mpirun_rsh {encoded_driver_addresses} {settings}" '
198             '{python} -m horovod.spark.task.mpirun_exec_fn {encoded_driver_addresses} {settings}'
199                 .format(num_proc=settings.num_proc,
200                         hosts=','.join('%s:%d' % (host_hash, len(driver.task_host_hash_indices()[host_hash]))
201                                        for host_hash in host_hashes),
202                         common_intfs=','.join(common_intfs),
203                         env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)),
204                         python=sys.executable,
205                         encoded_driver_addresses=codec.dumps_base64(driver.addresses()),
206                         settings=codec.dumps_base64(settings)))
207         if settings.verbose >= 2:
208             print('+ %s' % mpirun_command)
209         exit_code = safe_shell_exec.execute(mpirun_command, env, stdout, stderr)
210         if exit_code != 0:
211             raise Exception('mpirun exited with code %d, see the error above.' % exit_code)
212     except:
213         # Terminate Spark job.
214         spark_context.cancelJobGroup(spark_job_group)
215 
216         # Re-raise exception.
217         raise
218     finally:
219         spark_thread.join()
220         driver.shutdown()
221 
222     # Make sure Spark Job did not fail.
223     driver.check_for_spark_job_failure()
224 
225     # If there's no exception, execution results are in this queue.
226     results = result_queue.get_nowait()
227     return [results[index] for index in ranks_to_indices]
228
```
--- 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/spark/__init__.py b/horovod/spark/__init__.py
--- a/horovod/spark/__init__.py
+++ b/horovod/spark/__init__.py
@@ -13,6 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
+# Workaround for https://issues.apache.org/jira/browse/SPARK-22674
+# This fix also requires the user to make this same change at the top of their
+# training script before importing pyspark (on serialization).
+import collections
+collections.namedtuple.__hijack = 1
+
 import os
 import pyspark
 from six.moves import queue

{"golden_diff": "diff --git a/horovod/spark/__init__.py b/horovod/spark/__init__.py\n--- a/horovod/spark/__init__.py\n+++ b/horovod/spark/__init__.py\n@@ -13,6 +13,12 @@\n # limitations under the License.\n # ==============================================================================\n \n+# Workaround for https://issues.apache.org/jira/browse/SPARK-22674\n+# This fix also requires the user to make this same change at the top of their\n+# training script before importing pyspark (on serialization).\n+import collections\n+collections.namedtuple.__hijack = 1\n+\n import os\n import pyspark\n from six.moves import queue\n", "issue": "TF estimators and horovod.spark\n**Environment:**\r\n1. Framework: (TensorFlow, Keras, PyTorch, MXNet)\r\nTensorFlow\r\n2. Framework version:\r\n1.13.1\r\n3. Horovod version:\r\n0.16\r\n4. MPI version:\r\n3.0.0\r\n5. CUDA version:\r\nN/A\r\n6. NCCL version:\r\nN/A\r\n7. Python version:\r\n3.5\r\n8. OS and version:\r\nAmazon Linux AMI release 2018.03\r\n\r\n**Your question:**\r\nIs there any guidance on how to use horovod.spark with TF estimators? I've been able to train models using TF-Keras and horovod.spark in the style of the keras_spark_rossman.py example script, but I'm curious about how horovod.spark would work with Estimator's?\r\n\r\nMoreover, I think I'm just confused about how horovod.spark actually works - is it okay to instantiate an Estimator instance inside the input_fn for a horovod.spark.run() call, or will that mean each worker has a different starting point for the model?\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 os\nimport pyspark\nfrom six.moves import queue\nimport sys\nimport threading\n\nfrom horovod.spark.task import task_service\nfrom horovod.run.common.util import codec, env as env_util, safe_shell_exec, \\\n    timeout, host_hash, secret\nfrom horovod.run.common.util import settings as hvd_settings\nfrom horovod.spark.driver import driver_service, job_id\n\n\ndef _task_fn(index, driver_addresses, settings):\n    task = task_service.SparkTaskService(index, settings.key)\n    try:\n        driver_client = driver_service.SparkDriverClient(driver_addresses, settings.key, settings.verbose)\n        driver_client.register_task(index, task.addresses(), host_hash.host_hash())\n        task.wait_for_initial_registration(settings.timeout)\n        # Tasks ping each other in a circular fashion to determine interfaces reachable within\n        # the cluster.\n        next_task_index = (index + 1) % settings.num_proc\n        next_task_addresses = driver_client.all_task_addresses(next_task_index)\n        # We request interface matching to weed out all the NAT'ed interfaces.\n        next_task_client = \\\n            task_service.SparkTaskClient(next_task_index, next_task_addresses,\n                                         settings.key, settings.verbose,\n                                         match_intf=True)\n        driver_client.register_task_to_task_addresses(next_task_index, next_task_client.addresses())\n        task_indices_on_this_host = driver_client.task_host_hash_indices(\n            host_hash.host_hash())\n        if task_indices_on_this_host[0] == index:\n            # Task with first index will execute orted that will run mpirun_exec_fn for all tasks.\n            task.wait_for_command_start(settings.timeout)\n            task.wait_for_command_termination()\n        else:\n            # The rest of tasks need to wait for the first task to finish.\n            first_task_addresses = driver_client.all_task_addresses(task_indices_on_this_host[0])\n            first_task_client = \\\n                task_service.SparkTaskClient(task_indices_on_this_host[0],\n                                             first_task_addresses, settings.key,\n                                             settings.verbose)\n            first_task_client.wait_for_command_termination()\n        return task.fn_result()\n    finally:\n        task.shutdown()\n\n\ndef _make_mapper(driver_addresses, settings):\n    def _mapper(index, _):\n        yield _task_fn(index, driver_addresses, settings)\n    return _mapper\n\n\ndef _make_spark_thread(spark_context, spark_job_group, driver, result_queue,\n                       settings):\n    def run_spark():\n        try:\n            spark_context.setJobGroup(spark_job_group,\n                                      \"Horovod Spark Run\",\n                                      interruptOnCancel=True)\n            procs = spark_context.range(0, numSlices=settings.num_proc)\n            # We assume that folks caring about security will enable Spark RPC\n            # encryption, thus ensuring that key that is passed here remains\n            # secret.\n            result = procs.mapPartitionsWithIndex(_make_mapper(driver.addresses(), settings)).collect()\n            result_queue.put(result)\n        except:\n            driver.notify_spark_job_failed()\n            raise\n\n    spark_thread = threading.Thread(target=run_spark)\n    spark_thread.start()\n    return spark_thread\n\n\ndef run(fn, args=(), kwargs={}, num_proc=None, start_timeout=None, env=None,\n        stdout=None, stderr=None, verbose=1):\n    \"\"\"\n    Runs Horovod in Spark.  Runs `num_proc` processes executing `fn` using the same amount of Spark tasks.\n\n    Args:\n        fn: Function to run.\n        args: Arguments to pass to `fn`.\n        kwargs: Keyword arguments to pass to `fn`.\n        num_proc: Number of Horovod processes.  Defaults to `spark.default.parallelism`.\n        start_timeout: Timeout for Spark tasks to spawn, register and start running the code, in seconds.\n                       If not set, falls back to `HOROVOD_SPARK_START_TIMEOUT` environment variable value.\n                       If it is not set as well, defaults to 600 seconds.\n        env: Environment dictionary to use in Horovod run.  Defaults to `os.environ`.\n        stdout: Horovod stdout is redirected to this stream. Defaults to sys.stdout.\n        stderr: Horovod stderr is redirected to this stream. Defaults to sys.stderr.\n        verbose: Debug output verbosity (0-2). Defaults to 1.\n\n    Returns:\n        List of results returned by running `fn` on each rank.\n    \"\"\"\n\n    if start_timeout is None:\n        # Lookup default timeout from the environment variable.\n        start_timeout = int(os.getenv('HOROVOD_SPARK_START_TIMEOUT', '600'))\n\n    tmout = timeout.Timeout(start_timeout,\n                            message='Timed out waiting for {activity}. Please check that you have '\n                                    'enough resources to run all Horovod processes. Each Horovod '\n                                    'process runs in a Spark task. You may need to increase the '\n                                    'start_timeout parameter to a larger value if your Spark resources '\n                                    'are allocated on-demand.')\n    settings = hvd_settings.Settings(verbose=verbose,\n                                     key=secret.make_secret_key(),\n                                     timeout=tmout)\n\n    spark_context = pyspark.SparkContext._active_spark_context\n    if spark_context is None:\n        raise Exception('Could not find an active SparkContext, are you '\n                        'running in a PySpark session?')\n\n    if num_proc is None:\n        num_proc = spark_context.defaultParallelism\n        if settings.verbose >= 1:\n            print('Running %d processes (inferred from spark.default.parallelism)...' % num_proc)\n    else:\n        if settings.verbose >= 1:\n            print('Running %d processes...' % num_proc)\n    settings.num_proc = num_proc\n\n    result_queue = queue.Queue(1)\n\n    spark_job_group = 'horovod.spark.run.%d' % job_id.next_job_id()\n    driver = driver_service.SparkDriverService(settings.num_proc, fn, args, kwargs,\n                                               settings.key)\n    spark_thread = _make_spark_thread(spark_context, spark_job_group, driver,\n                                      result_queue, settings)\n    try:\n        driver.wait_for_initial_registration(settings.timeout)\n        if settings.verbose >= 2:\n            print('Initial Spark task registration is complete.')\n        task_clients = [\n            task_service.SparkTaskClient(index,\n                                         driver.task_addresses_for_driver(index),\n                                         settings.key, settings.verbose)\n            for index in range(settings.num_proc)]\n        for task_client in task_clients:\n            task_client.notify_initial_registration_complete()\n        driver.wait_for_task_to_task_address_updates(settings.timeout)\n        if settings.verbose >= 2:\n            print('Spark task-to-task address registration is complete.')\n\n        # Determine a set of common interfaces for task-to-task communication.\n        common_intfs = set(driver.task_addresses_for_tasks(0).keys())\n        for index in range(1, settings.num_proc):\n            common_intfs.intersection_update(driver.task_addresses_for_tasks(index).keys())\n        if not common_intfs:\n            raise Exception('Unable to find a set of common task-to-task communication interfaces: %s'\n                            % [(index, driver.task_addresses_for_tasks(index)) for index in range(settings.num_proc)])\n\n        # Determine the index grouping based on host hashes.\n        # Barrel shift until index 0 is in the first host.\n        host_hashes = list(driver.task_host_hash_indices().keys())\n        host_hashes.sort()\n        while 0 not in driver.task_host_hash_indices()[host_hashes[0]]:\n            host_hashes = host_hashes[1:] + host_hashes[:1]\n\n        ranks_to_indices = []\n        for host_hash in host_hashes:\n            ranks_to_indices += driver.task_host_hash_indices()[host_hash]\n        driver.set_ranks_to_indices(ranks_to_indices)\n\n        if env is None:\n            env = os.environ.copy()\n\n        # Pass secret key through the environment variables.\n        env[secret.HOROVOD_SECRET_KEY] = codec.dumps_base64(settings.key)\n\n        mpirun_command = (\n            'mpirun --allow-run-as-root --tag-output '\n            '-np {num_proc} -H {hosts} '\n            '-bind-to none -map-by slot '\n            '-mca pml ob1 -mca btl ^openib -mca btl_tcp_if_include {common_intfs} '\n            '-x NCCL_DEBUG=INFO -x NCCL_SOCKET_IFNAME={common_intfs} '\n            '{env} '  # expect a lot of environment variables\n            '-mca plm_rsh_agent \"{python} -m horovod.spark.driver.mpirun_rsh {encoded_driver_addresses} {settings}\" '\n            '{python} -m horovod.spark.task.mpirun_exec_fn {encoded_driver_addresses} {settings}'\n                .format(num_proc=settings.num_proc,\n                        hosts=','.join('%s:%d' % (host_hash, len(driver.task_host_hash_indices()[host_hash]))\n                                       for host_hash in host_hashes),\n                        common_intfs=','.join(common_intfs),\n                        env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)),\n                        python=sys.executable,\n                        encoded_driver_addresses=codec.dumps_base64(driver.addresses()),\n                        settings=codec.dumps_base64(settings)))\n        if settings.verbose >= 2:\n            print('+ %s' % mpirun_command)\n        exit_code = safe_shell_exec.execute(mpirun_command, env, stdout, stderr)\n        if exit_code != 0:\n            raise Exception('mpirun exited with code %d, see the error above.' % exit_code)\n    except:\n        # Terminate Spark job.\n        spark_context.cancelJobGroup(spark_job_group)\n\n        # Re-raise exception.\n        raise\n    finally:\n        spark_thread.join()\n        driver.shutdown()\n\n    # Make sure Spark Job did not fail.\n    driver.check_for_spark_job_failure()\n\n    # If there's no exception, execution results are in this queue.\n    results = result_queue.get_nowait()\n    return [results[index] for index in ranks_to_indices]\n", "path": "horovod/spark/__init__.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\n# Workaround for https://issues.apache.org/jira/browse/SPARK-22674\n# This fix also requires the user to make this same change at the top of their\n# training script before importing pyspark (on serialization).\nimport collections\ncollections.namedtuple.__hijack = 1\n\nimport os\nimport pyspark\nfrom six.moves import queue\nimport sys\nimport threading\n\nfrom horovod.spark.task import task_service\nfrom horovod.run.common.util import codec, env as env_util, safe_shell_exec, \\\n    timeout, host_hash, secret\nfrom horovod.run.common.util import settings as hvd_settings\nfrom horovod.spark.driver import driver_service, job_id\n\n\ndef _task_fn(index, driver_addresses, settings):\n    task = task_service.SparkTaskService(index, settings.key)\n    try:\n        driver_client = driver_service.SparkDriverClient(driver_addresses, settings.key, settings.verbose)\n        driver_client.register_task(index, task.addresses(), host_hash.host_hash())\n        task.wait_for_initial_registration(settings.timeout)\n        # Tasks ping each other in a circular fashion to determine interfaces reachable within\n        # the cluster.\n        next_task_index = (index + 1) % settings.num_proc\n        next_task_addresses = driver_client.all_task_addresses(next_task_index)\n        # We request interface matching to weed out all the NAT'ed interfaces.\n        next_task_client = \\\n            task_service.SparkTaskClient(next_task_index, next_task_addresses,\n                                         settings.key, settings.verbose,\n                                         match_intf=True)\n        driver_client.register_task_to_task_addresses(next_task_index, next_task_client.addresses())\n        task_indices_on_this_host = driver_client.task_host_hash_indices(\n            host_hash.host_hash())\n        if task_indices_on_this_host[0] == index:\n            # Task with first index will execute orted that will run mpirun_exec_fn for all tasks.\n            task.wait_for_command_start(settings.timeout)\n            task.wait_for_command_termination()\n        else:\n            # The rest of tasks need to wait for the first task to finish.\n            first_task_addresses = driver_client.all_task_addresses(task_indices_on_this_host[0])\n            first_task_client = \\\n                task_service.SparkTaskClient(task_indices_on_this_host[0],\n                                             first_task_addresses, settings.key,\n                                             settings.verbose)\n            first_task_client.wait_for_command_termination()\n        return task.fn_result()\n    finally:\n        task.shutdown()\n\n\ndef _make_mapper(driver_addresses, settings):\n    def _mapper(index, _):\n        yield _task_fn(index, driver_addresses, settings)\n    return _mapper\n\n\ndef _make_spark_thread(spark_context, spark_job_group, driver, result_queue,\n                       settings):\n    def run_spark():\n        try:\n            spark_context.setJobGroup(spark_job_group,\n                                      \"Horovod Spark Run\",\n                                      interruptOnCancel=True)\n            procs = spark_context.range(0, numSlices=settings.num_proc)\n            # We assume that folks caring about security will enable Spark RPC\n            # encryption, thus ensuring that key that is passed here remains\n            # secret.\n            result = procs.mapPartitionsWithIndex(_make_mapper(driver.addresses(), settings)).collect()\n            result_queue.put(result)\n        except:\n            driver.notify_spark_job_failed()\n            raise\n\n    spark_thread = threading.Thread(target=run_spark)\n    spark_thread.start()\n    return spark_thread\n\n\ndef run(fn, args=(), kwargs={}, num_proc=None, start_timeout=None, env=None,\n        stdout=None, stderr=None, verbose=1):\n    \"\"\"\n    Runs Horovod in Spark.  Runs `num_proc` processes executing `fn` using the same amount of Spark tasks.\n\n    Args:\n        fn: Function to run.\n        args: Arguments to pass to `fn`.\n        kwargs: Keyword arguments to pass to `fn`.\n        num_proc: Number of Horovod processes.  Defaults to `spark.default.parallelism`.\n        start_timeout: Timeout for Spark tasks to spawn, register and start running the code, in seconds.\n                       If not set, falls back to `HOROVOD_SPARK_START_TIMEOUT` environment variable value.\n                       If it is not set as well, defaults to 600 seconds.\n        env: Environment dictionary to use in Horovod run.  Defaults to `os.environ`.\n        stdout: Horovod stdout is redirected to this stream. Defaults to sys.stdout.\n        stderr: Horovod stderr is redirected to this stream. Defaults to sys.stderr.\n        verbose: Debug output verbosity (0-2). Defaults to 1.\n\n    Returns:\n        List of results returned by running `fn` on each rank.\n    \"\"\"\n\n    if start_timeout is None:\n        # Lookup default timeout from the environment variable.\n        start_timeout = int(os.getenv('HOROVOD_SPARK_START_TIMEOUT', '600'))\n\n    tmout = timeout.Timeout(start_timeout,\n                            message='Timed out waiting for {activity}. Please check that you have '\n                                    'enough resources to run all Horovod processes. Each Horovod '\n                                    'process runs in a Spark task. You may need to increase the '\n                                    'start_timeout parameter to a larger value if your Spark resources '\n                                    'are allocated on-demand.')\n    settings = hvd_settings.Settings(verbose=verbose,\n                                     key=secret.make_secret_key(),\n                                     timeout=tmout)\n\n    spark_context = pyspark.SparkContext._active_spark_context\n    if spark_context is None:\n        raise Exception('Could not find an active SparkContext, are you '\n                        'running in a PySpark session?')\n\n    if num_proc is None:\n        num_proc = spark_context.defaultParallelism\n        if settings.verbose >= 1:\n            print('Running %d processes (inferred from spark.default.parallelism)...' % num_proc)\n    else:\n        if settings.verbose >= 1:\n            print('Running %d processes...' % num_proc)\n    settings.num_proc = num_proc\n\n    result_queue = queue.Queue(1)\n\n    spark_job_group = 'horovod.spark.run.%d' % job_id.next_job_id()\n    driver = driver_service.SparkDriverService(settings.num_proc, fn, args, kwargs,\n                                               settings.key)\n    spark_thread = _make_spark_thread(spark_context, spark_job_group, driver,\n                                      result_queue, settings)\n    try:\n        driver.wait_for_initial_registration(settings.timeout)\n        if settings.verbose >= 2:\n            print('Initial Spark task registration is complete.')\n        task_clients = [\n            task_service.SparkTaskClient(index,\n                                         driver.task_addresses_for_driver(index),\n                                         settings.key, settings.verbose)\n            for index in range(settings.num_proc)]\n        for task_client in task_clients:\n            task_client.notify_initial_registration_complete()\n        driver.wait_for_task_to_task_address_updates(settings.timeout)\n        if settings.verbose >= 2:\n            print('Spark task-to-task address registration is complete.')\n\n        # Determine a set of common interfaces for task-to-task communication.\n        common_intfs = set(driver.task_addresses_for_tasks(0).keys())\n        for index in range(1, settings.num_proc):\n            common_intfs.intersection_update(driver.task_addresses_for_tasks(index).keys())\n        if not common_intfs:\n            raise Exception('Unable to find a set of common task-to-task communication interfaces: %s'\n                            % [(index, driver.task_addresses_for_tasks(index)) for index in range(settings.num_proc)])\n\n        # Determine the index grouping based on host hashes.\n        # Barrel shift until index 0 is in the first host.\n        host_hashes = list(driver.task_host_hash_indices().keys())\n        host_hashes.sort()\n        while 0 not in driver.task_host_hash_indices()[host_hashes[0]]:\n            host_hashes = host_hashes[1:] + host_hashes[:1]\n\n        ranks_to_indices = []\n        for host_hash in host_hashes:\n            ranks_to_indices += driver.task_host_hash_indices()[host_hash]\n        driver.set_ranks_to_indices(ranks_to_indices)\n\n        if env is None:\n            env = os.environ.copy()\n\n        # Pass secret key through the environment variables.\n        env[secret.HOROVOD_SECRET_KEY] = codec.dumps_base64(settings.key)\n\n        mpirun_command = (\n            'mpirun --allow-run-as-root --tag-output '\n            '-np {num_proc} -H {hosts} '\n            '-bind-to none -map-by slot '\n            '-mca pml ob1 -mca btl ^openib -mca btl_tcp_if_include {common_intfs} '\n            '-x NCCL_DEBUG=INFO -x NCCL_SOCKET_IFNAME={common_intfs} '\n            '{env} '  # expect a lot of environment variables\n            '-mca plm_rsh_agent \"{python} -m horovod.spark.driver.mpirun_rsh {encoded_driver_addresses} {settings}\" '\n            '{python} -m horovod.spark.task.mpirun_exec_fn {encoded_driver_addresses} {settings}'\n                .format(num_proc=settings.num_proc,\n                        hosts=','.join('%s:%d' % (host_hash, len(driver.task_host_hash_indices()[host_hash]))\n                                       for host_hash in host_hashes),\n                        common_intfs=','.join(common_intfs),\n                        env=' '.join('-x %s' % key for key in env.keys() if env_util.is_exportable(key)),\n                        python=sys.executable,\n                        encoded_driver_addresses=codec.dumps_base64(driver.addresses()),\n                        settings=codec.dumps_base64(settings)))\n        if settings.verbose >= 2:\n            print('+ %s' % mpirun_command)\n        exit_code = safe_shell_exec.execute(mpirun_command, env, stdout, stderr)\n        if exit_code != 0:\n            raise Exception('mpirun exited with code %d, see the error above.' % exit_code)\n    except:\n        # Terminate Spark job.\n        spark_context.cancelJobGroup(spark_job_group)\n\n        # Re-raise exception.\n        raise\n    finally:\n        spark_thread.join()\n        driver.shutdown()\n\n    # Make sure Spark Job did not fail.\n    driver.check_for_spark_job_failure()\n\n    # If there's no exception, execution results are in this queue.\n    results = result_queue.get_nowait()\n    return [results[index] for index in ranks_to_indices]\n", "path": "horovod/spark/__init__.py"}]}