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