Correct job name matching in SagemakerProcessingOperator

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):

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(

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")