def execute(self, context):
sagemaker = SageMakerHook(
sagemaker_conn_id=self.sagemaker_conn_id,
use_db_config=self.use_db_config,
region_name=self.region_name,
check_interval=self.check_interval,
max_ingestion_time=self.max_ingestion_time
)
if self.model_config:
self.log.info(
"Creating SageMaker Model %s for transform job"
% self.model_config['ModelName']
)
sagemaker.create_model(self.model_config)
self.log.info(
"Creating SageMaker transform Job %s."
% self.transform_job_config['TransformJobName']
)
response = sagemaker.create_transform_job(
self.transform_job_config,
wait_for_completion=self.wait_for_completion)
if not response['ResponseMetadata']['HTTPStatusCode'] \
== 200:
raise AirflowException(
'Sagemaker transform Job creation failed: %s' % response)
else:
return response
def test_create_transform_job(self, mock_client, mock_check_url):
mock_check_url.return_value = True
mock_session = mock.Mock()
attrs = {'create_transform_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_transform_job(create_transform_params,
wait_for_completion=False)
mock_session.create_transform_job.assert_called_once_with(
**create_transform_params)
self.assertEqual(response, test_arn_return)
def test_create_transform_job(self, mock_client, mock_check_url):
mock_check_url.return_value = True
mock_session = mock.Mock()
attrs = {'create_transform_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_transform_job(create_transform_params,
wait_for_completion=False)
mock_session.create_transform_job.assert_called_once_with(
**create_transform_params)
self.assertEqual(response, test_arn_return)
def test_create_transform_job_db_config(self, mock_client, mock_check_url):
mock_check_url.return_value = True
mock_session = mock.Mock()
attrs = {'create_transform_job.return_value': test_arn_return}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook_use_db_config = SageMakerHook(
sagemaker_conn_id='sagemaker_test_conn_id', use_db_config=True)
response = hook_use_db_config.create_transform_job(
create_transform_params, wait_for_completion=False)
updated_config = copy.deepcopy(create_transform_params)
updated_config.update(db_config)
mock_session.create_transform_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,
use_db_config=self.use_db_config,
region_name=self.region_name,
check_interval=self.check_interval,
max_ingestion_time=self.max_ingestion_time)
if self.model_config:
self.log.info("Creating SageMaker Model %s for transform job" %
self.model_config['ModelName'])
sagemaker.create_model(self.model_config)
self.log.info("Creating SageMaker transform Job %s." %
self.transform_job_config['TransformJobName'])
response = sagemaker.create_transform_job(
self.transform_job_config,
wait_for_completion=self.wait_for_completion)
if not response['ResponseMetadata']['HTTPStatusCode'] \
== 200:
raise AirflowException(
'Sagemaker transform Job creation failed: %s' % response)
else:
return response
