Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rsareddy0329 merged 1 commit into from
Oct 5, 2026
+127 −15
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 57 additions & 13 deletions sagemaker-core/src/sagemaker/core/processing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.

Expand All @@ -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,
)

Expand Down Expand Up @@ -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()

Expand Down
72 changes: 70 additions & 2 deletions sagemaker-core/tests/unit/test_processing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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')")
Expand All @@ -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",
Expand Down
Loading