Refactoring

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -268,10 +268,6 @@ def _tune_for_js(self, max_tuning_duration: int = 1800):
             )
             return self.pysdk_model
 
-        if not sharded_supported(self.model, self.js_model_config):
-            logger.warning("Sharded is not supported for this model. Returning original model.")
-            return self.pysdk_model
-
         num_shard_env_var_name = "SM_NUM_GPUS"
         if "OPTION_TENSOR_PARALLEL_DEGREE" in self.pysdk_model.env.keys():
             num_shard_env_var_name = "OPTION_TENSOR_PARALLEL_DEGREE"
@@ -281,6 +277,15 @@ def _tune_for_js(self, max_tuning_duration: int = 1800):
             self.js_model_config
         )
 
+        if len(admissible_tensor_parallel_degrees) > 1 and not sharded_supported(
+            self.model, self.js_model_config
+        ):
+            admissible_tensor_parallel_degrees = [1]
+            logger.warning(
+                "Sharded across multiple GPUs is not supported for this model."
+                "\nModel can only be sharded across [1] GPU"
+            )
+
         benchmark_results = {}
         best_tuned_combination = None
         timeout = datetime.now() + timedelta(seconds=max_tuning_duration)

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

@@ -484,16 +484,31 @@ def test_tune_for_djl_js_local_container_invoke_ex(
         "sagemaker.serve.builder.jumpstart_builder.prepare_djl_js_resources",
         return_value=(
             mock_set_serving_properties,
-            {"model_type": "sharded_not_supported", "n_head": 71},
+            {"model_type": "sharded_not_enabled", "n_head": 71},
             True,
         ),
     )
     @patch("sagemaker.serve.builder.jumpstart_builder._get_ram_usage_mb", return_value=1024)
     @patch(
         "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
     )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder._get_admissible_tensor_parallel_degrees",
+        return_value=[4, 2, 1],
+    )
+    @patch(
+        "sagemaker.serve.utils.tuning._serial_benchmark",
+        side_effect=[(5, 5, 25), (5.4, 5.4, 20), (5.2, 5.2, 15)],
+    )
+    @patch(
+        "sagemaker.serve.utils.tuning._concurrent_benchmark",
+        side_effect=[(0.9, 1), (0.10, 4), (0.13, 2)],
+    )
     def test_tune_for_djl_js_local_container_sharded_not_enabled(
         self,
+        mock_concurrent_benchmarks,
+        mock_serial_benchmarks,
+        mock_admissible_tensor_parallel_degrees,
         mock_get_nb_instance,
         mock_get_ram_usage_mb,
         mock_prepare_for_tgi,
@@ -502,7 +517,9 @@ def test_tune_for_djl_js_local_container_sharded_not_enabled(
         mock_telemetry,
     ):
         builder = ModelBuilder(
-            model=mock_model_id, schema_builder=mock_schema_builder, mode=Mode.LOCAL_CONTAINER
+            model=mock_model_id,
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
         )
 
         mock_pre_trained_model.return_value.image_uri = mock_djl_image_uri
@@ -513,4 +530,8 @@ def test_tune_for_djl_js_local_container_sharded_not_enabled(
         mock_pre_trained_model.return_value.env = mock_djl_model_serving_properties
 
         tuned_model = model.tune()
-        assert tuned_model.env == mock_djl_model_serving_properties
+        assert tuned_model.env == {
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
+            "OPTION_TENSOR_PARALLEL_DEGREE": "1",
+        }