Skip to content

Commit 8b84605

Browse files
committed
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.
1 parent ab5f478 commit 8b84605

2 files changed

Lines changed: 13 additions & 0 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3621,6 +3621,12 @@ def from_jumpstart_config(
36213621
)
36223622
mb_instance.inference_ami_version = deploy_kwargs.get("inference_ami_version")
36233623

3624+
# Apply network isolation from JumpStart model spec if not set by user via network param
3625+
if not mb_instance._enable_network_isolation and deploy_kwargs.get(
3626+
"enable_network_isolation"
3627+
):
3628+
mb_instance._enable_network_isolation = deploy_kwargs["enable_network_isolation"]
3629+
36243630
return mb_instance
36253631

36263632
@_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="model_builder.transformer")

sagemaker-serve/src/sagemaker/serve/model_builder_servers.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -877,6 +877,13 @@ def _build_for_jumpstart(self) -> Model:
877877
if hasattr(init_kwargs, "env") and init_kwargs.env:
878878
self.env_vars.update(init_kwargs.env)
879879

880+
# Apply network isolation from JumpStart model spec if not already set by user
881+
if (
882+
not self._enable_network_isolation
883+
and getattr(init_kwargs, "enable_network_isolation", None) is not None
884+
):
885+
self._enable_network_isolation = init_kwargs.enable_network_isolation
886+
880887
# Handle model artifacts for fine-tuned models
881888
if hasattr(init_kwargs, "model_data") and init_kwargs.model_data:
882889
if (

0 commit comments

Comments
 (0)