feat: support JumpStart proprietary models (aws#4467)

* feat: add proprietary manifest/specs parsing

add unittests for test_cache

small refactoring

address comments and more unittests

fix linting and fix more tests

fix: pylint

feat: JumpStartModel class for prop models

* remove unused imports and fix docstyle

* fix: remove unused args

* fix: remove unused args

* fix: more unused vars

* fix: slow tests

* fix: unittests

* added more tests to cover some lines

* remove estimator warn check

* chore: address comments re performance

* fix: address comments

* complete list experience and other fixes

* fix: pylint

* add doc utils and fix pylint

* fix: docstyle

* fix: doc

* fix: default payloads

* fix: doc and tags and enums

* fix: jumpstart doc

* rename to open_weights and fix filtering

* update filter name

* doc update

* fix: black

* rename to proprietary model and fix unittests

* address comments

* fix: docstyle and flake8

* address more comments and fix doc

* put back doc utils for future refactoring

* add prop model title in doc

* doc update

---------

Co-authored-by: liujiaor <[email protected]>

Author: 2 people

doc/doc_utils/jumpstart_doc_utils.py

@@ -74,9 +74,12 @@ class Frameworks(str, Enum):
 
 JUMPSTART_REGION = "eu-west-2"
 SDK_MANIFEST_FILE = "models_manifest.json"
+PROPRIETARY_SDK_MANIFEST_FILE = "proprietary-sdk-manifest.json"
 JUMPSTART_BUCKET_BASE_URL = "https://jumpstart-cache-prod-{}.s3.{}.amazonaws.com".format(
     JUMPSTART_REGION, JUMPSTART_REGION
 )
+PROPRIETARY_DOC_BUCKET = "https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com"
+
 TASK_MAP = {
     Tasks.IC: ProblemTypes.IMAGE_CLASSIFICATION,
     Tasks.IC_EMBEDDING: ProblemTypes.IMAGE_EMBEDDING,
@@ -152,18 +155,26 @@ class Frameworks(str, Enum):
 }
 
 
-def get_jumpstart_sdk_manifest():
-    url = "{}/{}".format(JUMPSTART_BUCKET_BASE_URL, SDK_MANIFEST_FILE)
+def get_public_s3_json_object(url):
     with request.urlopen(url) as f:
         models_manifest = f.read().decode("utf-8")
     return json.loads(models_manifest)
 
 
-def get_jumpstart_sdk_spec(key):
-    url = "{}/{}".format(JUMPSTART_BUCKET_BASE_URL, key)
-    with request.urlopen(url) as f:
-        model_spec = f.read().decode("utf-8")
-    return json.loads(model_spec)
+def get_jumpstart_sdk_manifest():
+    return get_public_s3_json_object(f"{JUMPSTART_BUCKET_BASE_URL}/{SDK_MANIFEST_FILE}")
+
+
+def get_proprietary_sdk_manifest():
+    return get_public_s3_json_object(f"{PROPRIETARY_DOC_BUCKET}/{PROPRIETARY_SDK_MANIFEST_FILE}")
+
+
+def get_jumpstart_sdk_spec(s3_key: str):
+    return get_public_s3_json_object(f"{JUMPSTART_BUCKET_BASE_URL}/{s3_key}")
+
+
+def get_proprietary_sdk_spec(s3_key: str):
+    return get_public_s3_json_object(f"{PROPRIETARY_DOC_BUCKET}/{s3_key}")
 
 
 def get_model_task(id):
@@ -196,6 +207,45 @@ def get_model_source(url):
         return "Source"
 
 
+def create_proprietary_model_table():
+    proprietary_content_intro = []
+    proprietary_content_intro.append("\n")
+    proprietary_content_intro.append(".. list-table:: Available Proprietary Models\n")
+    proprietary_content_intro.append("   :widths: 50 20 20 20 20\n")
+    proprietary_content_intro.append("   :header-rows: 1\n")
+    proprietary_content_intro.append("   :class: datatable\n")
+    proprietary_content_intro.append("\n")
+    proprietary_content_intro.append("   * - Model ID\n")
+    proprietary_content_intro.append("     - Fine Tunable?\n")
+    proprietary_content_intro.append("     - Supported Version\n")
+    proprietary_content_intro.append("     - Min SDK Version\n")
+    proprietary_content_intro.append("     - Source\n")
+
+    sdk_manifest = get_proprietary_sdk_manifest()
+    sdk_manifest_top_versions_for_models = {}
+
+    for model in sdk_manifest:
+        if model["model_id"] not in sdk_manifest_top_versions_for_models:
+            sdk_manifest_top_versions_for_models[model["model_id"]] = model
+        else:
+            if str(sdk_manifest_top_versions_for_models[model["model_id"]]["version"]) < str(
+                model["version"]
+            ):
+                sdk_manifest_top_versions_for_models[model["model_id"]] = model
+
+    proprietary_content_entries = []
+    for model in sdk_manifest_top_versions_for_models.values():
+        model_spec = get_proprietary_sdk_spec(model["spec_key"])
+        proprietary_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
+        proprietary_content_entries.append("     - {}\n".format(False))  # TODO: support training
+        proprietary_content_entries.append("     - {}\n".format(model["version"]))
+        proprietary_content_entries.append("     - {}\n".format(model["min_version"]))
+        proprietary_content_entries.append(
+            "     - `{} <{}>`__ |external-link|\n".format("Source", model_spec.get("url"))
+        )
+    return proprietary_content_intro + proprietary_content_entries + ["\n"]
+
+
 def create_jumpstart_model_table():
     sdk_manifest = get_jumpstart_sdk_manifest()
     sdk_manifest_top_versions_for_models = {}
@@ -249,19 +299,19 @@ def create_jumpstart_model_table():
     file_content_intro.append("     - Source\n")
 
     dynamic_table_files = []
-    file_content_entries = []
+    open_weight_content_entries = []
 
     for model in sdk_manifest_top_versions_for_models.values():
         model_spec = get_jumpstart_sdk_spec(model["spec_key"])
         model_task = get_model_task(model_spec["model_id"])
         string_model_task = get_string_model_task(model_spec["model_id"])
         model_source = get_model_source(model_spec["url"])
-        file_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
-        file_content_entries.append("     - {}\n".format(model_spec["training_supported"]))
-        file_content_entries.append("     - {}\n".format(model["version"]))
-        file_content_entries.append("     - {}\n".format(model["min_version"]))
-        file_content_entries.append("     - {}\n".format(model_task))
-        file_content_entries.append(
+        open_weight_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
+        open_weight_content_entries.append("     - {}\n".format(model_spec["training_supported"]))
+        open_weight_content_entries.append("     - {}\n".format(model["version"]))
+        open_weight_content_entries.append("     - {}\n".format(model["min_version"]))
+        open_weight_content_entries.append("     - {}\n".format(model_task))
+        open_weight_content_entries.append(
             "     - `{} <{}>`__ |external-link|\n".format(model_source, model_spec["url"])
         )
 
@@ -299,7 +349,10 @@ def create_jumpstart_model_table():
             f.writelines(file_content_single_entry)
             f.close()
 
+    proprietary_content_entries = create_proprietary_model_table()
+
     f = open("doc_utils/pretrainedmodels.rst", "a")
     f.writelines(file_content_intro)
-    f.writelines(file_content_entries)
+    f.writelines(open_weight_content_entries)
+    f.writelines(proprietary_content_entries)
     f.close()

src/sagemaker/accept_types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -75,6 +76,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default accept type for the model matching the given arguments.
 
@@ -114,4 +116,5 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )

src/sagemaker/base_predictor.py

@@ -58,7 +58,9 @@
 from sagemaker.model_monitor.model_monitoring import DEFAULT_REPOSITORY_NAME
 
 from sagemaker.lineage.context import EndpointContext
-from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
+from sagemaker.compute_resource_requirements.resource_requirements import (
+    ResourceRequirements,
+)
 
 LOGGER = logging.getLogger("sagemaker")
 

src/sagemaker/content_types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -75,6 +76,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default content type for the model matching the given arguments.
 
@@ -114,6 +116,7 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )
 
 

