def test_console_extra_link(self, mock_hook):
training_op = MLEngineStartTrainingJobOperator(
**self.TRAINING_DEFAULT_ARGS)
ti = TaskInstance(
task=training_op,
execution_date=DEFAULT_DATE,
)
job_id = self.TRAINING_DEFAULT_ARGS['job_id']
project_id = self.TRAINING_DEFAULT_ARGS['project_id']
gcp_metadata = {
"job_id": job_id,
"project_id": project_id,
}
ti.xcom_push(key='gcp_metadata', value=gcp_metadata)
self.assertEqual(
f"https://console.cloud.google.com/ai-platform/jobs/{job_id}?project={project_id}",
training_op.get_extra_links(DEFAULT_DATE,
AIPlatformConsoleLink.name),
)
self.assertEqual(
'',
training_op.get_extra_links(datetime.datetime(2019, 1, 1),
AIPlatformConsoleLink.name),
)
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_console_extra_link_serialized_field(self):
with self.dag:
training_op = MLEngineStartTrainingJobOperator(**self.TRAINING_DEFAULT_ARGS)
serialized_dag = SerializedDAG.to_dict(self.dag)
dag = SerializedDAG.from_dict(serialized_dag)
simple_task = dag.task_dict[self.TRAINING_DEFAULT_ARGS['task_id']]
# Check Serialized version of operator link
self.assertEqual(
serialized_dag["dag"]["tasks"][0]["_operator_extra_links"],
[{"airflow.providers.google.cloud.operators.mlengine.AIPlatformConsoleLink": {}}],
)
# Check DeSerialized version of operator link
self.assertIsInstance(list(simple_task.operator_extra_links)[0], AIPlatformConsoleLink)
job_id = self.TRAINING_DEFAULT_ARGS['job_id']
project_id = self.TRAINING_DEFAULT_ARGS['project_id']
gcp_metadata = {
"job_id": job_id,
"project_id": project_id,
}
ti = TaskInstance(task=training_op, execution_date=DEFAULT_DATE,)
ti.xcom_push(key='gcp_metadata', value=gcp_metadata)
self.assertEqual(
f"https://console.cloud.google.com/ai-platform/jobs/{job_id}?project={project_id}",
simple_task.get_extra_links(DEFAULT_DATE, AIPlatformConsoleLink.name),
)
self.assertEqual(
'', simple_task.get_extra_links(datetime.datetime(2019, 1, 1), AIPlatformConsoleLink.name),
)
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)
'categoricalValues': [
'BasicLSTMCell',
'BasicRNNCell',
'GRUCell',
'LSTMCell',
'LayerNormBasicLSTMCell',
],
})
# [START howto_operator_gcp_mlengine_training]
training = MLEngineStartTrainingJobOperator(
task_id="training",
project_id=PROJECT_ID,
region="us-central1",
job_id="training-job-{{ ts_nodash }}-{{ params.model_name }}",
runtime_version="1.15",
python_version="3.7",
job_dir=JOB_DIR,
package_uris=[TRAINER_URI],
training_python_module=TRAINER_PY_MODULE,
training_args=[],
labels={"job_type": "training"},
hyperparameters=hyperparams,
)
# [END howto_operator_gcp_mlengine_training]
# [START howto_operator_gcp_mlengine_create_model]
create_model = MLEngineCreateModelOperator(
task_id="create-model",
project_id=PROJECT_ID,
model={
"name": MODEL_NAME,
},