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_model(self, mock_client):
mock_session = mock.Mock()
attrs = {'create_model.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_model(create_model_params)
mock_session.create_model.assert_called_once_with(
**create_model_params)
self.assertEqual(response, test_arn_return)
def test_create_model(self, mock_client):
mock_session = mock.Mock()
attrs = {'create_model.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_model(create_model_params)
mock_session.create_model.assert_called_once_with(**create_model_params)
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
