From de26b5ce1314695ddfa24d4c88352a5b4a3d62b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Mon, 7 Nov 2022 11:50:08 -0800 Subject: [PATCH] Sagemaker hook: remove extra call at the end when waiting for completion also took the opportunity to speedup tests relying on this method by waiting zero --- airflow/providers/amazon/aws/hooks/sagemaker.py | 12 ++++-------- tests/providers/amazon/aws/hooks/test_sagemaker.py | 9 ++++----- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 6dc4e27cd2859..8bf442367316d 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -673,9 +673,8 @@ def check_status( non_terminal_states = self.non_terminal_states sec = 0 - running = True - while running: + while True: time.sleep(check_interval) sec += check_interval @@ -688,19 +687,16 @@ def check_status( except ClientError: raise AirflowException("AWS request failed, check logs for more info") - if status in non_terminal_states: - running = True - elif status in self.failed_states: + if status in self.failed_states: raise AirflowException(f"SageMaker job failed because {response['FailureReason']}") - else: - running = False + elif status not in non_terminal_states: + break if max_ingestion_time and sec > max_ingestion_time: # ensure that the job gets killed if the max ingestion time is exceeded raise AirflowException(f"SageMaker job took more than {max_ingestion_time} seconds") self.log.info("SageMaker Job completed") - response = describe_function(job_name) return response def check_training_status_with_log( diff --git a/tests/providers/amazon/aws/hooks/test_sagemaker.py b/tests/providers/amazon/aws/hooks/test_sagemaker.py index 72e58cf697db3..43e1facfa1d00 100644 --- a/tests/providers/amazon/aws/hooks/test_sagemaker.py +++ b/tests/providers/amazon/aws/hooks/test_sagemaker.py @@ -327,16 +327,15 @@ def test_training_ends_with_wait(self, mock_client, mock_check_training): DESCRIBE_TRAINING_INPROGRESS_RETURN, DESCRIBE_TRAINING_STOPPING_RETURN, DESCRIBE_TRAINING_COMPLETED_RETURN, - DESCRIBE_TRAINING_COMPLETED_RETURN, ], } mock_session.configure_mock(**attrs) mock_client.return_value = mock_session hook = SageMakerHook(aws_conn_id="sagemaker_test_conn_id_1") hook.create_training_job( - create_training_params, wait_for_completion=True, print_log=False, check_interval=1 + create_training_params, wait_for_completion=True, print_log=False, check_interval=0 ) - assert mock_session.describe_training_job.call_count == 4 + assert mock_session.describe_training_job.call_count == 3 @mock.patch.object(SageMakerHook, "check_training_config") @mock.patch.object(SageMakerHook, "get_conn") @@ -360,7 +359,7 @@ def test_training_throws_error_when_failed_with_wait(self, mock_client, mock_che create_training_params, wait_for_completion=True, print_log=False, - check_interval=1, + check_interval=0, ) assert mock_session.describe_training_job.call_count == 3 @@ -616,7 +615,7 @@ def test_training_with_logs(self, mock_describe, mock_client, mock_log_client, m mock_log_client.return_value = mock_log_session hook = SageMakerHook(aws_conn_id="sagemaker_test_conn_id_1") hook.create_training_job( - create_training_params, wait_for_completion=True, print_log=True, check_interval=1 + create_training_params, wait_for_completion=True, print_log=True, check_interval=0 ) assert mock_describe.call_count == 3 assert mock_session.describe_training_job.call_count == 1