def test_describe_training_job_with_logs_complete(self, mock_client,
mock_log_client):
mock_session = mock.Mock()
mock_log_session = mock.Mock()
attrs = {
'describe_training_job.return_value':
DESCRIBE_TRAINING_COMPLETED_RETURN
}
log_attrs = {
'describe_log_streams.side_effect': LIFECYCLE_LOG_STREAMS,
'get_log_events.side_effect': STREAM_LOG_EVENTS,
}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
mock_log_session.configure_mock(**log_attrs)
mock_log_client.return_value = mock_log_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.describe_training_job_with_log(
job_name=job_name,
positions={},
stream_names=[],
instance_count=1,
state=LogState.COMPLETE,
last_description={},
last_describe_job_call=0,
)
self.assertEqual(response, (LogState.COMPLETE, {}, 0))
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
