chore: modify metadata keys, add unit test for custom attributes

Author: evakravi

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -19,6 +19,9 @@
 from sagemaker.jumpstart.utils import (
     verify_model_region_and_return_specs,
 )
+from sagemaker.jumpstart.enums import (
+    JumpStartScriptScope,
+)
 
 
 def _retrieve_model_package_arn(
@@ -62,4 +65,7 @@ def _retrieve_model_package_arn(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
-    return model_specs.model_package_arn
+    if scope == JumpStartScriptScope.INFERENCE:
+        return model_specs.hosting_model_package_arn
+
+    raise NotImplementedError(f"Model Package ARN not supported for scope: '{scope}'")

src/sagemaker/jumpstart/types.py

@@ -351,8 +351,8 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "inference_enable_network_isolation",
         "training_enable_network_isolation",
         "resource_name_base",
-        "eula_key",
-        "model_package_arn",
+        "hosting_eula_key",
+        "hosting_model_package_arn",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -421,9 +421,9 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         )
         self.resource_name_base: bool = json_obj.get("resource_name_base")
 
-        self.eula_key: Optional[str] = json_obj.get("eula_key")
+        self.hosting_eula_key: Optional[str] = json_obj.get("hosting_eula_key")
 
-        self.model_package_arn: Optional[str] = json_obj.get("model_package_arn")
+        self.hosting_model_package_arn: Optional[str] = json_obj.get("hosting_model_package_arn")
 
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(

src/sagemaker/jumpstart/utils.py

@@ -402,13 +402,13 @@ def verify_model_region_and_return_specs(
             f"JumpStart model ID '{model_id}' and version '{version}' " "does not support training."
         )
 
-    if model_specs.eula_key:
+    if model_specs.hosting_eula_key and scope == constants.JumpStartScriptScope.INFERENCE.value:
         LOGGER.info(
             "Using model with end-user license agreement (EULA). "
             "See https://%s.s3.%s.amazonaws.com/%s for terms of use.",
             get_jumpstart_content_bucket(region=region),
             region,
-            model_specs.eula_key,
+            model_specs.hosting_eula_key,
         )
 
     if model_specs.deprecated:

tests/unit/sagemaker/jumpstart/constants.py

@@ -2299,8 +2299,8 @@
     "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
     "training_prepacked_script_key": None,
     "hosting_prepacked_artifact_key": None,
-    "model_package_arn": None,
-    "eula_key": False,
+    "hosting_model_package_arn": None,
+    "hosting_eula_key": False,
     "hyperparameters": [
         {
             "name": "epochs",

tests/unit/test_predictor.py

@@ -614,3 +614,22 @@ def test_setting_serializer_deserializer_atts_changes_content_accept_types():
     predictor.deserializer = PandasDeserializer()
     assert predictor.accept == ("text/csv", "application/json")
     assert predictor.content_type == "text/csv"
+
+
+def test_custom_attributes():
+    sagemaker_session = empty_sagemaker_session()
+    predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
+
+    sagemaker_session.sagemaker_runtime_client.invoke_endpoint = Mock(
+        return_value={"Body": io.StringIO("response")}
+    )
+
+    predictor.predict("payload", custom_attributes="custom-attribute")
+
+    sagemaker_session.sagemaker_runtime_client.invoke_endpoint.assert_called_once_with(
+        EndpointName=ENDPOINT,
+        ContentType="application/octet-stream",
+        Accept="*/*",
+        CustomAttributes="custom-attribute",
+        Body="payload",
+    )