src/sagemaker/deserializers.py

@@ -35,6 +35,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -95,6 +96,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> BaseDeserializer:
     """Retrieves the default deserializer for the model matching the given arguments.
 
@@ -135,4 +137,5 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )

src/sagemaker/instance_types.py

@@ -20,6 +20,7 @@
 from sagemaker.jumpstart import utils as jumpstart_utils
 from sagemaker.jumpstart import artifacts
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 logger = logging.getLogger(__name__)
@@ -34,6 +35,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     training_instance_type: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default instance type for the model matching the given arguments.
 
@@ -85,6 +87,7 @@ def retrieve_default(
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         training_instance_type=training_instance_type,
+        model_type=model_type,
     )
 
 

src/sagemaker/jumpstart/accessors.py

@@ -18,6 +18,7 @@
 
 from sagemaker.deprecations import deprecated
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.jumpstart import cache
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
@@ -197,7 +198,9 @@ def _set_cache_and_region(region: str, cache_kwargs: dict) -> None:
 
     @staticmethod
     def _get_manifest(
-        region: str = JUMPSTART_DEFAULT_REGION_NAME, s3_client: Optional[boto3.client] = None
+        region: str = JUMPSTART_DEFAULT_REGION_NAME,
+        s3_client: Optional[boto3.client] = None,
+        model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
     ) -> List[JumpStartModelHeader]:
         """Return entire JumpStart models manifest.
 
