def test_timeout(self):
with spark('test_timeout'):
with pytest.raises(
Exception,
match='^Timed out waiting for Spark tasks to start.'):
horovod.spark.run(None,
num_proc=4,
start_timeout=5,
env={'PATH': os.environ.get('PATH')},
verbose=0)
def test_spark_run_func_with_non_zero_exit(self):
run_func = MagicMock(return_value=1)
def fn():
return 1
with spark('test_spark_run_func', cores=4):
with pytest.raises(Exception,
match='^mpirun failed with exit code 1$') as e:
horovod.spark.run(fn, verbose=0, run_func=run_func)
def test_mpirun_not_found(self):
start = time.time()
with spark('test_mpirun_not_found'):
with pytest.raises(Exception,
match='^mpirun failed with exit code 127$'):
horovod.spark.run(None,
env={'PATH': '/nonexistent'},
verbose=0)
self.assertLessEqual(time.time() - start, 10,
'Failure propagation took too long')
def test_happy_run(self):
def fn():
hvd.init()
res = hvd.allgather(torch.tensor([hvd.rank()])).tolist()
return res, hvd.rank()
with spark('test_happy_run'):
res = horovod.spark.run(fn,
env={'PATH': os.environ.get('PATH')},
verbose=0)
self.assertListEqual([([0, 1], 0), ([0, 1], 1)], res)
def do_test_spark_run_func(self,
args=(),
kwargs={},
num_proc=1,
extra_mpi_args=None,
env={},
stdout=None,
stderr=None,
verbose=0,
cores=2,
expected_np=1,
expected_env=''):
def fn():
return 1
run_func = MagicMock(return_value=0)
with spark('test_spark_run_func', cores=cores):
with pytest.raises(Exception) as e:
# we need to timeout horovod because our mocked run_func will block spark otherwise
# this raises above exception, but allows us to catch run_func arguments
horovod.spark.run(fn,
args=args,
kwargs=kwargs,
num_proc=num_proc,
start_timeout=1,
extra_mpi_args=extra_mpi_args,
env=env,
stdout=stdout,
stderr=stderr,
verbose=verbose,
run_func=run_func)
self.assertFalse(
str(e.value).startswith(
'Timed out waiting for Spark tasks to start.'),
'Spark timed out before mpi_run was called, test setup is broken.')
self.assertEqual(str(e.value),
'Spark job has failed, see the error above.')
mpi_flags = _get_mpi_implementation_flags()
self.assertIsNotNone(mpi_flags)
expected_command = (
'mpirun '
'--allow-run-as-root --tag-output '
'-np {expected_np} -H [^ ]+ '
'-bind-to none -map-by slot '
'{mpi_flags} '
'-mca btl_tcp_if_include [^ ]+ -x NCCL_SOCKET_IFNAME=[^ ]+ '
'-x _HOROVOD_SECRET_KEY {expected_env}'
'{extra_mpi_args} '
'-x NCCL_DEBUG=INFO '
r'-mca plm_rsh_agent "[^"]+python[\d]* -m horovod.spark.driver.mpirun_rsh [^ ]+ [^ ]+" '
r'[^"]+python[\d]* -m horovod.spark.task.mpirun_exec_fn [^ ]+ [^ ]+'
.format(expected_np=expected_np,
expected_env=expected_env + ' ' if expected_env else '',
mpi_flags=' '.join(mpi_flags),
extra_mpi_args=extra_mpi_args if extra_mpi_args else ''))
run_func.assert_called_once()
run_func_args, run_func_kwargs = run_func.call_args
actual_command = run_func_kwargs.get('command')
actual_env = run_func_kwargs.get('env')
actual_stdout = run_func_kwargs.get('stdout')
actual_stderr = run_func_kwargs.get('stderr')
actual_secret = actual_env.pop('_HOROVOD_SECRET_KEY', None)
# for better comparison replace sections in actual_command that change across runs / hosts
for replacement in (
'-H [^ ]+', '-mca btl_tcp_if_include [^ ]+',
'-x NCCL_SOCKET_IFNAME=[^ ]+', r'"[^"]+python[\d]*',
r' [^"]+python[\d]*',
'-m horovod.spark.driver.mpirun_rsh [^ ]+ [^ ]+"',
'-m horovod.spark.task.mpirun_exec_fn [^ ]+ [^ ]+'):
actual_command = re.sub(replacement, replacement, actual_command,
1)
self.assertEqual(run_func_args, ())
self.assertEqual(actual_command, expected_command)
if env:
self.assertEqual(actual_env, env)
else:
self.assertIsNotNone(actual_env)
self.assertIsNotNone(actual_secret)
self.assertTrue(len(actual_secret) > 0)
self.assertEqual(actual_stdout, stdout)
self.assertEqual(actual_stderr, stderr)