Disable network isolation if using sharded models.

Author: Joseph Zhang

src/sagemaker/model.py

@@ -1607,6 +1607,12 @@ def deploy(
             )
             endpoint_type = EndpointType.INFERENCE_COMPONENT_BASED
 
+        if self._is_sharded_model and self._enable_network_isolation:
+            raise ValueError(
+                "EnableNetworkIsolation cannot be set to True since SageMaker Fast Model "
+                "Loading of model requires network access."
+            )
+
         # Support multiple models on same endpoint
         if endpoint_type == EndpointType.INFERENCE_COMPONENT_BASED:
             if endpoint_name:

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -795,6 +795,14 @@ def _optimize_for_jumpstart(
         optimization_env_vars = _update_environment_variables(optimization_env_vars, override_env)
         if optimization_env_vars:
             self.pysdk_model.env.update(optimization_env_vars)
+
+        if sharding_config and self.pysdk_model._enable_network_isolation:
+            logger.warning(
+                "EnableNetworkIsolation cannot be set to True since SageMaker Fast Model "
+                "Loading of model requires network access. Setting it to False."
+            )
+            self.pysdk_model._enable_network_isolation = False
+
         if quantization_config or sharding_config or is_compilation:
             return create_optimization_job_args
         return None

src/sagemaker/serve/builder/model_builder.py

@@ -1276,8 +1276,8 @@ def _model_builder_optimize_wrapper(
         ):
             raise ValueError(
                 (
-                    "OPTION_TENSOR_PARALLEL_DEGREE is required "
-                    "environment variable with Sharding config."
+                    "OPTION_TENSOR_PARALLEL_DEGREE is a required "
+                    "environment variable with sharding config."
                 )
             )
 

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -2701,7 +2701,7 @@ def test_optimize_exclusive_sharding_args(self, mock_get_serve_setting):
 
         self.assertRaisesRegex(
             ValueError,
-            "OPTION_TENSOR_PARALLEL_DEGREE is required environment variable with Sharding config.",
+            "OPTION_TENSOR_PARALLEL_DEGREE is a required environment variable with sharding config.",
             lambda: model_builder.optimize(
                 instance_type="ml.g5.24xlarge",
                 sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},