From 8b84605d6c443b6b904009b327190c98cf5efeac Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 6 May 2026 19:56:24 +0000 Subject: [PATCH 1/5] fix: apply JumpStart model spec enable_network_isolation in ModelBuilder In v2, JumpStart models deployed with EnableNetworkIsolation=True because the JumpStartModel class read inference_enable_network_isolation from the model spec JSON. In v3's ModelBuilder, this field was never extracted from the JumpStart init_kwargs or deploy_kwargs, causing all JumpStart models to deploy without network isolation (defaulting to False). This fix applies enable_network_isolation from the JumpStart model spec in both code paths: 1. _build_for_jumpstart() - regular build path via get_init_kwargs() 2. from_jumpstart_config() - classmethod path via _retrieve_model_deploy_kwargs() The user-provided value (via network= param) takes precedence. --- sagemaker-serve/src/sagemaker/serve/model_builder.py | 6 ++++++ .../src/sagemaker/serve/model_builder_servers.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 7c7af2defc..6f60d325ad 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -3621,6 +3621,12 @@ def from_jumpstart_config( ) mb_instance.inference_ami_version = deploy_kwargs.get("inference_ami_version") + # Apply network isolation from JumpStart model spec if not set by user via network param + if not mb_instance._enable_network_isolation and deploy_kwargs.get( + "enable_network_isolation" + ): + mb_instance._enable_network_isolation = deploy_kwargs["enable_network_isolation"] + return mb_instance @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.transformer") diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py index 64c76f59cf..7156ee20ee 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py @@ -877,6 +877,13 @@ def _build_for_jumpstart(self) -> Model: if hasattr(init_kwargs, "env") and init_kwargs.env: self.env_vars.update(init_kwargs.env) + # Apply network isolation from JumpStart model spec if not already set by user + if ( + not self._enable_network_isolation + and getattr(init_kwargs, "enable_network_isolation", None) is not None + ): + self._enable_network_isolation = init_kwargs.enable_network_isolation + # Handle model artifacts for fine-tuned models if hasattr(init_kwargs, "model_data") and init_kwargs.model_data: if ( From 0789d885546d307e904debc34d4b47f198fbc28d Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 6 May 2026 20:05:01 +0000 Subject: [PATCH 2/5] test: add unit tests for JumpStart network isolation fix Tests both code paths: 1. _build_for_jumpstart() - verifies enable_network_isolation is applied from init_kwargs and that user-set values are not overridden 2. from_jumpstart_config() - verifies enable_network_isolation is extracted from deploy_kwargs --- .../unit/test_model_builder_coverage_boost.py | 24 +++++++ .../test_model_builder_servers_coverage.py | 65 +++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py index 1f9aa2f3d4..14e2456dca 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py +++ b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py @@ -381,6 +381,30 @@ def test_from_jumpstart_config_basic(self): self.assertEqual(mb.model, "test-model") self.assertEqual(mb.model_version, "1.0.0") + @patch("sagemaker.core.jumpstart.artifacts.kwargs._retrieve_model_deploy_kwargs") + def test_from_jumpstart_config_applies_network_isolation(self, mock_deploy_kwargs): + """Test that enable_network_isolation from deploy kwargs is applied.""" + from sagemaker.core.jumpstart.configs import JumpStartConfig + from sagemaker.core.training.configs import Compute + + mock_deploy_kwargs.return_value = { + "model_data_download_timeout": 600, + "enable_network_isolation": True, + } + + js_config = JumpStartConfig( + model_id="test-model", + model_version="1.0.0" + ) + + mb = ModelBuilder.from_jumpstart_config( + jumpstart_config=js_config, + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + compute=Compute(instance_type="ml.g5.xlarge"), + ) + + self.assertTrue(mb._enable_network_isolation) + if __name__ == "__main__": unittest.main() diff --git a/sagemaker-serve/tests/unit/test_model_builder_servers_coverage.py b/sagemaker-serve/tests/unit/test_model_builder_servers_coverage.py index 3b00264ffa..fbd4366882 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_servers_coverage.py +++ b/sagemaker-serve/tests/unit/test_model_builder_servers_coverage.py @@ -501,6 +501,71 @@ def test_build_for_jumpstart_routes_to_mms(self, mock_prepare, mock_create, mock mock_create.assert_called_once() + @patch("sagemaker.core.jumpstart.factory.utils.get_init_kwargs") + @patch("sagemaker.serve.model_builder.ModelBuilder._create_model") + @patch("sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode") + def test_build_for_jumpstart_applies_network_isolation_from_spec( + self, mock_prepare, mock_create, mock_get_kwargs + ): + """Test that enable_network_isolation from JumpStart model spec is applied.""" + mock_init_kwargs = Mock() + mock_init_kwargs.image_uri = ( + "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.0-cu117" + ) + mock_init_kwargs.env = {} + mock_init_kwargs.model_data = "s3://jumpstart-cache/models/model.tar.gz" + mock_init_kwargs.enable_network_isolation = True + mock_get_kwargs.return_value = mock_init_kwargs + + mock_model = Mock(spec=Model) + mock_create.return_value = mock_model + + builder = ModelBuilder( + model="meta-textgeneration-llama-3-8b", + role_arn=MOCK_ROLE_ARN, + sagemaker_session=self.mock_session, + mode=Mode.SAGEMAKER_ENDPOINT, + ) + builder._optimizing = False + + builder._build_for_jumpstart() + + self.assertTrue(builder._enable_network_isolation) + + @patch("sagemaker.core.jumpstart.factory.utils.get_init_kwargs") + @patch("sagemaker.serve.model_builder.ModelBuilder._create_model") + @patch("sagemaker.serve.model_builder.ModelBuilder._prepare_for_mode") + def test_build_for_jumpstart_does_not_override_user_network_isolation( + self, mock_prepare, mock_create, mock_get_kwargs + ): + """Test that user-set network isolation is not overridden by spec.""" + mock_init_kwargs = Mock() + mock_init_kwargs.image_uri = ( + "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.0-cu117" + ) + mock_init_kwargs.env = {} + mock_init_kwargs.model_data = "s3://jumpstart-cache/models/model.tar.gz" + mock_init_kwargs.enable_network_isolation = False + mock_get_kwargs.return_value = mock_init_kwargs + + mock_model = Mock(spec=Model) + mock_create.return_value = mock_model + + builder = ModelBuilder( + model="meta-textgeneration-llama-3-8b", + role_arn=MOCK_ROLE_ARN, + sagemaker_session=self.mock_session, + mode=Mode.SAGEMAKER_ENDPOINT, + ) + builder._optimizing = False + builder._enable_network_isolation = True # User explicitly set + + builder._build_for_jumpstart() + + # User's True should not be overridden by spec's False + self.assertTrue(builder._enable_network_isolation) + + class TestDeployWrappers(unittest.TestCase): """Test deploy wrapper methods.""" From 0bbb83d738ae9cfe09bc329907a9df6282586fe7 Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 6 May 2026 20:05:01 +0000 Subject: [PATCH 3/5] test: add unit tests for JumpStart network isolation fix Tests both code paths: 1. _build_for_jumpstart() - verifies enable_network_isolation is applied from init_kwargs and that user-set values are not overridden 2. from_jumpstart_config() - verifies enable_network_isolation is extracted from deploy_kwargs --- .../tests/unit/test_model_builder_coverage_boost.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py index 14e2456dca..fafd132547 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py +++ b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py @@ -381,7 +381,7 @@ def test_from_jumpstart_config_basic(self): self.assertEqual(mb.model, "test-model") self.assertEqual(mb.model_version, "1.0.0") - @patch("sagemaker.core.jumpstart.artifacts.kwargs._retrieve_model_deploy_kwargs") + @patch("sagemaker.serve.model_builder._retrieve_model_deploy_kwargs") def test_from_jumpstart_config_applies_network_isolation(self, mock_deploy_kwargs): """Test that enable_network_isolation from deploy kwargs is applied.""" from sagemaker.core.jumpstart.configs import JumpStartConfig @@ -397,10 +397,14 @@ def test_from_jumpstart_config_applies_network_isolation(self, mock_deploy_kwarg model_version="1.0.0" ) + mock_session = Mock() + mock_session.boto_region_name = "us-west-2" + mb = ModelBuilder.from_jumpstart_config( jumpstart_config=js_config, role_arn="arn:aws:iam::123456789012:role/SageMakerRole", compute=Compute(instance_type="ml.g5.xlarge"), + sagemaker_session=mock_session, ) self.assertTrue(mb._enable_network_isolation) From 9c2fc8135272f2e55d618926daad7dfc65b8315a Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 6 May 2026 21:19:59 +0000 Subject: [PATCH 4/5] fix: add _enable_network_isolation to MockModelBuilderServers test fixture The mock directly inherits _ModelBuilderServers without going through ModelBuilder.__post_init__ which initializes this attribute. --- sagemaker-serve/tests/unit/servers/test_model_builder_servers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py index d905b7decf..053b14f416 100644 --- a/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py +++ b/sagemaker-serve/tests/unit/servers/test_model_builder_servers.py @@ -47,6 +47,7 @@ def __init__(self): self.framework_version = None self._is_mlflow_model = False self.config_name = None + self._enable_network_isolation = False def _deploy_local_endpoint(self, **kwargs): return Mock() From cb5e63d807381f9811a4fc14b7cf5d39ba5db86b Mon Sep 17 00:00:00 2001 From: Liam Neal Reilly Date: Wed, 6 May 2026 21:26:11 +0000 Subject: [PATCH 5/5] test: add integ test for JumpStart network isolation + fix mock fixture Integration test builds a real SageMaker Model (falcon-7b) and verifies EnableNetworkIsolation=True via describe_model API. No endpoint deployed. Also adds _enable_network_isolation to MockModelBuilderServers fixture to fix 7 pre-existing unit test failures. --- .../integ/test_jumpstart_network_isolation.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 sagemaker-serve/tests/integ/test_jumpstart_network_isolation.py diff --git a/sagemaker-serve/tests/integ/test_jumpstart_network_isolation.py b/sagemaker-serve/tests/integ/test_jumpstart_network_isolation.py new file mode 100644 index 0000000000..0dd2f214d2 --- /dev/null +++ b/sagemaker-serve/tests/integ/test_jumpstart_network_isolation.py @@ -0,0 +1,67 @@ +# +# 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. +from __future__ import absolute_import + +import uuid +import pytest +import logging + +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.core.jumpstart.configs import JumpStartConfig +from sagemaker.train.configs import Compute + +logger = logging.getLogger(__name__) + +MODEL_ID = "huggingface-llm-falcon-7b-bf16" +INSTANCE_TYPE = "ml.g5.2xlarge" +MODEL_NAME_PREFIX = "js-netiso-test" + + +@pytest.mark.slow_test +def test_jumpstart_build_enables_network_isolation(): + """Integration test verifying JumpStart models are built with EnableNetworkIsolation. + + JumpStart model specs define inference_enable_network_isolation=True for most models. + This test validates that ModelBuilder.build() propagates this setting to the + SageMaker Model resource, matching v2 behavior. + """ + logger.info("Starting JumpStart network isolation integration test...") + + compute = Compute(instance_type=INSTANCE_TYPE) + jumpstart_config = JumpStartConfig(model_id=MODEL_ID) + model_builder = ModelBuilder.from_jumpstart_config( + jumpstart_config=jumpstart_config, compute=compute + ) + unique_id = str(uuid.uuid4())[:8] + + core_model = model_builder.build(model_name=f"{MODEL_NAME_PREFIX}-{unique_id}") + logger.info(f"Model created: {core_model.model_name}") + + try: + # Verify ModelBuilder picked up network isolation from spec + assert model_builder._enable_network_isolation, ( + f"ModelBuilder._enable_network_isolation should be True for {MODEL_ID}, " + f"got {model_builder._enable_network_isolation}" + ) + + # Verify the actual SageMaker Model resource has EnableNetworkIsolation=True + sm_client = model_builder.sagemaker_session.sagemaker_client + desc = sm_client.describe_model(ModelName=core_model.model_name) + assert desc.get("EnableNetworkIsolation") is True, ( + f"SageMaker Model should have EnableNetworkIsolation=True, " + f"got {desc.get('EnableNetworkIsolation')}" + ) + + logger.info("✅ Network isolation correctly applied to SageMaker Model") + finally: + core_model.delete() + logger.info("Model deleted.")