def test_success_create_training_job_with_optional_args(self, mock_hook):
training_input = copy.deepcopy(self.TRAINING_INPUT)
training_input['trainingInput']['runtimeVersion'] = '1.6'
training_input['trainingInput']['pythonVersion'] = '3.5'
training_input['trainingInput'][
'jobDir'] = 'gs://some-bucket/jobs/test_training'
training_input['trainingInput'][
'serviceAccount'] = '*****@*****.**'
success_response = self.TRAINING_INPUT.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineStartTrainingJobOperator(
runtime_version='1.6',
python_version='3.5',
job_dir='gs://some-bucket/jobs/test_training',
service_account='*****@*****.**',
**self.TRAINING_DEFAULT_ARGS,
)
training_op.execute(MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default',
delegate_to=None,
impersonation_chain=None,
)
# Make sure only 'create_job' is invoked on hook instance
assert len(hook_instance.mock_calls) == 1
hook_instance.create_job.assert_called_once_with(
project_id='test-project',
job=training_input,
use_existing_job_fn=ANY)
def test_success_create_training_job_with_master_config(self, mock_hook):
custom_training_default_args: dict = copy.deepcopy(
self.TRAINING_DEFAULT_ARGS)
custom_training_default_args['scale_tier'] = 'CUSTOM'
training_input = copy.deepcopy(self.TRAINING_INPUT)
training_input['trainingInput']['runtimeVersion'] = '1.6'
training_input['trainingInput']['pythonVersion'] = '3.5'
training_input['trainingInput'][
'jobDir'] = 'gs://some-bucket/jobs/test_training'
training_input['trainingInput']['scaleTier'] = 'CUSTOM'
training_input['trainingInput']['masterType'] = 'n1-standard-4'
training_input['trainingInput']['masterConfig'] = {
'acceleratorConfig': {
'count': '1',
'type': 'NVIDIA_TESLA_P4'
},
}
success_response = training_input.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineStartTrainingJobOperator(
runtime_version='1.6',
python_version='3.5',
job_dir='gs://some-bucket/jobs/test_training',
master_type='n1-standard-4',
master_config={
'acceleratorConfig': {
'count': '1',
'type': 'NVIDIA_TESLA_P4'
},
},
**custom_training_default_args,
)
training_op.execute(MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default',
delegate_to=None,
impersonation_chain=None,
)
# Make sure only 'create_job' is invoked on hook instance
assert len(hook_instance.mock_calls) == 1
hook_instance.create_job.assert_called_once_with(
project_id='test-project',
job=training_input,
use_existing_job_fn=ANY)
def test_success_create_training_job(self, mock_hook):
success_response = self.TRAINING_INPUT.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineStartTrainingJobOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default', delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_once_with(
project_id='test-project', job=self.TRAINING_INPUT, use_existing_job_fn=ANY)
def test_success_create_training_job_with_master_image(self, hook):
arguments = {
'project_id': 'test-project',
'job_id': 'test_training',
'region': 'europe-west1',
'scale_tier': 'CUSTOM',
'master_type': 'n1-standard-8',
'master_config': {
'imageUri': 'eu.gcr.io/test-project/test-image:test-version',
},
'task_id': 'test-training',
'start_date': days_ago(1),
}
request = {
'jobId': 'test_training',
'trainingInput': {
'region': 'europe-west1',
'scaleTier': 'CUSTOM',
'masterType': 'n1-standard-8',
'masterConfig': {
'imageUri':
'eu.gcr.io/test-project/test-image:test-version',
},
},
}
response = request.copy()
response['state'] = 'SUCCEEDED'
hook_instance = hook.return_value
hook_instance.create_job.return_value = response
operator = MLEngineStartTrainingJobOperator(**arguments)
operator.execute(MagicMock())
hook.assert_called_once_with(
gcp_conn_id='google_cloud_default',
delegate_to=None,
impersonation_chain=None,
)
assert len(hook_instance.mock_calls) == 1
hook_instance.create_job.assert_called_once_with(
project_id='test-project',
job=request,
use_existing_job_fn=ANY,
)
def test_failed_job_error(self, mock_hook):
failure_response = self.TRAINING_INPUT.copy()
failure_response['state'] = 'FAILED'
failure_response['errorMessage'] = 'A failure message'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = failure_response
with self.assertRaises(RuntimeError) as context:
training_op = MLEngineStartTrainingJobOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default', delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_once_with(
project_id='test-project', job=self.TRAINING_INPUT, use_existing_job_fn=ANY)
self.assertEqual('A failure message', str(context.exception))
def test_http_error(self, mock_hook):
http_error_code = 403
hook_instance = mock_hook.return_value
hook_instance.create_job.side_effect = HttpError(
resp=httplib2.Response({'status': http_error_code}), content=b'Forbidden'
)
with self.assertRaises(HttpError) as context:
training_op = MLEngineStartTrainingJobOperator(**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default', delegate_to=None, impersonation_chain=None,
)
# Make sure only 'create_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_once_with(
project_id='test-project', job=self.TRAINING_INPUT, use_existing_job_fn=ANY
)
self.assertEqual(http_error_code, context.exception.resp.status)