@@ -215,13 +218,19 @@ def _get_manifest(
             additional_kwargs.update({"s3_client": s3_client})
 
         cache_kwargs = JumpStartModelsAccessor._validate_and_mutate_region_cache_kwargs(
-            {**JumpStartModelsAccessor._cache_kwargs, **additional_kwargs}, region
+            {**JumpStartModelsAccessor._cache_kwargs, **additional_kwargs},
+            region,
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
-        return JumpStartModelsAccessor._cache.get_manifest()  # type: ignore
+        return JumpStartModelsAccessor._cache.get_manifest(model_type)  # type: ignore
 
     @staticmethod
-    def get_model_header(region: str, model_id: str, version: str) -> JumpStartModelHeader:
+    def get_model_header(
+        region: str,
+        model_id: str,
+        version: str,
+        model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    ) -> JumpStartModelHeader:
         """Returns model header from JumpStart models cache.
 
         Args:
@@ -234,12 +243,18 @@ def get_model_header(region: str, model_id: str, version: str) -> JumpStartModel
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
         return JumpStartModelsAccessor._cache.get_header(  # type: ignore
-            model_id=model_id, semantic_version_str=version
+            model_id=model_id,
+            semantic_version_str=version,
+            model_type=model_type,
         )
 
     @staticmethod
     def get_model_specs(
-        region: str, model_id: str, version: str, s3_client: Optional[boto3.client] = None
+        region: str,
+        model_id: str,
+        version: str,
+        s3_client: Optional[boto3.client] = None,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     ) -> JumpStartModelSpecs:
         """Returns model specs from JumpStart models cache.
 
@@ -260,7 +275,7 @@ def get_model_specs(
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
-            model_id=model_id, semantic_version_str=version
+            model_id=model_id, version_str=version, model_type=model_type
         )
 
     @staticmethod

src/sagemaker/jumpstart/artifacts/instance_types.py

@@ -22,6 +22,7 @@
 )
 from sagemaker.jumpstart.enums import (
     JumpStartScriptScope,
+    JumpStartModelType,
 )
 from sagemaker.jumpstart.utils import (
     verify_model_region_and_return_specs,
@@ -38,6 +39,7 @@ def _retrieve_default_instance_type(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     training_instance_type: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default instance type for the model.
 
@@ -84,6 +86,7 @@ def _retrieve_default_instance_type(
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
+        model_type=model_type,
         sagemaker_session=sagemaker_session,
     )
 

src/sagemaker/jumpstart/artifacts/kwargs.py

@@ -22,6 +22,7 @@
 )
 from sagemaker.jumpstart.enums import (
     JumpStartScriptScope,
+    JumpStartModelType,
 )
 from sagemaker.jumpstart.utils import (
     verify_model_region_and_return_specs,
@@ -35,6 +36,7 @@ def _retrieve_model_init_kwargs(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> dict:
     """Retrieves kwargs for `Model`.
 
@@ -71,6 +73,7 @@ def _retrieve_model_init_kwargs(
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )
 
     kwargs = deepcopy(model_specs.model_kwargs)
@@ -89,6 +92,7 @@ def _retrieve_model_deploy_kwargs(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> dict:
     """Retrieves kwargs for `Model.deploy`.
 
@@ -128,6 +132,7 @@ def _retrieve_model_deploy_kwargs(
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )
 
     if volume_size_supported(instance_type) and model_specs.inference_volume_size is not None:

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -22,6 +22,7 @@
 )
 from sagemaker.jumpstart.enums import (
     JumpStartScriptScope,
+    JumpStartModelType,
 )
 from sagemaker.session import Session
 
@@ -35,6 +36,7 @@ def _retrieve_model_package_arn(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Optional[str]:
     """Retrieves associated model pacakge arn for the model.
 
@@ -74,6 +76,7 @@ def _retrieve_model_package_arn(
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )
 
     if scope == JumpStartScriptScope.INFERENCE: