From 8249ac512e3361e0c06afe90a22cdb2d1db94341 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Thu, 15 Oct 2026 15:12:22 -0700 Subject: [PATCH] fix: pass the configured session to ModelBuilder's Model reloads ModelBuilder.build() creates the SageMaker Model through the configured sagemaker_session, but reloaded it with a session-less Model.get(). That helper builds its own boto3 session from the ambient credential chain, so DescribeModel is signed with the default account's credentials: in a cross-account setup CreateModel succeeds in the target account and the reload then fails with "Could not find model". Single-account setups hide the defect because both sessions resolve to the same account. Pass sagemaker_session.boto_session to Model.get() at all three Model reload/lookup sites in model_builder.py (_create_sagemaker_model, _find_reusable_model and _get_model_for_endpoint - the latter two read the model name from the configured session's client and then described it session-lessly, the same defect on build()'s reuse path), matching the pattern the Endpoint.get() call sites already use after #6069. Model.get() also records the session on the resource it returns, so refresh, delete and wait calls on the Model that build() hands back stay in the configured account as well. Fixes #6459 --- .../src/sagemaker/serve/model_builder.py | 21 +- .../test_model_builder_session_propagation.py | 198 ++++++++++++++++++ 2 files changed, 216 insertions(+), 3 deletions(-) create mode 100644 sagemaker-serve/tests/unit/test_model_builder_session_propagation.py diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 10a29fa5f8..b10f922571 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2452,7 +2452,11 @@ def _get_model_for_endpoint(self, endpoint_name: str) -> Optional[Model]: model_name = variants[0].get("ModelName") if not model_name: return None - return Model.get(model_name=model_name, region=self.region) + return Model.get( + model_name=model_name, + session=self.sagemaker_session.boto_session, + region=self.region, + ) except Exception as e: logger.warning( "Could not resolve the Model backing endpoint %s: %s.", @@ -2482,7 +2486,11 @@ def _find_reusable_model(self) -> Optional["Model"]: if model_arn: # Extract model name from ARN: arn:aws:sagemaker:region:account:model/name model_name = model_arn.rsplit("/", 1)[-1] - return Model.get(model_name=model_name, region=self.region) + return Model.get( + model_name=model_name, + session=self.sagemaker_session.boto_session, + region=self.region, + ) except Exception as e: logger.warning("Could not search Models for reuse: %s", e) @@ -3168,7 +3176,14 @@ def _create_sagemaker_model(self): self.sagemaker_session.create_model(**create_model_args) if isinstance(self.sagemaker_session, PipelineSession): return - return Model.get(model_name=self.model_name, region=self.region) + # CreateModel above went through the configured session, so the reload + # must use it too: a session-less get() describes the model in the + # default boto session's account and fails cross-account. + return Model.get( + model_name=self.model_name, + session=self.sagemaker_session.boto_session, + region=self.region, + ) def _create_model(self): """Create a SageMaker Model instance from the current configuration.""" diff --git a/sagemaker-serve/tests/unit/test_model_builder_session_propagation.py b/sagemaker-serve/tests/unit/test_model_builder_session_propagation.py new file mode 100644 index 0000000000..7fed4c219d --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_session_propagation.py @@ -0,0 +1,198 @@ +"""Unit tests for session propagation on ModelBuilder's model-reload paths (issue #6459). + +``build()`` creates the SageMaker Model with the configured ``sagemaker_session`` +and then reloads it with ``Model.get()``. A session-less ``Model.get()`` builds its +own boto3 session, so the DescribeModel call is signed with the default +credentials: cross-account builds create the model in the target account and then +fail to describe it. The same applies to the Model lookups used for resource reuse. +""" + +import unittest +from unittest.mock import MagicMock, Mock, patch + +import boto3 + +from sagemaker.core.resources import Base, Model +from sagemaker.core.workflow.pipeline_context import PipelineSession +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.serve.utils.types import ModelServer + +ROLE_ARN = "arn:aws:iam::111111111111:role/TargetAccountSageMakerRole" +REGION = "us-east-2" +IMAGE_URI = "111111111111.dkr.ecr.us-east-2.amazonaws.com/inference:latest" +MODEL_NAME = "modelbuilder-session-propagation" +MODEL_ARN = f"arn:aws:sagemaker:{REGION}:111111111111:model/{MODEL_NAME}" + +# Two real, distinguishable boto3 sessions. CONFIGURED_BOTO_SESSION is what the +# caller hands to ModelBuilder (the target account); DEFAULT_BOTO_SESSION stands +# in for the process default session (a different account). +CONFIGURED_BOTO_SESSION = boto3.Session( + region_name=REGION, + aws_access_key_id="configured-key", + aws_secret_access_key="configured-secret", +) +DEFAULT_BOTO_SESSION = boto3.Session( + region_name=REGION, + aws_access_key_id="default-key", + aws_secret_access_key="default-secret", +) + + +def _configured_session(): + """A sagemaker_session wired to CONFIGURED_BOTO_SESSION (the target account).""" + session = Mock() + session.boto_region_name = REGION + session.boto_session = CONFIGURED_BOTO_SESSION + session.default_bucket.return_value = "test-bucket" + session.default_bucket_prefix = None + session.config = {} + session.sagemaker_config = {} + session.settings = Mock() + session.settings.include_jumpstart_tags = False + session.settings._local_download_dir = None + session.sagemaker_client = MagicMock() + return session + + +def _make_builder(session): + builder = ModelBuilder( + model=Mock(), + model_server=ModelServer.TORCHSERVE, + role_arn=ROLE_ARN, + sagemaker_session=session, + instance_type="ml.m5.large", + ) + builder.model_name = MODEL_NAME + builder.region = REGION + return builder + + +@patch("sagemaker.serve.model_builder.resolve_nested_dict_value_from_config") +@patch("sagemaker.serve.model_builder.resolve_value_from_config") +@patch.object(ModelBuilder, "_init_sagemaker_session_if_does_not_exist") +@patch.object(ModelBuilder, "_ensure_base_name_if_needed") +@patch.object(ModelBuilder, "_prepare_container_def") +class TestCreateSageMakerModelReloadSession(unittest.TestCase): + """The reload after CreateModel must use the configured session.""" + + def _run(self, session, mocks): + mock_prepare, mock_resolve, mock_resolve_nested = mocks + mock_prepare.return_value = {"Image": IMAGE_URI, "Environment": {}} + mock_resolve.side_effect = lambda value, *args, **kwargs: value + mock_resolve_nested.side_effect = lambda value, *args, **kwargs: value + + builder = _make_builder(_configured_session()) + if session is not None: + builder.sagemaker_session = session + return builder, builder._create_sagemaker_model() + + def test_reload_passes_configured_session_to_model_get( + self, mock_prepare, mock_base_name, mock_init, mock_resolve, mock_resolve_nested + ): + reloaded = Mock(model_name=MODEL_NAME, model_arn=MODEL_ARN) + + with patch( + "sagemaker.serve.model_builder.Model.get", return_value=reloaded + ) as mock_model_get: + builder, result = self._run(None, (mock_prepare, mock_resolve, mock_resolve_nested)) + + builder.sagemaker_session.create_model.assert_called_once() + mock_model_get.assert_called_once_with( + model_name=MODEL_NAME, + session=CONFIGURED_BOTO_SESSION, + region=REGION, + ) + self.assertIs(result, reloaded) + + def test_reload_describe_client_is_built_from_configured_session( + self, mock_prepare, mock_base_name, mock_init, mock_resolve, mock_resolve_nested + ): + """DescribeModel must be issued by a client built from the configured session. + + Model.get() resolves its client through Base.get_sagemaker_client; with no + session that helper creates one from the ambient credential chain, which is + the default account. + """ + captured = {} + + def fake_get_client(session=None, region_name=None, service_name="sagemaker"): + captured["session"] = session + captured["region_name"] = region_name + client = MagicMock() + client.describe_model.return_value = {"ModelName": MODEL_NAME} + return client + + with patch.object(Base, "get_sagemaker_client", side_effect=fake_get_client): + _builder, result = self._run(None, (mock_prepare, mock_resolve, mock_resolve_nested)) + + self.assertIs(captured["session"], CONFIGURED_BOTO_SESSION) + self.assertIsNot(captured["session"], DEFAULT_BOTO_SESSION) + self.assertEqual(captured["region_name"], REGION) + self.assertIsInstance(result, Model) + self.assertEqual(result.model_name, MODEL_NAME) + + def test_pipeline_session_build_still_returns_without_reloading( + self, mock_prepare, mock_base_name, mock_init, mock_resolve, mock_resolve_nested + ): + """A PipelineSession build returns before the reload and must stay that way.""" + with patch("sagemaker.serve.model_builder.Model.get") as mock_model_get: + _builder, result = self._run( + Mock(spec=PipelineSession), (mock_prepare, mock_resolve, mock_resolve_nested) + ) + + self.assertIsNone(result) + mock_model_get.assert_not_called() + + +class TestModelLookupSession(unittest.TestCase): + """The Model lookups used for resource reuse must use the configured session.""" + + def setUp(self): + self.session = _configured_session() + self.builder = _make_builder(self.session) + self.found_model = Mock(model_name=MODEL_NAME, model_arn=MODEL_ARN) + + def test_find_reusable_model_passes_configured_session(self): + with ( + patch.object( + ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model" + ), + patch( + "sagemaker.serve.model_reuse.find_sagemaker_model_arn_by_tag", + return_value=MODEL_ARN, + ), + patch( + "sagemaker.serve.model_builder.Model.get", return_value=self.found_model + ) as mock_model_get, + ): + result = self.builder._find_reusable_model() + + mock_model_get.assert_called_once_with( + model_name=MODEL_NAME, + session=CONFIGURED_BOTO_SESSION, + region=REGION, + ) + self.assertIs(result, self.found_model) + + def test_get_model_for_endpoint_passes_configured_session(self): + client = self.session.sagemaker_client + client.describe_endpoint.return_value = {"EndpointConfigName": "test-config"} + client.describe_endpoint_config.return_value = { + "ProductionVariants": [{"ModelName": MODEL_NAME}] + } + + with patch( + "sagemaker.serve.model_builder.Model.get", return_value=self.found_model + ) as mock_model_get: + result = self.builder._get_model_for_endpoint("test-endpoint") + + mock_model_get.assert_called_once_with( + model_name=MODEL_NAME, + session=CONFIGURED_BOTO_SESSION, + region=REGION, + ) + self.assertIs(result, self.found_model) + + +if __name__ == "__main__": + unittest.main()