def test_infer_no_job_id_error(mock_make_factory, mock_job):
# The `infer_backend_factory` should give a RuntimeError
# if there's no CIVIS_JOB_ID in the environment.
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_job
with mock.patch.dict('os.environ', {}, clear=True):
with pytest.raises(RuntimeError):
civis.parallel.infer_backend_factory(client=mock_client)
def test_outputs_succeeded(self):
poller = _create_poller_mock("succeeded")
mock_client = create_client_mock()
expected_return = [{'test': 'test_result'}]
mock_client.jobs.list_runs_outputs.return_value = expected_return
result = CivisFuture(poller, (1, 2), client=mock_client)
assert result.outputs() == expected_return
def test_future_job_id_run_id(poller_args, expected_job_id, expected_run_id):
result = CivisFuture(
poller=lambda x: x,
poller_args=poller_args,
client=create_client_mock(),
)
assert result.job_id == expected_job_id
assert result.run_id == expected_run_id
def test_future_job_id_run_id(poller_args, expected_job_id, expected_run_id):
result = CivisFuture(
poller=_create_poller_mock("succeeded"),
poller_args=poller_args,
client=create_client_mock(),
)
assert result.job_id == expected_job_id
assert result.run_id == expected_run_id
def test_container_future_job_id_run_id():
job_id, run_id = 123, 456
result = ContainerFuture(
job_id=job_id,
run_id=run_id,
client=create_client_mock(),
)
assert result.job_id == job_id
assert result.run_id == run_id
def test_infer_in_child_job(mock_make_factory, mock_child_job):
# Verify that infer_backend_factory doesn't include CIVIS_PARENT_JOB_ID and
# CIVIS_PARENT_RUN_ID since those will be automatically added later.
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_child_job
mock_env = {'CIVIS_JOB_ID': "test_job", 'CIVIS_RUN_ID': "test_run"}
with mock.patch.dict('os.environ', mock_env):
civis.parallel.infer_backend_factory(client=mock_client)
assert mock_make_factory.call_args[1]['params'] == [{'name': 'spam'}]
def test_outputs_succeeded(self):
poller = mock.Mock()
api_result = mock.Mock()
api_result.state = 'succeeded'
mock_client = create_client_mock()
expected_return = [{'test': 'test_result'}]
mock_client.jobs.list_runs_outputs.return_value = expected_return
result = CivisFuture(poller, (1, 2), client=mock_client)
result._set_api_result(api_result)
assert result.outputs() == expected_return
def _test_retries_helper(num_failures, max_submit_retries, should_fail,
from_template_id, mock_file_to_civis, mock_result_cls,
mock_custom_exec_cls, mock_executor_cls):
mock_file_to_civis.return_value = 0
mock_result_cls.return_value.get.return_value = [123]
# A function to raise fake API errors the first num_failures times it is
# called.
counter = {'n_failed': 0}
def mock_submit(fn='', *args, **kwargs):
if counter['n_failed'] < num_failures:
counter['n_failed'] += 1
raise CivisAPIError(mock.MagicMock())
else:
return mock.MagicMock(spec=ContainerFuture)
mock_custom_exec_cls.return_value.submit.side_effect = mock_submit
mock_executor_cls.return_value.submit.side_effect = mock_submit
if from_template_id:
factory = civis.parallel.make_backend_template_factory(
from_template_id=from_template_id,
max_submit_retries=max_submit_retries,
client=create_client_mock())
else:
factory = civis.parallel.make_backend_factory(
max_submit_retries=max_submit_retries, client=create_client_mock())
register_parallel_backend('civis', factory)
with parallel_backend('civis'):
# NB: joblib >v0.11 relies on callbacks from the result object to
# decide when it's done consuming inputs. We've mocked the result
# object here, so Parallel must be called either with n_jobs=1 or
# pre_dispatch='all' to consume the inputs all at once.
parallel = Parallel(n_jobs=1, pre_dispatch='n_jobs')
if should_fail:
with pytest.raises(civis.parallel.JobSubmissionError):
parallel(delayed(sqrt)(i**2) for i in range(3))
else:
parallel(delayed(sqrt)(i**2) for i in range(3))
def test_infer_update_resources(mock_make_factory, mock_job):
# Verify that users can modify requested resources for jobs.
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_job
with mock.patch.dict('os.environ', {
'CIVIS_JOB_ID': "test_job",
'CIVIS_RUN_ID': "test_run"
}):
civis.parallel.infer_backend_factory(client=mock_client,
required_resources={'cpu': -11})
assert mock_make_factory.call_args[1]['required_resources'] == \
{'cpu': -11}
def test_infer_update_args(mock_make_factory, mock_job):
# Verify that users can modify the existing job's
# arguments for sub-processes.
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_job
with mock.patch.dict('os.environ', {
'CIVIS_JOB_ID': "test_job",
'CIVIS_RUN_ID': "test_run"
}):
civis.parallel.infer_backend_factory(client=mock_client,
arguments={'foo': 'bar'})
assert mock_make_factory.call_args[1]['arguments'] == \
{'spam': 'eggs', 'foo': 'bar'}
def test_infer_extra_param(mock_make_factory, mock_job):
# Test adding a new parameter and keeping
# the existing parameter unchanged.
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_job
new_params = [{'name': 'foo', 'type': 'bar'}]
with mock.patch.dict('os.environ', {
'CIVIS_JOB_ID': "test_job",
'CIVIS_RUN_ID': "test_run"
}):
civis.parallel.infer_backend_factory(client=mock_client,
params=new_params)
assert mock_make_factory.call_args[1]['params'] == \
[{'name': 'spam'}, {'name': 'foo', 'type': 'bar'}]
def test_result_exception_no_result():
# If the job errored but didn't write an output, we should get
# a generic TransportableException back.
callback = mock.MagicMock()
# Passing the client mock as an argument instead of globally
# patching the client tests that the _CivisBackendResult
# uses the client object on the input CivisFuture.
mock_client = create_client_mock()
mock_client.scripts.list_containers_runs_outputs.return_value = []
fut = ContainerFuture(1, 2, client=mock_client)
fut._set_api_exception(Response({'state': 'failed'}))
res = civis.parallel._CivisBackendResult(fut, callback)
with pytest.raises(TransportableException) as exc:
res.get()
assert "{'state': 'failed'}" in str(exc.value)
assert callback.call_count == 0
def test_infer_new_params(mock_make_factory, mock_job):
# Test overwriting existing job parameters with new parameters
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_job
new_params = [{
'name': 'spam',
'type': 'fun'
}, {
'name': 'foo',
'type': 'bar'
}]
with mock.patch.dict('os.environ', {
'CIVIS_JOB_ID': "test_job",
'CIVIS_RUN_ID': "test_run"
}):
civis.parallel.infer_backend_factory(client=mock_client,
params=new_params)
assert mock_make_factory.call_args[1]['params'] == new_params
def test_infer(mock_make_factory, mock_job):
# Verify that `infer_backend_factory` passes through
# the expected arguments to `make_backend_factory`.
mock_client = create_client_mock()
mock_client.scripts.get_containers.return_value = mock_job
with mock.patch.dict('os.environ', {
'CIVIS_JOB_ID': "test_job",
'CIVIS_RUN_ID': "test_run"
}):
civis.parallel.infer_backend_factory(client=mock_client)
expected = dict(mock_job)
del expected['from_template_id']
del expected['id']
mock_make_factory.assert_called_once_with(client=mock_client,
setup_cmd=None,
polling_interval=None,
max_submit_retries=0,
max_job_retries=0,
hidden=True,
remote_backend='sequential',
**expected)
def test_template_submit(mock_file, mock_result, mock_pool):
# Verify that creating child jobs from a template looks like we expect
file_id = 17
mock_client = create_client_mock()
mock_file.return_value = file_id
factory = civis.parallel.make_backend_template_factory(
from_template_id=1234, client=mock_client)
n_calls = 3
register_parallel_backend('civis', factory)
with parallel_backend('civis'):
# NB: joblib >v0.11 relies on callbacks from the result object to
# decide when it's done consuming inputs. We've mocked the result
# object here, so Parallel must be called either with n_jobs=1 or
# pre_dispatch='all' to consume the inputs all at once.
parallel = Parallel(n_jobs=1, pre_dispatch='n_jobs')
parallel(delayed(sqrt)(i**2) for i in range(n_calls))
assert mock_file.call_count == 3, "Upload 3 functions to run"
assert mock_pool().submit.call_count == n_calls, "Run 3 functions"
for this_call in mock_pool().submit.call_args_list:
assert this_call == mock.call(JOBLIB_FUNC_FILE_ID=file_id)
assert mock_result.call_count == 3, "Create 3 results"
def test_infer_from_custom_job(mock_make_factory, mock_job):
# Test that `infer_backend_factory` can find needed
# parameters if it's run inside a custom job created
# from a template.
mock_client = create_client_mock()
mock_custom = Response(
dict(from_template_id=999,
id=42,
required_resources=None,
params=[{
'name': 'spam'
}],
arguments={'spam': 'eggs'},
docker_image_name='image_name',
docker_image_tag='tag',
repo_http_uri='cabbage',
repo_ref='servant'))
mock_template = Response(dict(id=999, script_id=171))
def _get_container(job_id):
if int(job_id) == 42:
return mock_custom
elif int(job_id) == 171:
return mock_job
else:
raise ValueError("Got job_id {}".format(job_id))
mock_client.scripts.get_containers.side_effect = _get_container
mock_client.templates.get_scripts.return_value = mock_template
with mock.patch.dict('os.environ', {
'CIVIS_JOB_ID': "42",
'CIVIS_RUN_ID': "test_run"
}):
civis.parallel.infer_backend_factory(client=mock_client)
# We should have called `get_containers` twice now -- once for
# the container we're running in, and a second time for the
# container which backs the template this job was created from.
# The backing script has settings which aren't visible from
# the container which was created from it.
assert mock_client.scripts.get_containers.call_count == 2
mock_client.templates.get_scripts.assert_called_once_with(999)
expected_kwargs = {
'required_resources': {
'cpu': 11
},
'params': [{
'name': 'spam'
}],
'arguments': {
'spam': 'eggs'
},
'client': mock.ANY,
'polling_interval': mock.ANY,
'setup_cmd': None,
'max_submit_retries': mock.ANY,
'max_job_retries': mock.ANY,
'hidden': True,
'remote_backend': 'sequential'
}
for key in civis.parallel.KEYS_TO_INFER:
expected_kwargs[key] = mock_job[key]
mock_make_factory.assert_called_once_with(**expected_kwargs)
def test__get_template_ids_all_versions():
m_client = create_client_mock()
m_client.aliases.list.return_value = TEST_TEMPLATE_ID_ALIAS_OBJECTS
actual_template_ids = _model._get_template_ids_all_versions(m_client)
expected_template_ids = TEST_TEMPLATE_IDS
assert actual_template_ids == expected_template_ids
