def test_create_tuning_job(self, mock_client, mock_check_tuning):
mock_session = mock.Mock()
attrs = {'create_hyper_parameter_tuning_job.return_value':
test_arn_return}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.create_tuning_job(create_tuning_params,
wait_for_completion=False)
mock_session.create_hyper_parameter_tuning_job.\
assert_called_once_with(**create_tuning_params)
self.assertEqual(response, test_arn_return)
def test_create_tuning_job(self, mock_client, mock_check_tuning):
mock_session = mock.Mock()
attrs = {'create_hyper_parameter_tuning_job.return_value':
test_arn_return}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.create_tuning_job(create_tuning_params,
wait_for_completion=False)
mock_session.create_hyper_parameter_tuning_job.\
assert_called_once_with(**create_tuning_params)
self.assertEqual(response, test_arn_return)
def test_create_tuning_job_db_config(self, mock_client, mock_check_tuning):
mock_check_tuning.return_value = True
mock_session = mock.Mock()
attrs = {
'create_hyper_parameter_tuning_job.return_value': test_arn_return
}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(sagemaker_conn_id='sagemaker_test_conn_id',
use_db_config=True)
response = hook.create_tuning_job(create_tuning_params)
updated_config = copy.deepcopy(create_tuning_params)
updated_config.update(db_config)
mock_session.create_hyper_parameter_tuning_job. \
assert_called_once_with(**updated_config)
self.assertEqual(response, test_arn_return)
def test_create_tuning_job_db_config(self, mock_client, mock_check_tuning):
mock_check_tuning.return_value = True
mock_session = mock.Mock()
attrs = {'create_hyper_parameter_tuning_job.return_value':
test_arn_return}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(sagemaker_conn_id='sagemaker_test_conn_id',
use_db_config=True)
response = hook.create_tuning_job(create_tuning_params,
wait_for_completion=False)
updated_config = copy.deepcopy(create_tuning_params)
updated_config.update(db_config)
mock_session.create_hyper_parameter_tuning_job. \
assert_called_once_with(**updated_config)
self.assertEqual(response, test_arn_return)
def execute(self, context):
sagemaker = SageMakerHook(sagemaker_conn_id=self.sagemaker_conn_id,
region_name=self.region_name,
use_db_config=self.use_db_config,
check_interval=self.check_interval,
max_ingestion_time=self.max_ingestion_time)
self.log.info("Creating SageMaker Hyper Parameter Tunning Job %s" %
self.tuning_job_config['HyperParameterTuningJobName'])
response = sagemaker.create_tuning_job(
self.tuning_job_config,
wait_for_completion=self.wait_for_completion)
if not response['ResponseMetadata']['HTTPStatusCode'] \
== 200:
raise AirflowException("Sagemaker Tuning Job creation failed: %s" %
response)
else:
return response
def execute(self, context):
sagemaker = SageMakerHook(sagemaker_conn_id=self.sagemaker_conn_id,
region_name=self.region_name,
use_db_config=self.use_db_config,
check_interval=self.check_interval,
max_ingestion_time=self.max_ingestion_time
)
self.log.info(
"Creating SageMaker Hyper Parameter Tunning Job %s"
% self.tuning_job_config['HyperParameterTuningJobName']
)
response = sagemaker.create_tuning_job(
self.tuning_job_config,
wait_for_completion=self.wait_for_completion
)
if not response['ResponseMetadata']['HTTPStatusCode'] \
== 200:
raise AirflowException(
"Sagemaker Tuning Job creation failed: %s" % response)
else:
return response
