diff --git a/sagemaker-core/src/sagemaker/core/transformer.py b/sagemaker-core/src/sagemaker/core/transformer.py index d6b18c30ed..787978c8d6 100644 --- a/sagemaker-core/src/sagemaker/core/transformer.py +++ b/sagemaker-core/src/sagemaker/core/transformer.py @@ -155,6 +155,7 @@ def __init__( self.base_transform_job_name = base_transform_job_name self._current_job_name = None self.latest_transform_job = None + self._latest_transform_job_is_local = False self._reset_output_path = False self.sagemaker_session = sagemaker_session or Session() @@ -412,10 +413,21 @@ def submit(request): transformed = transform_util(serialized_request, "CreateTransformJobRequest") # Remove tags from transformed dict as TransformJob resource doesn't accept it transformed.pop("tags", None) - self.latest_transform_job = TransformJob(**transformed) + # Submission went through the configured session, so bind that session and + # region to the resource: a session-less TransformJob builds its client from + # the process default session, so wait()/stop() poll the wrong account or + # region and fail with ResourceNotFound, NoRegionError or an authorization + # error even though the job was created correctly (#6503). + self.latest_transform_job = TransformJob(**transformed)._set_client_context( + session=self.sagemaker_session.boto_session, + region=self.sagemaker_session.boto_region_name, + ) + # A local transform job runs to completion inside create_transform_job and + # exists only in the local session, so it must never be polled over AWS. + self._latest_transform_job_is_local = local_mode if wait: - self.latest_transform_job.wait(logs=logs) + self.wait(logs=logs) def delete_model(self): """Delete the corresponding SageMaker model for this Transformer.""" @@ -462,13 +474,26 @@ def _retrieve_image_uri(self): ) def wait(self, logs=True): - """Placeholder docstring""" + """Wait for the latest transform job to reach a terminal state.""" self._ensure_last_transform_job() + if self._latest_transform_job_is_local: + logger.info( + "Local transform job %s already ran to completion; nothing to wait for.", + self.latest_transform_job.get_name(), + ) + return self.latest_transform_job.wait(logs=logs) def stop_transform_job(self, wait=True): """Stop latest running batch transform job.""" self._ensure_last_transform_job() + if self._latest_transform_job_is_local: + logger.warning( + "Transform job %s ran locally and has already finished; there is nothing " + "to stop.", + self.latest_transform_job.get_name(), + ) + return self.latest_transform_job.stop() if wait: self.latest_transform_job.wait() @@ -495,17 +520,20 @@ def attach(cls, transform_job_name, sagemaker_session=None): """ sagemaker_session = sagemaker_session or Session() + # TransformJob.get() builds its client from a boto3 session, so pass the + # session's boto_session rather than the sagemaker Session itself, and keep + # the resource that describe returned instead of describing it a second time. transform_job = TransformJob.get( - transform_job_name=transform_job_name, session=sagemaker_session + transform_job_name=transform_job_name, + session=sagemaker_session.boto_session, + region=sagemaker_session.boto_region_name, ) if not transform_job: raise ValueError(f"Transform job {transform_job_name} not found") job_details = transform_job.__dict__ init_params = cls._prepare_init_params_from_job_description(job_details) transformer = cls(sagemaker_session=sagemaker_session, **init_params) - transformer.latest_transform_job = TransformJob.get( - transform_job_name=init_params["base_transform_job_name"], session=sagemaker_session - ) + transformer.latest_transform_job = transform_job return transformer diff --git a/sagemaker-core/tests/unit/test_transformer_session_propagation.py b/sagemaker-core/tests/unit/test_transformer_session_propagation.py new file mode 100644 index 0000000000..94818dcae0 --- /dev/null +++ b/sagemaker-core/tests/unit/test_transformer_session_propagation.py @@ -0,0 +1,229 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Transformer keeps its configured session after submission (issue #6503). + +``transform()`` submitted through the configured ``sagemaker_session`` and then +built ``TransformJob(**transformed)`` with no session, so ``wait()``, ``stop()`` +and the log stream fell back to the process default boto session: the job was +created in the target account and the wait failed against the default one. +Local mode was worse -- the wait could poll AWS for a job that only ever +existed in the local session. +""" + +from unittest.mock import MagicMock, Mock, patch + +import boto3 +import pytest + +from sagemaker.core.resources import Base, TransformJob +from sagemaker.core.shapes import TransformOutput, TransformResources +from sagemaker.core.transformer import Transformer +from sagemaker.core.utils.utils import SageMakerClient + +REGION = "us-east-2" +DEFAULT_REGION = "us-east-1" +MODEL_NAME = "transformer-session-propagation-model" +JOB_NAME = "transformer-session-propagation-job" +INPUT_URI = "s3://target-bucket/input/" +OUTPUT_URI = "s3://target-bucket/output/" + +# The session the caller configures (target account) and the process default one. +TARGET_BOTO_SESSION = boto3.Session( + aws_access_key_id="AKIDTARGET", aws_secret_access_key="secret", region_name=REGION +) +DEFAULT_BOTO_SESSION = boto3.Session( + aws_access_key_id="AKIDDEFAULT", aws_secret_access_key="secret", region_name=DEFAULT_REGION +) + + +@pytest.fixture(autouse=True) +def _reset_client_cache(): + SageMakerClient.reset() + yield + SageMakerClient.reset() + + +def _configured_session(local_mode=False): + session = MagicMock() + session.boto_session = TARGET_BOTO_SESSION + session.boto_region_name = REGION + session.local_mode = local_mode + session.default_bucket.return_value = "target-bucket" + session.default_bucket_prefix = None + session.sagemaker_config = {} + # Run the submit callback inline; the sagemaker_client is a mock, so no AWS call. + session._intercept_create_request.side_effect = lambda request, submit, name: ( + submit(request) if submit else None + ) + return session + + +def _described_job(): + """A TransformJob resource shaped like a DescribeTransformJob response.""" + return TransformJob( + transform_job_name=JOB_NAME, + model_name=MODEL_NAME, + transform_resources=TransformResources(instance_type="ml.m5.large", instance_count=1), + transform_output=TransformOutput(s3_output_path=OUTPUT_URI), + ) + + +def _transformer(session): + return Transformer( + model_name=MODEL_NAME, + instance_count=1, + instance_type="ml.m5.large", + output_path=OUTPUT_URI, + sagemaker_session=session, + ) + + +@pytest.fixture +def recorded_clients(): + """Record the (session, region) every SageMaker client request is built with.""" + calls = [] + client = MagicMock() + client.describe_transform_job.return_value = { + "TransformJobName": JOB_NAME, + "TransformJobStatus": "Completed", + } + + def fake(session=None, region_name=None, service_name="sagemaker"): + calls.append((session, region_name)) + return client + + with patch.object(Base, "get_sagemaker_client", side_effect=fake): + yield calls, client + + +class TestRemoteTransformBindsConfiguredSession: + def test_transform_binds_session_and_region_to_the_resource(self): + transformer = _transformer(_configured_session()) + + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=False) + + job = transformer.latest_transform_job + assert job._session is TARGET_BOTO_SESSION + assert job._region == REGION + + def test_waiting_inside_transform_uses_the_configured_session(self, recorded_clients): + calls, _ = recorded_clients + transformer = _transformer(_configured_session()) + + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=True, logs=False) + + assert calls, "the wait made no client request" + assert all(call == (TARGET_BOTO_SESSION, REGION) for call in calls), calls + + def test_wait_after_transform_uses_the_configured_session(self, recorded_clients): + calls, _ = recorded_clients + transformer = _transformer(_configured_session()) + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=False) + calls.clear() + + transformer.wait(logs=False) + + assert all(call == (TARGET_BOTO_SESSION, REGION) for call in calls), calls + + def test_stop_transform_job_uses_the_configured_session(self, recorded_clients): + calls, _ = recorded_clients + transformer = _transformer(_configured_session()) + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=False) + calls.clear() + + transformer.stop_transform_job(wait=False) + + assert calls, "stop made no client request" + assert all(call == (TARGET_BOTO_SESSION, REGION) for call in calls), calls + + def test_bound_client_is_signed_with_the_configured_credentials(self): + """End to end through SageMakerClient, with a default client built first.""" + default = SageMakerClient(session=DEFAULT_BOTO_SESSION) + transformer = _transformer(_configured_session()) + + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=False) + client = transformer.latest_transform_job._get_client() + + assert client is not default.sagemaker_client + assert client._request_signer._credentials.access_key == "AKIDTARGET" + + +class TestLocalTransformNeverCallsAws: + def test_local_wait_does_not_request_an_aws_client(self, recorded_clients): + calls, _ = recorded_clients + transformer = _transformer(_configured_session(local_mode=True)) + + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=True, logs=False) + + assert calls == [], f"a local transform must not poll AWS, got {calls}" + + def test_local_stop_does_not_request_an_aws_client(self, recorded_clients): + calls, _ = recorded_clients + transformer = _transformer(_configured_session(local_mode=True)) + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=False) + calls.clear() + + transformer.stop_transform_job(wait=False) + + assert calls == [], f"a local transform has nothing to stop in AWS, got {calls}" + + def test_remote_transform_is_not_marked_local(self): + transformer = _transformer(_configured_session()) + + transformer.transform(data=INPUT_URI, job_name=JOB_NAME, wait=False) + + assert transformer._latest_transform_job_is_local is False + + +class TestAttachUsesTheConfiguredSession: + def test_attach_passes_the_boto_session_and_reuses_the_resource(self): + session = _configured_session() + described = _described_job() + + with patch.object(TransformJob, "get", return_value=described) as mock_get: + transformer = Transformer.attach(JOB_NAME, sagemaker_session=session) + + mock_get.assert_called_once_with( + transform_job_name=JOB_NAME, + session=TARGET_BOTO_SESSION, + region=REGION, + ) + assert transformer.latest_transform_job is described + + def test_attach_raises_when_the_job_is_missing(self): + with patch.object(TransformJob, "get", return_value=None): + with pytest.raises(ValueError, match="not found"): + Transformer.attach(JOB_NAME, sagemaker_session=_configured_session()) + + def test_attached_transformer_is_not_marked_local(self): + session = _configured_session() + described = _described_job() + + with patch.object(TransformJob, "get", return_value=described): + transformer = Transformer.attach(JOB_NAME, sagemaker_session=session) + + assert transformer._latest_transform_job_is_local is False + + +class TestEnsureLastTransformJob: + def test_wait_without_a_job_raises(self): + transformer = _transformer(_configured_session()) + + with pytest.raises(ValueError, match="No transform job available"): + transformer.wait() + + def test_stop_without_a_job_raises(self): + transformer = _transformer(Mock(local_mode=False, sagemaker_config={})) + + with pytest.raises(ValueError, match="No transform job available"): + transformer.stop_transform_job()