def test_success_create_training_job(self):
with patch('airflow.gcp.operators.mlengine.MLEngineHook') \
as 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 = MLEngineTrainingOperator(
**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(
'test-project', self.TRAINING_INPUT, ANY)
def test_failed_job_error(self):
with patch('airflow.gcp.operators.mlengine.MLEngineHook') \
as 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 = MLEngineTrainingOperator(
**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(
'test-project', self.TRAINING_INPUT, ANY)
self.assertEqual('A failure message', str(context.exception))
def test_http_error(self):
http_error_code = 403
with patch('airflow.gcp.operators.mlengine.MLEngineHook') \
as mock_hook:
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 = MLEngineTrainingOperator(
**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(
'test-project', self.TRAINING_INPUT, ANY)
self.assertEqual(http_error_code, context.exception.resp.status)
"params": {
"model_name": MODEL_NAME
}
}
with models.DAG(
"example_gcp_mlengine",
default_args=default_args,
schedule_interval=None # Override to match your needs
) as dag:
training = MLEngineTrainingOperator(
task_id="training",
project_id=PROJECT_ID,
region="us-central1",
job_id="training-job-{{ ts_nodash }}-{{ params.model_name }}",
runtime_version="1.14",
python_version="3.5",
job_dir=JOB_DIR,
package_uris=[TRAINER_URI],
training_python_module=TRAINER_PY_MODULE,
training_args=[],
)
create_model = MLEngineModelOperator(
task_id="create-model",
project_id=PROJECT_ID,
operation='create',
model={
"name": MODEL_NAME,
},
)
