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_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_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_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)
'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, },