From 91c51931f15dd2304a752c1bcc07972208fc1f7d Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 1 Nov 2022 15:23:18 -0700 Subject: [PATCH 1/3] Correct job name matching in SagemakerProcessingOperator SagemakerProcessingOperator "increment if job name exists" was matching any job name that contains the proposed name when really it should have been matching names which fit the specific pattern ^{proposed_name}{suffix_template}$ Mocking sleep() shaved test run times from 9.98s to 457ms --- .../providers/amazon/aws/hooks/sagemaker.py | 31 ++++++++++++++----- .../amazon/aws/operators/sagemaker.py | 5 ++- .../amazon/aws/hooks/test_sagemaker.py | 29 +++++++++++++---- 3 files changed, 51 insertions(+), 14 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 0c6abffba80e0..21680d25d1d64 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -19,6 +19,7 @@ import collections import os +import re import tarfile import tempfile import time @@ -954,33 +955,49 @@ def find_processing_job_by_name(self, processing_job_name: str) -> bool: ) return bool(self.count_processing_jobs_by_name(processing_job_name)) + @staticmethod + def _name_matches_pattern( + processing_job_name: str, + found_name: str, + job_name_suffix: str | None = None, + ) -> bool: + pattern = re.compile(f"^{processing_job_name}({job_name_suffix})?$") + return bool(pattern.search(found_name)) + def count_processing_jobs_by_name( self, processing_job_name: str, + job_name_suffix: str | None = None, throttle_retry_delay: int = 2, retries: int = 3, ) -> int: """ Returns the number of processing jobs found with the provided name prefix. :param processing_job_name: The prefix to look for. + :param job_name_suffix: The optional suffix which may be appended to dedupe an existing job name. :param throttle_retry_delay: Seconds to wait if a ThrottlingException is hit. :param retries: The max number of times to retry. :returns: The number of processing jobs that start with the provided prefix. """ try: jobs = self.get_conn().list_processing_jobs(NameContains=processing_job_name) - return len(jobs["ProcessingJobSummaries"]) + # We want to make sure the job name starts with the provided name, not just contains it. + matching_jobs = [ + job["ProcessingJobName"] + for job in jobs["ProcessingJobSummaries"] + if self._name_matches_pattern(processing_job_name, job["ProcessingJobName"], job_name_suffix) + ] + return len(matching_jobs) except ClientError as e: if e.response["Error"]["Code"] == "ResourceNotFound": # No jobs found with that name. This is good, return 0. return 0 - if e.response["Error"]["Code"] == "ThrottlingException": + if e.response["Error"]["Code"] == "ThrottlingException" and retries: # If we hit a ThrottlingException, back off a little and try again. - if retries: - time.sleep(throttle_retry_delay) - return self.count_processing_jobs_by_name( - processing_job_name, throttle_retry_delay * 2, retries - 1 - ) + time.sleep(throttle_retry_delay) + return self.count_processing_jobs_by_name( + processing_job_name, job_name_suffix, throttle_retry_delay * 2, retries - 1 + ) raise def delete_model(self, model_name: str): diff --git a/airflow/providers/amazon/aws/operators/sagemaker.py b/airflow/providers/amazon/aws/operators/sagemaker.py index 80adc3665b9b4..c5f08db049206 100644 --- a/airflow/providers/amazon/aws/operators/sagemaker.py +++ b/airflow/providers/amazon/aws/operators/sagemaker.py @@ -181,7 +181,10 @@ def expand_role(self) -> None: def execute(self, context: Context) -> dict: self.preprocess_config() processing_job_name = self.config["ProcessingJobName"] - existing_jobs_found = self.hook.count_processing_jobs_by_name(processing_job_name) + processing_job_dedupe_pattern = "-[0-9]+$" + existing_jobs_found = self.hook.count_processing_jobs_by_name( + processing_job_name, processing_job_dedupe_pattern + ) if existing_jobs_found: if self.action_if_job_exists == "fail": raise AirflowException( diff --git a/tests/providers/amazon/aws/hooks/test_sagemaker.py b/tests/providers/amazon/aws/hooks/test_sagemaker.py index adcfd5397f1c9..52906031d7a4f 100644 --- a/tests/providers/amazon/aws/hooks/test_sagemaker.py +++ b/tests/providers/amazon/aws/hooks/test_sagemaker.py @@ -318,7 +318,8 @@ def test_create_training_job(self, mock_client, mock_check_training): @mock.patch.object(SageMakerHook, "check_training_config") @mock.patch.object(SageMakerHook, "get_conn") - def test_training_ends_with_wait(self, mock_client, mock_check_training): + @mock.patch("time.sleep", return_value=None) + def test_training_ends_with_wait(self, _, mock_client, mock_check_training): mock_check_training.return_value = True mock_session = mock.Mock() attrs = { @@ -339,7 +340,8 @@ def test_training_ends_with_wait(self, mock_client, mock_check_training): @mock.patch.object(SageMakerHook, "check_training_config") @mock.patch.object(SageMakerHook, "get_conn") - def test_training_throws_error_when_failed_with_wait(self, mock_client, mock_check_training): + @mock.patch("time.sleep", return_value=None) + def test_training_throws_error_when_failed_with_wait(self, _, mock_client, mock_check_training): mock_check_training.return_value = True mock_session = mock.Mock() attrs = { @@ -592,7 +594,8 @@ def test_describe_training_job_with_complete_states(self, mock_client, mock_log_ @mock.patch.object(AwsLogsHook, "get_conn") @mock.patch.object(SageMakerHook, "get_conn") @mock.patch.object(SageMakerHook, "describe_training_job_with_log") - def test_training_with_logs(self, mock_describe, mock_client, mock_log_client, mock_check_training): + @mock.patch("time.sleep", return_value=None) + def test_training_with_logs(self, _, mock_describe, mock_client, mock_log_client, mock_check_training): mock_check_training.return_value = True mock_describe.side_effect = [ (LogState.WAIT_IN_PROGRESS, DESCRIBE_TRAINING_INPROGRESS_RETURN, 0), @@ -653,6 +656,20 @@ def test_count_processing_jobs_by_name(self, mock_conn): ret = hook.count_processing_jobs_by_name(existing_job_name) assert ret == 1 + @mock.patch.object(SageMakerHook, "get_conn") + def test_count_processing_jobs_by_name_only_counts_actual_hits(self, mock_conn): + hook = SageMakerHook(aws_conn_id="sagemaker_test_conn_id") + existing_job_name = "existing_job" + mock_conn().list_processing_jobs.return_value = { + "ProcessingJobSummaries": [ + {"ProcessingJobName": existing_job_name}, + {"ProcessingJobName": f"contains_but_does_not_start_with_{existing_job_name}"}, + {"ProcessingJobName": f"{existing_job_name}_with_different_suffix-123"}, + ] + } + ret = hook.count_processing_jobs_by_name(existing_job_name) + assert ret == 1 + @mock.patch.object(SageMakerHook, "get_conn") @mock.patch("time.sleep", return_value=None) def test_count_processing_jobs_by_name_retries_on_throttle_exception(self, _, mock_conn): @@ -676,12 +693,12 @@ def test_count_processing_jobs_by_name_fails_after_max_retries(self, _, mock_con error_response={"Error": {"Code": "ThrottlingException"}}, operation_name="empty" ) hook = SageMakerHook(aws_conn_id="sagemaker_test_conn_id") + retries = 3 with pytest.raises(ClientError) as raised_exception: - hook.count_processing_jobs_by_name("existing_job") + hook.count_processing_jobs_by_name("existing_job", retries=retries) - # One initial call plus retries - assert mock_conn().list_processing_jobs.call_count == 4 + assert mock_conn().list_processing_jobs.call_count == retries + 1 assert raised_exception.value.response["Error"]["Code"] == "ThrottlingException" @mock.patch.object(SageMakerHook, "get_conn") From f0cf68451c2c1c18edb6cfd9dd41bd503f48a01b Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Thu, 17 Nov 2022 12:26:33 -0800 Subject: [PATCH 2/3] fixes --- airflow/providers/amazon/aws/hooks/sagemaker.py | 2 +- docs/spelling_wordlist.txt | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 21680d25d1d64..f17be7da87580 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -962,7 +962,7 @@ def _name_matches_pattern( job_name_suffix: str | None = None, ) -> bool: pattern = re.compile(f"^{processing_job_name}({job_name_suffix})?$") - return bool(pattern.search(found_name)) + return pattern.fullmatch(found_name) is not None def count_processing_jobs_by_name( self, diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index eaa0553ed6d1d..52ab3a689d7ba 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -389,6 +389,7 @@ Decrypt decrypt decrypted Decrypts +dedupe deduplicate deduplication deepcopy From 15c44a4dcdd8d465925dfadea574f8a66d4df09b Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 22 Nov 2022 12:38:19 -0800 Subject: [PATCH 3/3] Unfix - replace 'dedupe' with 'deduplicate' --- airflow/providers/amazon/aws/hooks/sagemaker.py | 2 +- docs/spelling_wordlist.txt | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index f17be7da87580..ae57e832d1b90 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -974,7 +974,7 @@ def count_processing_jobs_by_name( """ Returns the number of processing jobs found with the provided name prefix. :param processing_job_name: The prefix to look for. - :param job_name_suffix: The optional suffix which may be appended to dedupe an existing job name. + :param job_name_suffix: The optional suffix which may be appended to deduplicate an existing job name. :param throttle_retry_delay: Seconds to wait if a ThrottlingException is hit. :param retries: The max number of times to retry. :returns: The number of processing jobs that start with the provided prefix. diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 52ab3a689d7ba..eaa0553ed6d1d 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -389,7 +389,6 @@ Decrypt decrypt decrypted Decrypts -dedupe deduplicate deduplication deepcopy