diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 0c6abffba80e0..ae57e832d1b90 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 pattern.fullmatch(found_name) is not None + 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 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. """ 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")