feat: jumpstart EULA models (initial commit)

Author: evakravi

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart model uris."""
 from __future__ import absolute_import
+import logging
 import os
 from typing import Optional
 
@@ -26,6 +27,23 @@
     get_jumpstart_content_bucket,
     verify_model_region_and_return_specs,
 )
+from sagemaker.s3_utils import parse_s3_url
+
+logger = logging.getLogger(__name__)
+
+
+_HOSTING_MODEL_ARTIFACT_URIS_PRIORITY_LIST = [
+    "hosting_prepacked_artifact_uri",
+    "hosting_artifact_uri",
+]
+_HOSTING_MODEL_ARTIFACT_KEYS_PRIORITY_LIST = [
+    "hosting_prepacked_artifact_key",
+    "hosting_artifact_key",
+]
+_MODEL_ARTIFACT_PRIORITY_LIST = [
+    *_HOSTING_MODEL_ARTIFACT_URIS_PRIORITY_LIST,
+    *_HOSTING_MODEL_ARTIFACT_KEYS_PRIORITY_LIST,
+]
 
 
 def _retrieve_model_uri(
@@ -76,18 +94,35 @@ def _retrieve_model_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
+    bucket_override = os.environ.get(ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE)
+
+    if bucket_override:
+        logger.info("Using JumpStart bucket override for model URI: '%s'", bucket_override)
+
     if model_scope == JumpStartScriptScope.INFERENCE:
-        model_artifact_key = (
-            getattr(model_specs, "hosting_prepacked_artifact_key", None)
-            or model_specs.hosting_artifact_key
-        )
+
+        model_artifact_key = None
+
+        for model_artifact_field_name in _MODEL_ARTIFACT_PRIORITY_LIST:
+            field_value = getattr(model_specs, model_artifact_field_name, None)
+            if field_value is not None:
+                if model_artifact_field_name in _HOSTING_MODEL_ARTIFACT_URIS_PRIORITY_LIST:
+                    bucket, key = parse_s3_url(field_value)
+                    bucket_to_use = bucket_override or bucket
+                    field_value_to_use = f"s3://{bucket_to_use}/{key}"
+
+                    return field_value_to_use
+
+                model_artifact_key = field_value
+                break
+
+        if model_artifact_key is None:
+            raise RuntimeError(f"Unable to find model artifact for '{model_id}'.")
 
     elif model_scope == JumpStartScriptScope.TRAINING:
         model_artifact_key = model_specs.training_artifact_key
 
-    bucket = os.environ.get(
-        ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE
-    ) or get_jumpstart_content_bucket(region)
+    bucket = bucket_override or get_jumpstart_content_bucket(region)
 
     model_s3_uri = f"s3://{bucket}/{model_artifact_key}"
 

src/sagemaker/jumpstart/artifacts/script_uris.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart script uris."""
 from __future__ import absolute_import
+import logging
 import os
 from typing import Optional
 from sagemaker.jumpstart.constants import (
@@ -25,6 +26,9 @@
     get_jumpstart_content_bucket,
     verify_model_region_and_return_specs,
 )
+from sagemaker.s3_utils import parse_s3_url
+
+logger = logging.getLogger(__name__)
 
 
 def _retrieve_script_uri(
@@ -76,16 +80,27 @@ def _retrieve_script_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
+    bucket_override = os.environ.get(ENV_VARIABLE_JUMPSTART_SCRIPT_ARTIFACT_BUCKET_OVERRIDE)
+
+    if bucket_override:
+        logger.info("Using JumpStart bucket override for script URI: '%s'", bucket_override)
+
     if script_scope == JumpStartScriptScope.INFERENCE:
+        script_uri = getattr(model_specs, "hosting_script_uri", None)
+        if script_uri is not None:
+            bucket, key = parse_s3_url(script_uri)
+            bucket_to_use = bucket_override or bucket
+            script_uri_to_use = f"s3://{bucket_to_use}/{key}"
+
+            return script_uri_to_use
+
         model_script_key = model_specs.hosting_script_key
     elif script_scope == JumpStartScriptScope.TRAINING:
         model_script_key = (
             getattr(model_specs, "training_prepacked_script_key") or model_specs.training_script_key
         )
 
-    bucket = os.environ.get(
-        ENV_VARIABLE_JUMPSTART_SCRIPT_ARTIFACT_BUCKET_OVERRIDE
-    ) or get_jumpstart_content_bucket(region)
+    bucket = bucket_override or get_jumpstart_content_bucket(region)
 
     script_s3_uri = f"s3://{bucket}/{model_script_key}"
 

src/sagemaker/jumpstart/types.py

@@ -351,6 +351,10 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "inference_enable_network_isolation",
         "training_enable_network_isolation",
         "resource_name_base",
+        "eula_model",
+        "hosting_prepacked_artifact_uri",
+        "hosting_artifact_uri",
+        "hosting_script_uri",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -419,6 +423,13 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         )
         self.resource_name_base: bool = json_obj.get("resource_name_base")
 
+        self.eula_model: bool = json_obj.get("eula_model", False)
+        self.hosting_prepacked_artifact_uri: Optional[str] = json_obj.get(
+            "hosting_prepacked_artifact_uri"
+        )
+        self.hosting_artifact_uri: Optional[str] = json_obj.get("hosting_artifact_uri")
+        self.hosting_script_uri: Optional[str] = json_obj.get("hosting_script_uri")
+
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(
                 json_obj["training_ecr_specs"]

src/sagemaker/jumpstart/utils.py

@@ -402,6 +402,12 @@ def verify_model_region_and_return_specs(
             f"JumpStart model ID '{model_id}' and version '{version}' " "does not support training."
         )
 
+    if model_specs.eula_model:
+        LOGGER.info(
+            "Using model with end-user license agreement (EULA). "
+            "Deploying this model requires accepting EULA terms."
+        )
+
     if model_specs.deprecated:
         if not tolerate_deprecated_model:
             raise DeprecatedJumpStartModelError(model_id=model_id, version=version)

tests/unit/sagemaker/jumpstart/constants.py

@@ -2294,11 +2294,15 @@
         "py_version": "py3",
     },
     "hosting_artifact_key": "pytorch-infer/infer-pytorch-ic-mobilenet-v2.tar.gz",
+    "hosting_artifact_uri": None,
     "training_artifact_key": "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz",
     "hosting_script_key": "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
+    "hosting_script_uri": None,
     "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,
+    "hosting_prepacked_artifact_uri": None,
+    "eula_model": False,
     "hyperparameters": [
         {
             "name": "epochs",