add more UTs for sharding

Author: gwang111

src/sagemaker/serve/builder/model_builder.py

@@ -1234,7 +1234,11 @@ def _model_builder_optimize_wrapper(
         Returns:
             Model: A deployable ``Model`` object.
         """
-        if hasattr(self, "enable_network_isolation") and sharding_config:
+        if (
+            hasattr(self, "enable_network_isolation")
+            and self.enable_network_isolation
+            and sharding_config
+        ):
             raise ValueError(
                 "EnableNetworkIsolation cannot be set to True since SageMaker Fast Model "
                 "Loading of model requires network access."
@@ -1272,19 +1276,25 @@ def _model_builder_optimize_wrapper(
                 )
             )
 
-        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"]
+        if sharding_config:
+            has_tensor_parallel_degree_in_env_vars = (
+                env_vars and "OPTION_TENSOR_PARALLEL_DEGREE" in env_vars
             )
-        ):
-            raise ValueError(
-                (
-                    "OPTION_TENSOR_PARALLEL_DEGREE is a required "
-                    "environment variable with sharding config."
-                )
+            has_tensor_parallel_degree_in_overrides = (
+                sharding_config
+                and sharding_config.get("OverrideEnvironment")
+                and "OPTION_TENSOR_PARALLEL_DEGREE" in sharding_config.get("OverrideEnvironment")
             )
+            if (
+                not has_tensor_parallel_degree_in_env_vars
+                and not has_tensor_parallel_degree_in_overrides
+            ):
+                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
@@ -1399,10 +1409,11 @@ def _optimize_for_hf(
             create_optimization_job_args["ModelSource"] = model_source
 
             optimization_config, override_env = _extract_optimization_config_and_env(
-                quantization_config, compilation_config
+                quantization_config, compilation_config, sharding_config
             )
             create_optimization_job_args["OptimizationConfigs"] = [optimization_config]
-            self.pysdk_model.env.update(override_env)
+            if override_env:
+                self.pysdk_model.env.update(override_env)
 
             output_config = {"S3OutputLocation": output_path}
             if kms_key: