def testRunUnexpectedStopFromRegularVm_WithPreemptibles(self, mock_dsub_call):
self._argv.extend([
'--jobs_to_run',
'postprocess_variants',
'--preemptible',
'--max_preemptible_tries',
'1',
'--max_non_preemptible_tries',
'2',
])
preemptible_exception = dsub_errors.JobExecutionError(
'error', ['Error in job - code 10: 14: VM stopped unexpectedly'])
unexpected_stop_exception = dsub_errors.JobExecutionError(
'error', ['Error in job - code 10: 13: VM stopped unexpectedly'])
mock_dsub_call.side_effect = [
preemptible_exception, unexpected_stop_exception,
unexpected_stop_exception
]
try:
gcp_deepvariant_runner.run(self._argv)
self.fail(
'Too many unexpected stops from regular VMs should throw exception.')
except RuntimeError:
pass
mock_dsub_call.assert_has_calls([
mock.call(_HasAllOf('postprocess_variants', '--preemptible')),
mock.call(_HasAllOf('postprocess_variants')),
mock.call(_HasAllOf('postprocess_variants'))
])
def testRunPipeline(self, mock_pool, mock_dsub_call):
mock_apply_async = mock_pool.return_value.apply_async
mock_apply_async.return_value = None
self._argv.extend(
['--make_examples_workers', '1', '--call_variants_workers', '1'])
gcp_deepvariant_runner.run(self._argv)
mock_apply_async.assert_has_calls([
mock.call(
func=mock.ANY,
args=(_HasAllOf('make_examples', 'gcr.io/dockerimage',
'INPUT_BAM=gs://bam', 'INPUT_BAI=gs://bam.bai',
'INPUT_REF=gs://ref', 'INPUT_REF_FAI=gs://ref.fai',
'EXAMPLES=gs://staging/examples/0'), mock.ANY, 0)),
mock.call(
func=mock.ANY,
args=(_HasAllOf(
'call_variants', 'gcr.io/dockerimage', 'MODEL=gs://model',
'EXAMPLES=gs://staging/examples/0',
'CALLED_VARIANTS=gs://staging/called_variants'), mock.ANY, 0)),
],)
mock_dsub_call.assert_called_once_with(
_HasAllOf('postprocess_variants', 'gcr.io/dockerimage',
'CALLED_VARIANTS=gs://staging/called_variants',
'OUTFILE=gs://output.vcf'))
def testRunWithPreemptibles_NonPreemptibleFailure(self, mock_dsub_call):
self._argv.extend([
'--jobs_to_run',
'postprocess_variants',
'--preemptible',
'--max_preemptible_tries',
'2',
])
preemptible_exception = dsub_errors.JobExecutionError(
'error', ['Error in job - code 10: 13: VM stopped unexpectedly'])
# Note that the error code changed to '130' from '13'.
non_preemptible_exception = dsub_errors.JobExecutionError(
'error', ['Error in job - code 10: 130: VM stopped unexpectedly'])
mock_dsub_call.side_effect = [
preemptible_exception, non_preemptible_exception
]
try:
gcp_deepvariant_runner.run(self._argv)
self.fail('Non-preemptible failures should throw exception.')
except RuntimeError:
pass
# Two preemptible tries should have still happened.
mock_dsub_call.assert_has_calls([
mock.call(_HasAllOf('postprocess_variants', '--preemptible')),
mock.call(_HasAllOf('postprocess_variants', '--preemptible'))
])
def testRunPostProcessVariants(self, mock_dsub_call):
self._argv.extend([
'--jobs_to_run',
'postprocess_variants',
'--shards',
'15',
'--gpu', # GPU should not have any effect.
'--docker_image_gpu',
'gcr.io/dockerimage_gpu',
])
gcp_deepvariant_runner.run(self._argv)
mock_dsub_call.assert_called_once_with(
_HasAllOf('postprocess_variants', 'gcr.io/dockerimage',
'CALLED_VARIANTS=gs://staging/called_variants',
'INPUT_REF=gs://ref', 'INPUT_REF_FAI=gs://ref.fai',
'OUTFILE=gs://output.vcf'))
def testRunUnexpectedStopFromRegularVm(self, mock_dsub_call):
self._argv.extend([
'--jobs_to_run',
'postprocess_variants',
'--max_non_preemptible_tries',
'3',
])
unexpected_stop_exception = dsub_errors.JobExecutionError(
'error', ['Error in job - code 10: 13: VM stopped unexpectedly'])
mock_dsub_call.side_effect = [
unexpected_stop_exception, unexpected_stop_exception, None
]
gcp_deepvariant_runner.run(self._argv)
mock_dsub_call.assert_has_calls([
mock.call(_HasAllOf('postprocess_variants')),
mock.call(_HasAllOf('postprocess_variants')),
mock.call(_HasAllOf('postprocess_variants'))
])
def testRunCallVariants_GPU(self, mock_pool):
mock_apply_async = mock_pool.return_value.apply_async
mock_apply_async.return_value = None
self._argv.extend([
'--jobs_to_run',
'call_variants',
'--call_variants_workers',
'3',
'--call_variants_cores_per_worker',
'5',
'--call_variants_cores_per_shard',
'5',
'--shards',
'15',
'--gpu',
'--docker_image_gpu',
'gcr.io/dockerimage_gpu',
])
gcp_deepvariant_runner.run(self._argv)
mock_apply_async.assert_has_calls(
[
mock.call(
func=mock.ANY,
args=(_HasAllOf('call_variants', 'gcr.io/dockerimage_gpu',
'nvidia-tesla-k80', 'SHARD_START_INDEX=0',
'SHARD_END_INDEX=4', 'CONCURRENT_JOBS=1'),
mock.ANY, 0)),
mock.call(
func=mock.ANY,
args=(_HasAllOf('call_variants', 'gcr.io/dockerimage_gpu',
'nvidia-tesla-k80', 'SHARD_START_INDEX=5',
'SHARD_END_INDEX=9', 'CONCURRENT_JOBS=1'),
mock.ANY, 1)),
mock.call(
func=mock.ANY,
args=(_HasAllOf('call_variants', 'gcr.io/dockerimage_gpu',
'nvidia-tesla-k80', 'SHARD_START_INDEX=10',
'SHARD_END_INDEX=14', 'CONCURRENT_JOBS=1'),
mock.ANY, 2)),
],
any_order=True,
)
def testRunWithPreemptibles(self, mock_dsub_call):
self._argv.extend([
'--jobs_to_run',
'postprocess_variants',
'--preemptible',
'--max_preemptible_tries',
'2',
])
preemptible_exception = dsub_errors.JobExecutionError(
'error', ['Error in job - code 10: 14: VM stopped unexpectedly'])
mock_dsub_call.side_effect = [
preemptible_exception, preemptible_exception, None
]
gcp_deepvariant_runner.run(self._argv)
mock_dsub_call.assert_has_calls([
mock.call(_HasAllOf('postprocess_variants', '--preemptible')),
mock.call(_HasAllOf('postprocess_variants', '--preemptible')),
mock.call(_HasAllOf('postprocess_variants'))
])
def testRunMakeExamples(self, mock_pool):
mock_apply_async = mock_pool.return_value.apply_async
mock_apply_async.return_value = None
self._argv.extend([
'--jobs_to_run',
'make_examples',
'--make_examples_workers',
'3',
'--shards',
'15',
'--gpu', # GPU should not have any effect.
'--docker_image_gpu',
'image_gpu',
'--job_name_prefix',
'prefix_',
])
gcp_deepvariant_runner.run(self._argv)
mock_apply_async.assert_has_calls(
[
mock.call(
func=mock.ANY,
args=(_HasAllOf('prefix_make_examples', 'gcr.io/dockerimage',
'SHARD_START_INDEX=0', 'SHARD_END_INDEX=4',
'EXAMPLES=gs://staging/examples'), mock.ANY,
0)),
mock.call(
func=mock.ANY,
args=(_HasAllOf('prefix_make_examples', 'gcr.io/dockerimage',
'SHARD_START_INDEX=5', 'SHARD_END_INDEX=9'),
mock.ANY, 1)),
mock.call(
func=mock.ANY,
args=(_HasAllOf('prefix_make_examples', 'gcr.io/dockerimage',
'SHARD_START_INDEX=10', 'SHARD_END_INDEX=14'),
mock.ANY, 2)),
],
any_order=True,
)
def testRunCallVariants(self, mock_pool):
mock_apply_async = mock_pool.return_value.apply_async
mock_apply_async.return_value = None
self._argv.extend([
'--jobs_to_run',
'call_variants',
'--call_variants_workers',
'3',
'--call_variants_cores_per_worker',
'5',
'--call_variants_cores_per_shard',
'2',
'--shards',
'15',
])
gcp_deepvariant_runner.run(self._argv)
mock_apply_async.assert_has_calls(
[
mock.call(
func=mock.ANY,
args=(_HasAllOf('call_variants', 'gcr.io/dockerimage',
'SHARD_START_INDEX=0', 'SHARD_END_INDEX=4',
'CONCURRENT_JOBS=2'), mock.ANY, 0)),
mock.call(
func=mock.ANY,
args=(_HasAllOf('call_variants', 'gcr.io/dockerimage',
'SHARD_START_INDEX=5', 'SHARD_END_INDEX=9',
'CONCURRENT_JOBS=2'), mock.ANY, 1)),
mock.call(func=mock.ANY,
args=(_HasAllOf(
'call_variants', 'gcr.io/dockerimage',
'SHARD_START_INDEX=10', 'SHARD_END_INDEX=14',
'CONCURRENT_JOBS=2'), mock.ANY, 2)),
],
any_order=True,
)
def testRunPostProcessVariants(self, mock_can_write_to_bucket,
mock_obj_exist, mock_run_job):
mock_obj_exist.return_value = True
mock_can_write_to_bucket.return_value = True
self._argv.extend([
'--jobs_to_run',
'postprocess_variants',
'--shards',
'15',
'--gpu', # GPU should not have any effect.
'--gvcf_outfile',
'gvcf-folder-path',
'--docker_image_gpu',
'gcr.io/dockerimage_gpu',
])
gcp_deepvariant_runner.run(self._argv)
mock_run_job.assert_called_once_with(
_HasAllOf('postprocess_variants', 'gcr.io/dockerimage',
'CALLED_VARIANTS=gs://bucket/staging/called_variants/*',
'INPUT_REF=gs://bucket/ref', 'SHARDS=15',
'CALL_VARIANTS_SHARDS=1',
'INPUT_REF_FAI=gs://bucket/ref.fai',
'OUTFILE=gs://bucket/output.vcf'),
'gs://bucket/staging/logs/postprocess_variants')
def testRunCallVariants_TPU(self, mock_can_write_to_bucket, mock_obj_exist,
mock_cluster_exists, mock_deploy_pod,
mock_init):
mock_obj_exist.return_value = True
mock_can_write_to_bucket.return_value = True
mock_cluster_exists.return_value = True
self._argv.extend([
'--jobs_to_run',
'call_variants',
'--call_variants_workers',
'1',
'--shards',
'15',
'--tpu',
'--gke_cluster_name',
'foo-cluster',
'--gke_cluster_zone',
'us-central1-c',
'--docker_image',
'gcr.io/dockerimage',
])
gcp_deepvariant_runner.run(self._argv)
mock_init.assert_has_calls([
mock.call('foo-cluster',
None,
'us-central1-c',
create_if_not_exist=False),
mock.call('foo-cluster', None, 'us-central1-c')
])
self.assertEqual(mock_init.call_count, 2)
mock_deploy_pod.assert_called_once_with(
pod_config=mock.ANY,
pod_name=AnyStringWith('deepvariant-'),
retries=1,
wait=True)
def testRunMakeExamples(self, mock_pool):
mock_apply_async = mock_pool.return_value.apply_async
mock_apply_async.return_value = None
self._argv.extend([
'--jobs_to_run',
'make_examples',
'--make_examples_workers',
'3',
'--shards',
'15',
'--gpu', # GPU should not have any effect.
'--docker_image_gpu',
'image_gpu',
'--job_name_prefix',
'prefix_',
])
gcp_deepvariant_runner.run(self._argv)
mock_apply_async.assert_has_calls(
[
mock.call(mock.ANY, [
_HasAllOf('prefix_make_examples', 'gcr.io/dockerimage',
'SHARD_START_INDEX=0', 'SHARD_END_INDEX=4',
'EXAMPLES=gs://bucket/staging/examples/*')
]),
mock.call(mock.ANY, [
_HasAllOf('prefix_make_examples', 'gcr.io/dockerimage',
'SHARD_START_INDEX=5', 'SHARD_END_INDEX=9')
]),
mock.call(mock.ANY, [
_HasAllOf('prefix_make_examples', 'gcr.io/dockerimage',
'SHARD_START_INDEX=10', 'SHARD_END_INDEX=14')
]),
],
any_order=True,
)
def testRunCallVariants(self, mock_pool):
mock_apply_async = mock_pool.return_value.apply_async
mock_apply_async.return_value = None
self._argv.extend([
'--make_examples_workers',
'3',
'--jobs_to_run',
'call_variants',
'--call_variants_workers',
'3',
'--call_variants_cores_per_worker',
'5',
'--call_variants_cores_per_shard',
'2',
'--shards',
'15',
])
gcp_deepvariant_runner.run(self._argv)
mock_apply_async.assert_has_calls(
[
mock.call(mock.ANY, [
_HasAllOf('call_variants', 'gcr.io/dockerimage',
'CALL_VARIANTS_SHARD_INDEX=0')
]),
mock.call(mock.ANY, [
_HasAllOf('call_variants', 'gcr.io/dockerimage',
'CALL_VARIANTS_SHARD_INDEX=1')
]),
mock.call(mock.ANY, [
_HasAllOf('call_variants', 'gcr.io/dockerimage',
'CALL_VARIANTS_SHARD_INDEX=2')
]),
],
any_order=True,
)
