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

src/sagemaker/jumpstart/cache.py

@@ -411,7 +411,6 @@ def _retrieval_function(
         """
 
         data_type, id_info = key.data_type, key.id_info
-
         if data_type in {
             JumpStartS3FileType.OPEN_WEIGHT_MANIFEST,
             JumpStartS3FileType.PROPRIETARY_MANIFEST,
@@ -425,7 +424,6 @@ def _retrieval_function(
                 formatted_content=utils.get_formatted_manifest(formatted_body),
                 md5_hash=etag,
             )
-
         if data_type in {
             JumpStartS3FileType.OPEN_WEIGHT_SPECS,
             JumpStartS3FileType.PROPRIETARY_SPECS,

src/sagemaker/jumpstart/types.py

@@ -2368,6 +2368,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "model_version",
         "model_type",
         "hub_arn",
+        "model_type",
         "region",
         "tolerate_deprecated_model",
         "tolerate_vulnerable_model",

tests/unit/sagemaker/instance_types/jumpstart/test_instance_types.py

@@ -126,6 +126,7 @@ def test_jumpstart_instance_types(patched_get_model_specs, patched_validate_mode
         model_type=JumpStartModelType.OPEN_WEIGHTS,
         hub_arn=None,
         s3_client=mock_client,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     )
 
     patched_get_model_specs.reset_mock()

tests/unit/sagemaker/jumpstart/utils.py

@@ -31,7 +31,6 @@
     HubContentType,
 )
 from sagemaker.jumpstart.enums import JumpStartModelType
-
 from sagemaker.jumpstart.utils import get_formatted_manifest
 from tests.unit.sagemaker.jumpstart.constants import (
     PROTOTYPICAL_MODEL_SPECS_DICT,

tests/unit/sagemaker/resource_requirements/jumpstart/test_resource_requirements.py

@@ -110,8 +110,12 @@ def test_jumpstart_resource_requirements_instance_type_variants(patched_get_mode
     }
 
 
+@patch("sagemaker.jumpstart.utils.validate_model_id_and_get_type")
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_no_supported_resource_requirements(patched_get_model_specs):
+def test_jumpstart_no_supported_resource_requirements(
+    patched_get_model_specs, patched_validate_model_id_and_get_type
+):
+
     patched_get_model_specs.side_effect = get_special_model_spec
     region = "us-west-2"
     mock_client = boto3.client("s3")

tests/unit/sagemaker/script_uris/jumpstart/test_common.py

@@ -54,6 +54,7 @@ def test_jumpstart_common_script_uri(
         s3_client=mock_client,
         model_type=JumpStartModelType.OPEN_WEIGHTS,
         hub_arn=None,
+        model_type=JumpStartModelType.OPEN_WEIGHTS,
     )
     patched_verify_model_region_and_return_specs.assert_called_once()