rename to proprietary model and fix unittests

Author: Captainia

doc/doc_utils/jumpstart_doc_utils.py

@@ -206,10 +206,10 @@ def get_model_source(url):
         return "Source"
 
 
-def create_marketplace_model_table():
+def create_proprietary_model_table():
     marketpkace_content_intro = []
     marketpkace_content_intro.append("\n")
-    marketpkace_content_intro.append(".. list-table:: Available Models\n")
+    marketpkace_content_intro.append(".. list-table:: Available Proprietary Models\n")
     marketpkace_content_intro.append("   :widths: 50 20 20 20 20\n")
     marketpkace_content_intro.append("   :header-rows: 1\n")
     marketpkace_content_intro.append("   :class: datatable\n")
@@ -232,17 +232,17 @@ def create_marketplace_model_table():
             ):
                 sdk_manifest_top_versions_for_models[model["model_id"]] = model
 
-    marketplace_content_entries = []
+    proprietary_content_entries = []
     for model in sdk_manifest_top_versions_for_models.values():
         model_spec = get_jumpstart_sdk_spec(model["spec_key"])
-        marketplace_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
-        marketplace_content_entries.append("     - {}\n".format(False))  # TODO: support training
-        marketplace_content_entries.append("     - {}\n".format(model["version"]))
-        marketplace_content_entries.append("     - {}\n".format(model["min_version"]))
-        marketplace_content_entries.append(
+        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 marketpkace_content_intro + marketplace_content_entries + ["\n"]
+    return marketpkace_content_intro + proprietary_content_entries + ["\n"]
 
 
 def create_jumpstart_model_table():
@@ -348,10 +348,10 @@ def create_jumpstart_model_table():
             f.writelines(file_content_single_entry)
             f.close()
 
-    marketplace_content_entries = create_marketplace_model_table()
+    proprietary_content_entries = create_proprietary_model_table()
 
     f = open("doc_utils/pretrainedmodels.rst", "a")
     f.writelines(file_content_intro)
     f.writelines(open_weight_content_entries)
-    f.writelines(marketplace_content_entries)
+    f.writelines(proprietary_content_entries)
     f.close()

src/sagemaker/jumpstart/enums.py

@@ -91,8 +91,6 @@ class JumpStartTag(str, Enum):
     MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
     MODEL_TYPE = "sagemaker-sdk:jumpstart-model-type"
 
-    MARKETPLACE_MODEL_TYPE_VALUE = "SageMakerJumpStartMarketplace"
-
 
 class SerializerType(str, Enum):
     """Enum class for serializers associated with JumpStart models."""

src/sagemaker/jumpstart/exceptions.py

@@ -62,10 +62,10 @@ def get_proprietary_model_subscription_msg(
     model_id: str,
     subscription_link: str,
 ) -> str:
-    """Returns customer-facing message for using a Marketplace model."""
+    """Returns customer-facing message for using a proprietary model."""
 
     return (
-        f"INFO: Using Marketplace model '{model_id}'. "
+        f"INFO: Using proprietary model '{model_id}'. "
         f"Please make sure to subscribe to the model from {subscription_link}"
     )
 
@@ -75,7 +75,7 @@ def get_wildcard_proprietary_model_version_msg(
 ) -> str:
     """Returns customer-facing message for passing wildcard version to proprietary models."""
     msg = (
-        f"Marketplace model '{model_id}' does not support "
+        f"Proprietary model '{model_id}' does not support "
         f"wildcard version identifier '{wildcard_model_version}'. "
     )
     if len(available_versions) > 0:

src/sagemaker/jumpstart/factory/model.py

@@ -169,7 +169,7 @@ def _log_model_type(kwargs: JumpStartModelInitKwargs) -> None:
     """Log the model type being used"""
     if kwargs.model_type == JumpStartModelType.PROPRIETARY:
         JUMPSTART_LOGGER.info(
-            "Marketplace model %s of version %s is being used.",
+            "Proprietary model %s of version %s is being used.",
             kwargs.model_id,
             kwargs.model_version,
         )

src/sagemaker/jumpstart/filters.py

@@ -430,6 +430,22 @@ def __init__(self, key: str, value: str, operator: str):
         self.value = value
         self.operator = operator
 
+    def set_key(self, key: str) -> None:
+        """Sets the key for the model filter.
+
+        Args:
+            key (str): The key to be set.
+        """
+        self.key = key
+
+    def set_value(self, value: str) -> None:
+        """Sets the value for the model filter.
+
+        Args:
+            value (str): The value to be set.
+        """
+        self.value = value
+
 
 def parse_filter_string(filter_string: str) -> ModelFilter:
     """Parse filter string and return a serialized ``ModelFilter`` object.

src/sagemaker/jumpstart/notebook_utils.py

@@ -25,7 +25,6 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     JUMPSTART_DEFAULT_REGION_NAME,
     PROPRIETARY_MODEL_SPEC_PREFIX,
-    PROPRIETARY_MODEL_FILTER_NAME,
 )
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
 from sagemaker.jumpstart.filters import (
@@ -130,7 +129,7 @@ def extract_framework_task_model(model_id: str) -> Tuple[str, str, str]:
     """
     _id_parts = model_id.split("-")
 
-    if len(_id_parts) != 3:
+    if len(_id_parts) < 3:
         return "", "", ""
 
     framework = _id_parts[0]
@@ -149,7 +148,7 @@ def extract_model_type(spec_key: str) -> str:
     model_spec_prefix = spec_key.split("/")[0]
 
     if model_spec_prefix == PROPRIETARY_MODEL_SPEC_PREFIX:
-        return PROPRIETARY_MODEL_FILTER_NAME
+        return JumpStartModelType.PROPRIETARY.value
 
     return JumpStartModelType.OPEN_WEIGHT.value
 
@@ -228,6 +227,7 @@ def list_jumpstart_scripts(  # pylint: disable=redefined-builtin
         sagemaker_session (sagemaker.session.Session): Optional. The SageMaker Session to
             use to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
+
     if (isinstance(filter, Constant) and filter.resolved_value == BooleanValues.TRUE) or (
         isinstance(filter, str) and filter.lower() == BooleanValues.TRUE.lower()
     ):
@@ -279,8 +279,6 @@ def list_jumpstart_models(  # pylint: disable=redefined-builtin
             versions should be included in the returned result. (Default: False).
         list_versions (bool): Optional. True if versions for models should be returned in addition
             to the id of the model. (Default: False).
-        marketplace_models (bool): Optional. True if only listing JumpStart Marketplace models.
-            (Default: False).
         sagemaker_session (sagemaker.session.Session): Optional. The SageMaker Session to use
             to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
@@ -361,6 +359,11 @@ def _generate_jumpstart_model_versions(  # pylint: disable=redefined-builtin
         model_filter = operator.unresolved_value
         key = model_filter.key
         all_keys.add(key)
+        if model_filter.key == SpecialSupportedFilterKeys.MODEL_TYPE and model_filter.value in [
+            "marketplace",
+            "proprietary",
+        ]:
+            model_filter.set_value(JumpStartModelType.PROPRIETARY.value)
         model_filters.add(model_filter)
 
     for key in all_keys:

src/sagemaker/jumpstart/utils.py

@@ -369,7 +369,7 @@ def add_jumpstart_model_id_version_tags(
     )
     if model_type == enums.JumpStartModelType.PROPRIETARY:
         tags = add_single_jumpstart_tag(
-            enums.JumpStartTag.MARKETPLACE_MODEL_TYPE_VALUE,
+            enums.JumpStartModelType.PROPRIETARY.value,
             enums.JumpStartTag.MODEL_TYPE,
             tags,
             is_uri=False,

src/sagemaker/payloads.py

@@ -137,8 +137,8 @@ def retrieve_example(
             the model payload.
         model_version (str): The version of the JumpStart model for which to retrieve
             the model payload.
-        model_type (str): The model type of the JumpStart model, either is open source
-            or marketplace (proprietary).
+        model_type (str): The model type of the JumpStart model, either is open weight
+            or proprietary.
         serialize (bool): Whether to serialize byte-stream valued payloads by downloading
             binary files from s3 and applying encoding, or to keep payload in pre-serialized
             state. Set this option to False if you want to avoid s3 downloads or if you

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -492,7 +492,7 @@ def test_proprietary_model_endpoint(
             tags=[
                 {"Key": JumpStartTag.MODEL_ID, "Value": "ai21-summarization"},
                 {"Key": JumpStartTag.MODEL_VERSION, "Value": "2.0.004"},
-                {"Key": JumpStartTag.MODEL_TYPE, "Value": "SageMakerJumpStartMarketplace"},
+                {"Key": JumpStartTag.MODEL_TYPE, "Value": "proprietary"},
             ],
             endpoint_logging=False,
             model_data_download_timeout=3600,

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -181,7 +181,7 @@ def test_jumpstart_cache_get_header():
             model_type=JumpStartModelType.PROPRIETARY,
         )
     assert (
-        "Marketplace model 'ai21-summarization' does not support wildcard version identifier '3.*'. "
+        "Proprietary model 'ai21-summarization' does not support wildcard version identifier '3.*'. "
         "You can pin to version '1.1.003'. "
         "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html "
         "for list of supported model IDs. " in str(e.value)
@@ -941,7 +941,7 @@ def test_jumpstart_cache_get_specs():
             model_type=JumpStartModelType.PROPRIETARY,
         )
     assert (
-        "Marketplace model 'ai21-summarization' does not support wildcard version identifier '3.*'. "
+        "Proprietary model 'ai21-summarization' does not support wildcard version identifier '3.*'. "
         "You can pin to version '1.1.003'. "
         "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html "
         "for list of supported model IDs. " in str(e.value)

tests/unit/sagemaker/jumpstart/test_notebook_utils.py

@@ -640,6 +640,7 @@ def test_list_jumpstart_proprietary_models(
         ]
 
         assert list_jumpstart_models("model_type == proprietary") == all_prop_model_ids
+        assert list_jumpstart_models("model_type == marketplace") == all_prop_model_ids
         assert list_jumpstart_models("model_type == open_weight") == all_open_weight_model_ids
 
         assert list_jumpstart_models(list_versions=False) == sorted(