def get_sagemaker_response(self):
sagemaker_hook = SageMakerHook(aws_conn_id=self.aws_conn_id)
if self.print_log:
if not self.log_resource_inited:
self.init_log_resource(sagemaker_hook)
self.state, self.last_description, self.last_describe_job_call = \
sagemaker_hook.describe_training_job_with_log(self.job_name,
self.positions, self.stream_names,
self.instance_count, self.state,
self.last_description,
self.last_describe_job_call)
else:
self.last_description = sagemaker_hook.describe_training_job(
self.job_name)
status = self.state_from_response(self.last_description)
if status not in self.non_terminal_states(
) and status not in self.failed_states():
billable_time = \
(self.last_description['TrainingEndTime'] - self.last_description['TrainingStartTime']) * \
self.last_description['ResourceConfig']['InstanceCount']
self.log.info('Billable seconds: %s',
int(billable_time.total_seconds()) + 1)
return self.last_description
def get_sagemaker_response(self):
sagemaker = SageMakerHook(
aws_conn_id=self.aws_conn_id,
region_name=self.region_name
)
self.log.info('Poking Sagemaker Training Job %s', self.job_name)
return sagemaker.describe_training_job(self.job_name)
def test_describe_training_job(self, mock_client):
mock_session = mock.Mock()
attrs = {'describe_training_job.return_value': 'InProgress'}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.describe_training_job(job_name)
mock_session.describe_training_job.\
assert_called_once_with(TrainingJobName=job_name)
self.assertEqual(response, 'InProgress')
def test_describe_training_job(self, mock_client):
mock_session = mock.Mock()
attrs = {'describe_training_job.return_value': 'InProgress'}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.describe_training_job(job_name)
mock_session.describe_training_job.\
assert_called_once_with(TrainingJobName=job_name)
self.assertEqual(response, 'InProgress')
def get_sagemaker_response(self):
sagemaker_hook = SageMakerHook(aws_conn_id=self.aws_conn_id)
if self.print_log:
if not self.log_resource_inited:
self.init_log_resource(sagemaker_hook)
self.state, self.last_description, self.last_describe_job_call = \
sagemaker_hook.describe_training_job_with_log(self.job_name,
self.positions, self.stream_names,
self.instance_count, self.state,
self.last_description,
self.last_describe_job_call)
else:
self.last_description = sagemaker_hook.describe_training_job(self.job_name)
status = self.state_from_response(self.last_description)
if status not in self.non_terminal_states() and status not in self.failed_states():
billable_time = \
(self.last_description['TrainingEndTime'] - self.last_description['TrainingStartTime']) * \
self.last_description['ResourceConfig']['InstanceCount']
self.log.info('Billable seconds: %s', int(billable_time.total_seconds()) + 1)
return self.last_description
def get_sagemaker_response(self):
sagemaker = SageMakerHook(aws_conn_id=self.aws_conn_id,
region_name=self.region_name)
self.log.info('Poking Sagemaker Training Job %s', self.job_name)
return sagemaker.describe_training_job(self.job_name)