fix(processing): respect session region when waiting for processing jobs (#6369)

Processor.run / FrameworkProcessor.run waited via self.latest_job.wait(),
whose generated refresh() resolves a default, global SageMaker client
(Base.get_sagemaker_client()) rather than the region of the sagemaker_session
that created the job. With a non-default-region session the wait/describe
and log tailing hit the wrong region and fail (issue #5796).

Route the wait path through the session instead:
- logs=True -> logs_for_processing_job(sagemaker_session, ...), which already
  resolves clients from sagemaker_session; also fix it to describe via
  sagemaker_session.sagemaker_client (region-aware) instead of
  ProcessingJob.get(...).refresh() (the chained refresh used the global client
  and read PascalCase keys off a snake_case resource __dict__).
- logs=False -> new _wait_for_processing_job(sagemaker_session, job_name),
  which polls describe_processing_job through the session's region-aware client
  and raises on failure.

Add unit tests covering both wait paths and _wait_for_processing_job.

Relates-to #5796

Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>

Author: rsareddy0329

sagemaker-core/src/sagemaker/core/processing.py

@@ -361,7 +361,17 @@ def run(
         if not isinstance(self.sagemaker_session, PipelineSession):
             self.jobs.append(self.latest_job)
             if wait:
-                self.latest_job.wait(logs=logs)
+                if logs:
+                    logs_for_processing_job(
+                        sagemaker_session=self.sagemaker_session,
+                        job_name=self.latest_job.processing_job_name,
+                        wait=True,
+                    )
+                else:
+                    _wait_for_processing_job(
+                        sagemaker_session=self.sagemaker_session,
+                        job_name=self.latest_job.processing_job_name,
+                    )
 
     def _extend_processing_args(self, inputs, outputs, **kwargs):  # pylint: disable=W0613
         """Extend inputs and outputs based on extra parameters"""
@@ -929,7 +939,17 @@ def run(
         if not isinstance(self.sagemaker_session, PipelineSession):
             self.jobs.append(self.latest_job)
             if wait:
-                self.latest_job.wait(logs=logs)
+                if logs:
+                    logs_for_processing_job(
+                        sagemaker_session=self.sagemaker_session,
+                        job_name=self.latest_job.processing_job_name,
+                        wait=True,
+                    )
+                else:
+                    _wait_for_processing_job(
+                        sagemaker_session=self.sagemaker_session,
+                        job_name=self.latest_job.processing_job_name,
+                    )
 
     def _include_code_in_inputs(self, inputs, code, kms_key=None):
         """Converts code to appropriate input and includes in input list.
@@ -1832,6 +1852,36 @@ def _get_process_request(
     return process_request
 
 
+def _wait_for_processing_job(sagemaker_session, job_name, poll=10):
+    """Wait for a processing job to reach a terminal state, respecting the session region.
+
+    Unlike ``ProcessingJob.wait()`` (which resolves a default, global SageMaker client),
+    this describes the job through ``sagemaker_session.sagemaker_client`` so the job is
+    polled in the same region it was created in (issue #5796).
+
+    Args:
+        sagemaker_session: The session used to create the job; its region-aware
+            ``sagemaker_client`` is used to describe the job.
+        job_name (str): Name of the processing job to wait for.
+        poll (int): The interval in seconds between polling for job completion.
+
+    Raises:
+        ValueError: If the processing job fails.
+    """
+    terminal_states = ("Completed", "Failed", "Stopped")
+
+    def _describe_if_terminal():
+        description = sagemaker_session.sagemaker_client.describe_processing_job(
+            ProcessingJobName=job_name
+        )
+        if description["ProcessingJobStatus"] in terminal_states:
+            return description
+        return None
+
+    description = _wait_until(_describe_if_terminal, poll)
+    _check_job_status(job_name, description, "ProcessingJobStatus")
+
+
 def logs_for_processing_job(sagemaker_session, job_name, wait=False, poll=10):
     """Display logs for a given processing job, optionally tailing them until the is complete.
 
@@ -1847,11 +1897,9 @@ def logs_for_processing_job(sagemaker_session, job_name, wait=False, poll=10):
     """
 
     description = _wait_until(
-        lambda: ProcessingJob.get(
-            processing_job_name=job_name, session=sagemaker_session.boto_session
-        )
-        .refresh()
-        .__dict__,
+        lambda: sagemaker_session.sagemaker_client.describe_processing_job(
+            ProcessingJobName=job_name
+        ),
         poll,
     )
 
@@ -1901,12 +1949,8 @@ def logs_for_processing_job(sagemaker_session, job_name, wait=False, poll=10):
         if state == LogState.JOB_COMPLETE:
             state = LogState.COMPLETE
         elif time.time() - last_describe_job_call >= 30:
-            description = (
-                ProcessingJob.get(
-                    processing_job_name=job_name, session=sagemaker_session.boto_session
-                )
-                .refresh()
-                .__dict__
+            description = sagemaker_session.sagemaker_client.describe_processing_job(
+                ProcessingJobName=job_name
             )
             last_describe_job_call = time.time()
 

sagemaker-core/tests/unit/test_processing.py

@@ -23,6 +23,7 @@
     _processing_output_to_request_dict,
     _get_process_request,
     logs_for_processing_job,
+    _wait_for_processing_job,
 )
 from sagemaker.core.shapes import (
     ProcessingInput,
@@ -1099,6 +1100,7 @@ def test_run_with_wait(self, mock_session):
 
         mock_job = Mock()
         mock_job.wait = Mock()
+        mock_job.processing_job_name = "test-processing-job"
 
         with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as f:
             f.write("print('test')")
@@ -1110,12 +1112,78 @@ def test_run_with_wait(self, mock_session):
                     with patch(
                         "sagemaker.core.s3.S3Uploader.upload", return_value="s3://bucket/code.py"
                     ):
-                        processor.run(code=temp_file, wait=True, logs=False)
-                        mock_job.wait.assert_called_once()
+                        with patch(
+                            "sagemaker.core.processing._wait_for_processing_job"
+                        ) as mock_wait:
+                            processor.run(code=temp_file, wait=True, logs=False)
+                            # The wait path must go through the session-aware helper
+                            # (which uses the session's region), not latest_job.wait()
+                            # which resolves a default-region client (issue #5796).
+                            mock_job.wait.assert_not_called()
+                            mock_wait.assert_called_once_with(
+                                sagemaker_session=mock_session,
+                                job_name="test-processing-job",
+                            )
         finally:
             if os.path.exists(temp_file):
                 os.unlink(temp_file)
 
+    def test_run_with_wait_and_logs_uses_session(self, mock_session):
+        processor = ScriptProcessor(
+            role="arn:aws:iam::123456789012:role/SageMakerRole",
+            image_uri="test-image:latest",
+            command=["python3"],
+            instance_count=1,
+            instance_type="ml.m5.xlarge",
+            sagemaker_session=mock_session,
+        )
+
+        mock_job = Mock()
+        mock_job.processing_job_name = "test-processing-job"
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as f:
+            f.write("print('test')")
+            temp_file = f.name
+
+        try:
+            with patch.object(processor, "_start_new", return_value=mock_job):
+                with patch("os.path.isfile", return_value=True):
+                    with patch(
+                        "sagemaker.core.s3.S3Uploader.upload", return_value="s3://bucket/code.py"
+                    ):
+                        with patch(
+                            "sagemaker.core.processing.logs_for_processing_job"
+                        ) as mock_logs:
+                            processor.run(code=temp_file, wait=True, logs=True)
+                            # logs path must also be session-aware (issue #5796).
+                            mock_logs.assert_called_once_with(
+                                sagemaker_session=mock_session,
+                                job_name="test-processing-job",
+                                wait=True,
+                            )
+        finally:
+            if os.path.exists(temp_file):
+                os.unlink(temp_file)
+
+    def test_wait_for_processing_job_uses_session_region(self, mock_session):
+        """_wait_for_processing_job describes via the session's region-aware client."""
+        mock_session.sagemaker_client.describe_processing_job = Mock(
+            return_value={"ProcessingJobStatus": "Completed"}
+        )
+        with patch("sagemaker.core.processing._wait_until", side_effect=lambda fn, poll: fn()):
+            _wait_for_processing_job(mock_session, "test-job", poll=1)
+        mock_session.sagemaker_client.describe_processing_job.assert_called_with(
+            ProcessingJobName="test-job"
+        )
+
+    def test_wait_for_processing_job_raises_on_failure(self, mock_session):
+        mock_session.sagemaker_client.describe_processing_job = Mock(
+            return_value={"ProcessingJobStatus": "Failed", "FailureReason": "boom"}
+        )
+        with patch("sagemaker.core.processing._wait_until", side_effect=lambda fn, poll: fn()):
+            with pytest.raises(Exception):
+                _wait_for_processing_job(mock_session, "test-job", poll=1)
+
     def test_run_without_wait(self, mock_session):
         processor = ScriptProcessor(
             role="arn:aws:iam::123456789012:role/SageMakerRole",