Bug fixes (#1496)

* Bug fixes

* refcatore

* ENV update

* Remove code duplication

* Fix Integ tests

* Fix MB EULA bug

---------

Co-authored-by: Jonathan Makunga <[email protected]>

Author: makungaj1

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -45,6 +45,7 @@
     _extracts_and_validates_speculative_model_source,
     _generate_channel_name,
     _generate_additional_model_data_sources,
+    _is_s3_uri,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -750,6 +751,8 @@ def _optimize_for_jumpstart(
 
         if pysdk_model_env_vars:
             self.pysdk_model.env.update(pysdk_model_env_vars)
+        if accept_eula:
+            self.pysdk_model.accept_eula = accept_eula
 
         if quantization_config or compilation_config:
             return create_optimization_job_args
@@ -787,8 +790,9 @@ def _set_additional_model_source(
         if speculative_decoding_config:
             model_provider = _extract_speculative_draft_model_provider(speculative_decoding_config)
             channel_name = _generate_channel_name(self.pysdk_model.additional_model_data_sources)
+            speculative_draft_model = f"/opt/ml/additional-model-data-sources/{channel_name}"
 
-            if model_provider.lower() == "sagemaker":
+            if model_provider == "sagemaker":
                 additional_model_data_sources = self.pysdk_model.deployment_config.get(
                     "DeploymentArgs", {}
                 ).get("AdditionalDataSources")
@@ -805,27 +809,31 @@ def _set_additional_model_source(
                         raise ValueError(
                             "Cannot find deployment config compatible for optimization job."
                         )
-
-                self.pysdk_model.add_tags(
-                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "sagemaker"},
-                )
             else:
-                s3_uri = _extracts_and_validates_speculative_model_source(
+                model_source = _extracts_and_validates_speculative_model_source(
                     speculative_decoding_config
                 )
 
-                self.pysdk_model.additional_model_data_sources = (
-                    _generate_additional_model_data_sources(s3_uri, channel_name, accept_eula)
-                )
-                self.pysdk_model.add_tags(
-                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "customer"},
-                )
+                if _is_s3_uri(model_source):
+                    self.pysdk_model.additional_model_data_sources = (
+                        _generate_additional_model_data_sources(
+                            model_source, channel_name, accept_eula
+                        )
+                    )
+                else:
+                    speculative_draft_model = model_source
 
-            speculative_draft_model = f"/opt/ml/additional-model-data-sources/{channel_name}"
             self.pysdk_model.env = _update_environment_variables(
                 self.pysdk_model.env,
                 {"OPTION_SPECULATIVE_DRAFT_MODEL": speculative_draft_model},
             )
+            self.pysdk_model.add_tags(
+                {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": model_provider},
+            )
+            if accept_eula and isinstance(self.pysdk_model.model_data, dict):
+                self.pysdk_model.model_data["S3DataSource"]["ModelAccessConfig"] = {
+                    "AcceptEula": True
+                }
 
     def _find_compatible_deployment_config(
         self, speculative_decoding_config: Optional[Dict] = None

src/sagemaker/serve/builder/model_builder.py

@@ -408,7 +408,8 @@ def _prepare_for_mode(self, should_upload_artifacts: bool = False):
                 getattr(self, "model_hub", None) == ModelHub.JUMPSTART,
                 should_upload=should_upload_artifacts,
             )
-            self.env_vars.update(env_vars_sagemaker)
+            if env_vars_sagemaker:
+                self.env_vars.update(env_vars_sagemaker)
             return self.s3_upload_path, env_vars_sagemaker
         if self.mode == Mode.LOCAL_CONTAINER:
             # init the LocalContainerMode object
@@ -1026,6 +1027,12 @@ def _model_builder_optimize_wrapper(
         )
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
+
+        if instance_type:
+            self.instance_type = instance_type
+        if role:
+            self.role = role
+
         self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
         job_name = job_name or f"modelbuilderjob-{uuid.uuid4().hex}"
 

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -70,7 +70,7 @@ def prepare(
                 + "session to be created or supply `sagemaker_session` into @serve.invoke."
             ) from e
 
-        upload_artifacts = None
+        upload_artifacts = None, None
         if self.model_server == ModelServer.TORCHSERVE:
             upload_artifacts = self._upload_torchserve_artifacts(
                 model_path=model_path,

src/sagemaker/serve/utils/optimize_utils.py

@@ -25,23 +25,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _is_inferentia_or_trainium(instance_type: Optional[str]) -> bool:
-    """Checks whether an instance is compatible with Inferentia.
-
-    Args:
-        instance_type (str): The instance type used for the compilation job.
-
-    Returns:
-        bool: Whether the given instance type is Inferentia or Trainium.
-    """
-    if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
-        if match:
-            if match[1].startswith("inf") or match[1].startswith("trn"):
-                return True
-    return False
-
-
 def _is_image_compatible_with_optimization_job(image_uri: Optional[str]) -> bool:
     """Checks whether an instance is compatible with an optimization job.
 
@@ -69,13 +52,16 @@ def _generate_optimized_model(pysdk_model: Model, optimization_response: dict) -
     """
     recommended_image_uri = optimization_response["OptimizationOutput"]["RecommendedInferenceImage"]
     optimized_environment = optimization_response["OptimizationEnvironment"]
-    s3_uri = optimization_response["ModelSource"]["S3"]
+    s3_uri = optimization_response["OutputConfig"]["S3OutputLocation"]
     deployment_instance_type = optimization_response["DeploymentInstanceType"]
 
     if recommended_image_uri:
         pysdk_model.image_uri = recommended_image_uri
     if optimized_environment:
-        pysdk_model.env = optimized_environment
+        if pysdk_model.env:
+            pysdk_model.env.update(optimized_environment)
+        else:
+            pysdk_model.env = optimized_environment
     if s3_uri:
         pysdk_model.model_data["S3DataSource"]["S3Uri"] = s3_uri
     if deployment_instance_type:
@@ -258,3 +244,18 @@ def _generate_additional_model_data_sources(
         additional_model_data_source["S3DataSource"]["ModelAccessConfig"] = {"ACCEPT_EULA": True}
 
     return [additional_model_data_source]
+
+
+def _is_s3_uri(s3_uri: Optional[str]) -> bool:
+    """Checks whether an S3 URI is valid.
+
+    Args:
+        s3_uri (Optional[str]): The S3 URI.
+
+    Returns:
+        bool: Whether the S3 URI is valid.
+    """
+    if s3_uri is None:
+        return False
+
+    return re.match("^s3://([^/]+)/?(.*)$", s3_uri) is not None

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

@@ -19,19 +19,21 @@
 from sagemaker.enums import Tag
 from sagemaker.serve.utils.optimize_utils import (
     _generate_optimized_model,
-    _is_inferentia_or_trainium,
     _update_environment_variables,
     _is_image_compatible_with_optimization_job,
     _extract_speculative_draft_model_provider,
     _validate_optimization_inputs,
     _extracts_and_validates_speculative_model_source,
+    _is_s3_uri,
+    _generate_additional_model_data_sources,
+    _generate_channel_name,
 )
 
 mock_optimization_job_output = {
-    "OptimizationJobArn": "arn:aws:sagemaker:us-west-2:312206380606:"
-    "optimization-job/modelbuilderjob-6b09ffebeb0741b8a28b85623fd9c968",
+    "OptimizationJobArn": "arn:aws:sagemaker:us-west-2:312206380606:optimization-job/"
+    "modelbuilderjob-3cbf9c40b63c455d85b60033f9a01691",
     "OptimizationJobStatus": "COMPLETED",
-    "OptimizationJobName": "modelbuilderjob-6b09ffebeb0741b8a28b85623fd9c968",
+    "OptimizationJobName": "modelbuilderjob-3cbf9c40b63c455d85b60033f9a01691",
     "ModelSource": {
         "S3": {
             "S3Uri": "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/"
@@ -46,7 +48,7 @@
         "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
         "SAGEMAKER_PROGRAM": "inference.py",
     },
-    "DeploymentInstanceType": "ml.g5.48xlarge",
+    "DeploymentInstanceType": "ml.g5.2xlarge",
     "OptimizationConfigs": [
         {
             "ModelQuantizationConfig": {
@@ -55,40 +57,26 @@
             }
         }
     ],
-    "OutputConfig": {
-        "S3OutputLocation": "s3://dont-delete-ss-jarvis-integ-test-312206380606-us-west-2/"
-    },
+    "OutputConfig": {"S3OutputLocation": "s3://quicksilver-model-data/llama-3-8b/quantized-1/"},
     "OptimizationOutput": {
         "RecommendedInferenceImage": "763104351884.dkr.ecr.us-west-2.amazonaws.com/djl-inference:0.28.0-lmi10.0.0-cu124"
     },
-    "RoleArn": "arn:aws:iam::312206380606:role/service-role/AmazonSageMaker-ExecutionRole-20230707T131628",
+    "RoleArn": "arn:aws:iam::312206380606:role/service-role/AmazonSageMaker-ExecutionRole-20240116T151132",
     "StoppingCondition": {"MaxRuntimeInSeconds": 36000},
     "ResponseMetadata": {
-        "RequestId": "17ae151f-b51d-4194-8ba9-edbba068c90b",
+        "RequestId": "a95253d5-c045-4708-8aac-9f0d327515f7",
         "HTTPStatusCode": 200,
         "HTTPHeaders": {
-            "x-amzn-requestid": "17ae151f-b51d-4194-8ba9-edbba068c90b",
+            "x-amzn-requestid": "a95253d5-c045-4708-8aac-9f0d327515f7",
             "content-type": "application/x-amz-json-1.1",
-            "content-length": "1380",
-            "date": "Thu, 20 Jun 2024 19:25:53 GMT",
+            "content-length": "1371",
+            "date": "Fri, 21 Jun 2024 04:27:42 GMT",
         },
         "RetryAttempts": 0,
     },
 }
 
 
-@pytest.mark.parametrize(
-    "instance, expected",
-    [
-        ("ml.trn1.2xlarge", True),
-        ("ml.inf2.xlarge", True),
-        ("ml.c7gd.4xlarge", False),
-    ],
-)
-def test_is_inferentia_or_trainium(instance, expected):
-    assert _is_inferentia_or_trainium(instance) == expected
-
-
 @pytest.mark.parametrize(
     "image_uri, expected",
     [
@@ -124,17 +112,21 @@ def test_generate_optimized_model():
             "meta-textgeneration-llama-3-8b/artifacts/inference-prepack/v1.0.1/"
         }
     }
+    pysdk_model.env = {"OPTION_QUANTIZE": "awq"}
 
     optimized_model = _generate_optimized_model(pysdk_model, mock_optimization_job_output)
 
     assert (
         optimized_model.image_uri
         == mock_optimization_job_output["OptimizationOutput"]["RecommendedInferenceImage"]
     )
-    assert optimized_model.env == mock_optimization_job_output["OptimizationEnvironment"]
+    assert optimized_model.env == {
+        "OPTION_QUANTIZE": "awq",
+        **mock_optimization_job_output["OptimizationEnvironment"],
+    }
     assert (
         optimized_model.model_data["S3DataSource"]["S3Uri"]
-        == mock_optimization_job_output["ModelSource"]["S3"]
+        == mock_optimization_job_output["OutputConfig"]["S3OutputLocation"]
     )
     assert optimized_model.instance_type == mock_optimization_job_output["DeploymentInstanceType"]
     pysdk_model.add_tags.assert_called_once_with(
@@ -209,3 +201,61 @@ def test_extract_speculative_draft_model_s3_uri():
 def test_extract_speculative_draft_model_s3_uri_ex():
     with pytest.raises(ValueError):
         _extracts_and_validates_speculative_model_source({"ModelSource": None})
+
+
+def test_generate_channel_name():
+    assert _generate_channel_name(None) is not None
+
+    additional_model_data_sources = _generate_additional_model_data_sources(
+        "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/", "channel_name", True
+    )
+
+    assert _generate_channel_name(additional_model_data_sources) == "channel_name"
+
+
+def test_generate_additional_model_data_sources():
+    model_source = _generate_additional_model_data_sources(
+        "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/", "channel_name", True
+    )
+
+    assert model_source == [
+        {
+            "ChannelName": "channel_name",
+            "S3DataSource": {
+                "S3Uri": "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/",
+                "S3DataType": "S3Prefix",
+                "CompressionType": "None",
+                "ModelAccessConfig": {"ACCEPT_EULA": True},
+            },
+        }
+    ]
+
+    model_source = _generate_additional_model_data_sources(
+        "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/", "channel_name", False
+    )
+
+    assert model_source == [
+        {
+            "ChannelName": "channel_name",
+            "S3DataSource": {
+                "S3Uri": "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/",
+                "S3DataType": "S3Prefix",
+                "CompressionType": "None",
+            },
+        }
+    ]
+
+
+@pytest.mark.parametrize(
+    "s3_uri, expected",
+    [
+        (
+            "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/"
+            "meta-textgeneration-llama-3-8b/artifacts/inference-prepack/v1.0.1/",
+            True,
+        ),
+        ("invalid://", False),
+    ],
+)
+def test_is_s3_uri(s3_uri, expected):
+    assert _is_s3_uri(s3_uri) == expected