Skip to content

Commit 1f09f26

Browse files
authored
fix(serve): Support pipeline variable image_uri in ModelBuilder.build (#6347)
ModelBuilder.image_uri is typed Optional[StrPipeVar], but build() treated it as a plain string in two places, so a ParameterString image failed under a PipelineSession: - is_1p_image_uri() sliced the value (image_uri[0:12]), raising "'ParameterString' object is not subscriptable" in _build_validations. - _create_sagemaker_model() ran "nova-" in image_uri for the Nova network-isolation default, raising "argument of type 'ParameterString' is not iterable". A pipeline variable is only resolved at execution time, so is_1p_image_uri() now treats non-string values as not first-party, and the Nova network-isolation check only inspects string images. The pipeline variable is passed through unchanged to CreateModel. Fixes #5760
1 parent c0bc32f commit 1f09f26

4 files changed

Lines changed: 159 additions & 3 deletions

File tree

‎sagemaker-serve/src/sagemaker/serve/model_builder.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3141,13 +3141,16 @@ def _create_sagemaker_model(self):
31413141
sagemaker_session=self.sagemaker_session,
31423142
)
31433143

3144-
# Nova 1P images require network isolation on the Model resource
3144+
# Nova 1P images require network isolation on the Model resource.
3145+
# A pipeline-variable image is only resolved at execution time, so it
3146+
# cannot be inspected here; set enable_network_isolation explicitly then.
31453147
enable_network_isolation = self._enable_network_isolation
31463148
resolved_image_uri = (
31473149
container_def["Image"] if isinstance(container_def, dict) else container_def[0]["Image"]
31483150
)
31493151
if (
31503152
not enable_network_isolation
3153+
and isinstance(resolved_image_uri, str)
31513154
and "nova-" in resolved_image_uri
31523155
and is_1p_image_uri(resolved_image_uri)
31533156
):

‎sagemaker-serve/src/sagemaker/serve/validations/check_image_uri.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44

55
import logging
66
import re
7+
from typing import Optional
8+
9+
from sagemaker.core.helper.pipeline_variable import StrPipeVar
710

811
logger = logging.getLogger(__name__)
912

@@ -318,8 +321,15 @@
318321
}
319322

320323

