diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 9dec7060bc..d0e2f65ae0 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/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. @@ -1826,6 +1846,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. @@ -1841,11 +1891,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, ) @@ -1895,12 +1943,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() diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index 559ac4179c..8f677eca3b 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/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",