update filter name

Author: Captainia

src/sagemaker/jumpstart/constants.py

@@ -190,6 +190,9 @@
 
 SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY = "SageMakerGatedModelS3Uri"
 
+PROPRIETARY_MODEL_SPEC_PREFIX = "proprietary-models"
+PROPRIETARY_MODEL_FILTER_NAME = "marketplace"
+
 CONTENT_TYPE_TO_SERIALIZER_TYPE_MAP: Dict[MIMEType, SerializerType] = {
     MIMEType.X_IMAGE: SerializerType.RAW_BYTES,
     MIMEType.LIST_TEXT: SerializerType.JSON,

src/sagemaker/jumpstart/notebook_utils.py

@@ -24,6 +24,8 @@
 from sagemaker.jumpstart.constants import (
     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 (
@@ -128,7 +130,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]
@@ -144,10 +146,10 @@ def extract_model_type(spec_key: str) -> str:
     Args:
         spek_key (str): The model spec key for which to extract the model type.
     """
-    model_type = spec_key.split("/")[0]
+    model_spec_prefix = spec_key.split("/")[0]
 
-    if model_type == "proprietary-models":
-        return JumpStartModelType.PROPRIETARY.value
+    if model_spec_prefix == PROPRIETARY_MODEL_SPEC_PREFIX:
+        return PROPRIETARY_MODEL_FILTER_NAME
 
     return JumpStartModelType.OPEN_WEIGHT.value
 

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -21,7 +21,6 @@
 from sagemaker.predictor import retrieve_default
 
 import tests.integ
-from sagemaker.jumpstart import notebook_utils
 
 from sagemaker.jumpstart.model import JumpStartModel
 from tests.integ.sagemaker.jumpstart.constants import (