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)