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 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)
assert response == 'InProgress'
def init_log_resource(self, hook: SageMakerHook) -> None:
"""Set tailing LogState for associated training job."""
description = hook.describe_training_job(self.job_name)
self.instance_count = description['ResourceConfig']['InstanceCount']
status = description['TrainingJobStatus']
job_already_completed = status not in self.non_terminal_states()
self.state = LogState.COMPLETE if job_already_completed else LogState.TAILING
self.last_description = description
self.last_describe_job_call = time.monotonic()
self.log_resource_inited = True