321-
def is_1p_image_uri(image_uri: str) -> bool:
322-
"""Shows if the given image_uri is owned by a 1st party account"""
324+
def is_1p_image_uri(image_uri: Optional[StrPipeVar]) -> bool:
325+
"""Shows if the given image_uri is owned by a 1st party account.
326+
327+
A pipeline variable (e.g. ``ParameterString``) is only resolved at pipeline
328+
execution time, so its account cannot be inspected at build time. Such values,
329+
and any other non-string value, are treated as not 1st party.
330+
"""
331+
if not isinstance(image_uri, str):
332+
return False
323333
image_uri_account = image_uri[0:12]
324334
return image_uri_account in all_accounts
325335

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
"""Unit tests for ModelBuilder with a pipeline-variable image_uri (issue #5760)."""
2+
3+
import unittest
4+
from unittest.mock import Mock, patch
5+
6+
import boto3
7+
8+
from sagemaker.core.workflow.parameters import ParameterString
9+
from sagemaker.core.workflow.pipeline_context import PipelineSession, _ModelStepArguments
10+
from sagemaker.serve.model_builder import ModelBuilder
11+
12+
ROLE_ARN = "arn:aws:iam::123456789012:role/TestRole"
13+
TELEMETRY_MODULE = "sagemaker.core.telemetry.telemetry_logging"
14+
NOVA_1P_IMAGE = "708977205387.dkr.ecr.us-east-1.amazonaws.com/nova-inference:latest"
15+
16+
17+
def _mock_session():
18+
session = Mock()
19+
session.boto_region_name = "us-west-2"
20+
session.default_bucket.return_value = "test-bucket"
21+
session.default_bucket_prefix = None
22+
session.config = {}
23+
session.sagemaker_config = {}
24+
return session
25+
26+
27+
class TestBuildValidationsWithPipelineVariableImageUri(unittest.TestCase):
28+
"""_build_validations must not slice a pipeline-variable image_uri."""
29+
30+
def test_image_only_pipeline_variable_is_passthrough(self):
31+
builder = ModelBuilder(
32+
image_uri=ParameterString(name="ServingImageUri"),
33+
role_arn=ROLE_ARN,
34+
sagemaker_session=_mock_session(),
35+
)
36+
37+
builder._build_validations()
38+
39+
self.assertTrue(builder._passthrough)
40+
41+
def test_pipeline_variable_with_model_requires_model_server(self):
42+
builder = ModelBuilder(
43+
model=Mock(),
44+
image_uri=ParameterString(name="ServingImageUri"),
45+
role_arn=ROLE_ARN,
46+
sagemaker_session=_mock_session(),
47+
)
48+
49+
with self.assertRaises(ValueError) as context:
50+
builder._build_validations()
51+
52+
self.assertIn("Model_server must be set", str(context.exception))
53+
54+
55+
@patch("sagemaker.serve.model_builder.resolve_nested_dict_value_from_config")
56+
@patch("sagemaker.serve.model_builder.resolve_value_from_config")
57+
@patch.object(ModelBuilder, "_init_sagemaker_session_if_does_not_exist")
58+
@patch.object(ModelBuilder, "_prepare_container_def")
59+
class TestCreateSageMakerModelNetworkIsolation(unittest.TestCase):
60+
"""The Nova network-isolation check must tolerate a pipeline-variable image."""
61+
62+
def _run(self, image_uri, mock_prepare, mock_resolve, mock_resolve_nested):
63+
mock_prepare.return_value = {"Image": image_uri, "Environment": {}}
64+
mock_resolve.side_effect = lambda value, *args, **kwargs: value
65+
mock_resolve_nested.side_effect = lambda value, *args, **kwargs: value
66+
67+
session = Mock(spec=PipelineSession)
68+
builder = ModelBuilder(
69+
image_uri=image_uri,
70+
role_arn=ROLE_ARN,
71+
sagemaker_session=_mock_session(),
72+
)
73+
builder.sagemaker_session = session
74+
builder.model_name = "test-model"
75+
76+
builder._create_sagemaker_model()
77+
78+
session.create_model.assert_called_once()
79+
return session.create_model.call_args.kwargs
80+
81+
def test_pipeline_variable_image_does_not_raise(
82+
self, mock_prepare, mock_init, mock_resolve, mock_resolve_nested
83+
):
84+
image_uri = ParameterString(name="ServingImageUri")
85+
86+
kwargs = self._run(image_uri, mock_prepare, mock_resolve, mock_resolve_nested)
87+
88+
self.assertIs(kwargs["container_defs"]["Image"], image_uri)
89+
self.assertFalse(kwargs["enable_network_isolation"])
90+
91+
def test_nova_1p_string_image_still_enables_network_isolation(
92+
self, mock_prepare, mock_init, mock_resolve, mock_resolve_nested
93+
):
94+
kwargs = self._run(NOVA_1P_IMAGE, mock_prepare, mock_resolve, mock_resolve_nested)
95+
96+
self.assertTrue(kwargs["enable_network_isolation"])
97+
98+
99+
@patch(f"{TELEMETRY_MODULE}.resolve_value_from_config", return_value=False)
100+
@patch(f"{TELEMETRY_MODULE}._send_telemetry_request")
101+
@patch("sagemaker.serve.model_builder.resolve_and_validate_role", return_value=ROLE_ARN)
102+
@patch.object(PipelineSession, "default_bucket", return_value="test-bucket")
103+
class TestBuildWithPipelineVariableImageUri(unittest.TestCase):
104+
"""End-to-end build() under a PipelineSession, as reported in #5760."""
105+
106+
def test_build_keeps_pipeline_variable_in_create_model_request(self, *_mocks):
107+
image_uri = ParameterString(name="ServingImageUri")
108+
session = PipelineSession(
109+
boto_session=boto3.Session(
110+
region_name="us-west-2",
111+
aws_access_key_id="testing",
112+
aws_secret_access_key="testing",
113+
)
114+
)
115+
builder = ModelBuilder(image_uri=image_uri, role_arn=ROLE_ARN, sagemaker_session=session)
116+
117+
step_args = builder.build()
118+
119+
self.assertIsInstance(step_args, _ModelStepArguments)
120+
self.assertIs(step_args.create_model_request["PrimaryContainer"]["Image"], image_uri)
121+
122+
123+
if __name__ == "__main__":
124+
unittest.main()

‎sagemaker-serve/tests/unit/validations/test_check_image_uri.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import unittest
2+
from sagemaker.core.workflow.functions import Join
3+
from sagemaker.core.workflow.parameters import ParameterString
24
from sagemaker.serve.validations.check_image_uri import (
35
is_1p_image_uri,
46
all_accounts,
@@ -73,6 +75,23 @@ def test_all_accounts_contains_known_accounts(self):
7375
self.assertIn("763104351884", all_accounts)
7476
self.assertIn("246618743249", all_accounts)
7577

78+
def test_is_1p_image_uri_pipeline_variable_returns_false(self):
79+
# A pipeline variable is only resolved at execution time, so it must not be
80+
# sliced (TypeError, #5760); it is treated as not 1P, even if its default
81+
# value would be a 1P image.
82+
image_uri = ParameterString(
83+
name="ServingImageUri",
84+
default_value="763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch:latest",
85+
)
86+
self.assertFalse(is_1p_image_uri(image_uri))
87+
88+
def test_is_1p_image_uri_pipeline_function_returns_false(self):
89+
image_uri = Join(on="", values=["763104351884", ParameterString(name="Suffix")])
90+
self.assertFalse(is_1p_image_uri(image_uri))
91+
92+
def test_is_1p_image_uri_none_returns_false(self):
93+
self.assertFalse(is_1p_image_uri(None))
94+
7695

7796
if __name__ == "__main__":
7897
unittest.main()

0 commit comments

Comments
 (0)