add unit tests

Author: Ashish Gupta

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

@@ -1198,6 +1198,57 @@ def test_optimize_quantize_for_jumpstart(
 
         self.assertIsNotNone(out_put)
 
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_optimize_sharding_for_jumpstart(
+        self,
+        mock_serve_settings,
+        mock_telemetry,
+    ):
+        mock_sagemaker_session = Mock()
+
+        mock_pysdk_model = Mock()
+        mock_pysdk_model.env = {"SAGEMAKER_ENV": "1"}
+        mock_pysdk_model.model_data = mock_model_data
+        mock_pysdk_model.image_uri = mock_tgi_image_uri
+        mock_pysdk_model.list_deployment_configs.return_value = DEPLOYMENT_CONFIGS
+        mock_pysdk_model.deployment_config = DEPLOYMENT_CONFIGS[0]
+
+        sample_input = {
+            "inputs": "The diamondback terrapin or simply terrapin is a species "
+            "of turtle native to the brackish coastal tidal marshes of the",
+            "parameters": {"max_new_tokens": 1024},
+        }
+        sample_output = [
+            {
+                "generated_text": "The diamondback terrapin or simply terrapin is a "
+                "species of turtle native to the brackish coastal "
+                "tidal marshes of the east coast."
+            }
+        ]
+
+        model_builder = ModelBuilder(
+            model="meta-textgeneration-llama-3-70b",
+            schema_builder=SchemaBuilder(sample_input, sample_output),
+            sagemaker_session=mock_sagemaker_session,
+        )
+
+        model_builder.pysdk_model = mock_pysdk_model
+
+        out_put = model_builder._optimize_for_jumpstart(
+            accept_eula=True,
+            sharding_config={
+                "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
+            },
+            env_vars={
+                "OPTION_TENSOR_PARALLEL_DEGREE": "1",
+                "OPTION_MAX_ROLLING_BATCH_SIZE": "2",
+            },
+            output_path="s3://bucket/code/",
+        )
+
+        self.assertIsNotNone(out_put)
+
     @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
     @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
     @patch(

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

@@ -2667,6 +2667,39 @@ def test_optimize_exclusive_args(self, mock_get_serve_setting):
             ),
         )
 
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_optimize_exclusive_sharding(self, mock_get_serve_setting):
+        mock_sagemaker_session = Mock()
+        model_builder = ModelBuilder(
+            model="meta-textgeneration-llama-3-70b",
+            sagemaker_session=mock_sagemaker_session,
+        )
+
+        self.assertRaisesRegex(
+            ValueError,
+            "Sharding config is mutually exclusive and cannot be combined with any other optimization.",
+            lambda: model_builder.optimize(
+                compilation_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
+                sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
+            ),
+        )
+
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_optimize_exclusive_sharding_args(self, mock_get_serve_setting):
+        mock_sagemaker_session = Mock()
+        model_builder = ModelBuilder(
+            model="meta-textgeneration-llama-3-70b",
+            sagemaker_session=mock_sagemaker_session,
+        )
+
+        self.assertRaisesRegex(
+            ValueError,
+            "OPTION_TENSOR_PARALLEL_DEGREE is required environment variable with Sharding config.",
+            lambda: model_builder.optimize(
+                sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
+            ),
+        )
+
     @patch.object(ModelBuilder, "_prepare_for_mode")
     @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
     def test_optimize_for_hf_with_custom_s3_path(

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

@@ -261,7 +261,7 @@ def test_is_s3_uri(s3_uri, expected):
 
 
 @pytest.mark.parametrize(
-    "quantization_config, compilation_config, expected_config, expected_env",
+    "quantization_config, compilation_config, sharding_config, expected_config, expected_env",
     [
         (
             None,
@@ -270,6 +270,7 @@ def test_is_s3_uri(s3_uri, expected):
                     "OPTION_TENSOR_PARALLEL_DEGREE": "2",
                 }
             },
+            None,
             {
                 "ModelCompilationConfig": {
                     "OverrideEnvironment": {
@@ -288,6 +289,7 @@ def test_is_s3_uri(s3_uri, expected):
                 }
             },
             None,
+            None,
             {
                 "ModelQuantizationConfig": {
                     "OverrideEnvironment": {
@@ -299,13 +301,32 @@ def test_is_s3_uri(s3_uri, expected):
                 "OPTION_TENSOR_PARALLEL_DEGREE": "2",
             },
         ),
-        (None, None, None, None),
+        (
+            None,
+            None,
+            {
+                "OverrideEnvironment": {
+                    "OPTION_TENSOR_PARALLEL_DEGREE": "2",
+                }
+            },
+            {
+                "ModelShardingConfig": {
+                    "OverrideEnvironment": {
+                        "OPTION_TENSOR_PARALLEL_DEGREE": "2",
+                    }
+                },
+            },
+            {
+                "OPTION_TENSOR_PARALLEL_DEGREE": "2",
+            },
+        ),
+        (None, None, None, None, None),
     ],
 )
 def test_extract_optimization_config_and_env(
-    quantization_config, compilation_config, expected_config, expected_env
+    quantization_config, compilation_config, sharding_config, expected_config, expected_env
 ):
-    assert _extract_optimization_config_and_env(quantization_config, compilation_config) == (
+    assert _extract_optimization_config_and_env(quantization_config, compilation_config, sharding_config) == (
         expected_config,
         expected_env,
     )