Disable network isolation if using sharded models.

Author: Joseph Zhang

src/sagemaker/model.py

@@ -1601,10 +1601,18 @@ def deploy(
                 self._base_name = "-".join((self._base_name, compiled_model_suffix))
 
         if self._is_sharded_model and endpoint_type != EndpointType.INFERENCE_COMPONENT_BASED:
-            logging.warning("Forcing INFERENCE_COMPONENT_BASED endpoint for sharded model. ADVISORY - "
-                            "Use INFERENCE_COMPONENT_BASED endpoints over MODEL_BASED endpoints.")
+            logging.warning(
+                "Forcing INFERENCE_COMPONENT_BASED endpoint for sharded model. ADVISORY - "
+                "Use INFERENCE_COMPONENT_BASED endpoints over MODEL_BASED endpoints."
+            )
             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

@@ -1245,11 +1245,25 @@ def _model_builder_optimize_wrapper(
         if quantization_config and compilation_config:
             raise ValueError("Quantization config and compilation config are mutually exclusive.")
 
-        if sharding_config and (quantization_config or compilation_config or speculative_decoding_config):
-            raise ValueError("Sharding config is mutually exclusive and cannot be combined with any other optimization.")
+        if sharding_config and (
+            quantization_config or compilation_config or speculative_decoding_config
+        ):
+            raise ValueError(
+                "Sharding config is mutually exclusive and cannot be combined with any "
+                "other optimization."
+            )
 
-        if sharding_config and ((env_vars and "OPTION_TENSOR_PARALLEL_DEGREE" not in env_vars) or (sharding_config.get("OverrideEnvironment") and "OPTION_TENSOR_PARALLEL_DEGREE" not in sharding_config["OverrideEnvironment"])):
-            raise ValueError("OPTION_TENSOR_PARALLEL_DEGREE is required environment variable with Sharding config.")
+        if sharding_config and (
+            (env_vars and "OPTION_TENSOR_PARALLEL_DEGREE" not in env_vars)
+            or (
+                sharding_config.get("OverrideEnvironment")
+                and "OPTION_TENSOR_PARALLEL_DEGREE" not in sharding_config["OverrideEnvironment"]
+            )
+        ):
+            raise ValueError(
+                "OPTION_TENSOR_PARALLEL_DEGREE is a required environment variable with "
+                "sharding config."
+            )
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
         self.instance_type = instance_type or self.instance_type

src/sagemaker/serve/utils/optimize_utils.py

@@ -259,8 +259,9 @@ def _is_s3_uri(s3_uri: Optional[str]) -> bool:
 
 
 def _extract_optimization_config_and_env(
-    quantization_config: Optional[Dict] = None, compilation_config: Optional[Dict] = None,
-    sharding_config: Optional[Dict] = None
+    quantization_config: Optional[Dict] = None,
+    compilation_config: Optional[Dict] = None,
+    sharding_config: Optional[Dict] = None,
 ) -> Optional[Tuple[Optional[Dict], Optional[Dict]]]:
     """Extracts optimization config and environment variables.
 
@@ -282,9 +283,7 @@ def _extract_optimization_config_and_env(
             "OverrideEnvironment"
         )
     if sharding_config:
-        return {"ModelShardingConfig": sharding_config}, sharding_config.get(
-            "OverrideEnvironment"
-        )
+        return {"ModelShardingConfig": sharding_config}, sharding_config.get("OverrideEnvironment")
     return None, None
 
 

tests/unit/sagemaker/model/test_model.py

@@ -958,6 +958,7 @@ def test_all_framework_models_inference_component_based_endpoint_deploy_path(
         sagemaker_session.endpoint_in_service_or_not.reset_mock()
         sagemaker_session.create_model.reset_mock()
 
+
 @patch("sagemaker.utils.repack_model")
 @patch("sagemaker.fw_utils.tar_and_upload_dir")
 def test_sharded_model_force_inference_component_based_endpoint_deploy_path(
@@ -967,7 +968,7 @@ def test_sharded_model_force_inference_component_based_endpoint_deploy_path(
         HuggingFaceModel: {
             "pytorch_version": "1.7.1",
             "py_version": "py36",
-            "transformers_version": "4.6.1"
+            "transformers_version": "4.6.1",
         },
     }
 
@@ -1007,6 +1008,7 @@ def test_sharded_model_force_inference_component_based_endpoint_deploy_path(
         sagemaker_session.endpoint_in_service_or_not.reset_mock()
         sagemaker_session.create_model.reset_mock()
 
+
 @patch("sagemaker.utils.repack_model")
 def test_repack_code_location_with_key_prefix(repack_model, sagemaker_session):
 

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

@@ -2694,7 +2694,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(
                 sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
             ),

tests/unit/sagemaker/serve/utils/test_optimize_utils.py

@@ -326,7 +326,9 @@ def test_is_s3_uri(s3_uri, expected):
 def test_extract_optimization_config_and_env(
     quantization_config, compilation_config, sharding_config, expected_config, expected_env
 ):
-    assert _extract_optimization_config_and_env(quantization_config, compilation_config, sharding_config) == (
+    assert _extract_optimization_config_and_env(
+        quantization_config, compilation_config, sharding_config
+    ) == (
         expected_config,
         expected_env,
     )