chore: cleanup jumpstart factory (#4840)

* chore: cleanup jumpstart factory

* fix: typing

* chore: address pr comments, fix formatting

* fix: failing config tests

Author: evakravi

src/sagemaker/environment_variables.py

@@ -20,7 +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 JumpStartScriptScope
+from sagemaker.jumpstart.enums import JumpStartModelType, JumpStartScriptScope
 from sagemaker.session import Session
 
 logger = logging.getLogger(__name__)
@@ -38,6 +38,7 @@ def retrieve_default(
     instance_type: Optional[str] = None,
     script: JumpStartScriptScope = JumpStartScriptScope.INFERENCE,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Dict[str, str]:
     """Retrieves the default container environment variables for the model matching the arguments.
 
@@ -70,6 +71,8 @@ def retrieve_default(
         script (JumpStartScriptScope): The JumpStart script for which to retrieve environment
             variables.
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         dict: The variables to use for the model.
 
@@ -94,4 +97,5 @@ def retrieve_default(
         instance_type=instance_type,
         script=script,
         config_name=config_name,
+        model_type=model_type,
     )

src/sagemaker/hyperparameters.py

@@ -20,7 +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 HyperparameterValidationMode
+from sagemaker.jumpstart.enums import HyperparameterValidationMode, JumpStartModelType
 from sagemaker.jumpstart.validators import validate_hyperparameters
 from sagemaker.session import Session
 
@@ -38,6 +38,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Dict[str, str]:
     """Retrieves the default training hyperparameters for the model matching the given arguments.
 
@@ -71,6 +72,8 @@ def retrieve_default(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         dict: The hyperparameters to use for the model.
 
@@ -93,6 +96,7 @@ def retrieve_default(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
 

src/sagemaker/image_uris.py

@@ -22,6 +22,7 @@
 
 from sagemaker import utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.jumpstart.utils import is_jumpstart_model_input
 from sagemaker.spark import defaults
 from sagemaker.jumpstart import artifacts
@@ -72,6 +73,7 @@ def retrieve(
     serverless_inference_config=None,
     sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name=None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the ECR URI for the Docker image matching the given arguments.
 
@@ -128,6 +130,8 @@ def retrieve(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
 
     Returns:
         str: The ECR URI for the corresponding SageMaker Docker image.
@@ -169,6 +173,7 @@ def retrieve(
             tolerate_deprecated_model,
             sagemaker_session=sagemaker_session,
             config_name=config_name,
+            model_type=model_type,
         )
 
     if training_compiler_config and (framework in [HUGGING_FACE_FRAMEWORK, "pytorch"]):

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -19,6 +19,7 @@
     SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
+    JumpStartModelType,
     JumpStartScriptScope,
 )
 from sagemaker.jumpstart.utils import (
@@ -41,6 +42,7 @@ def _retrieve_default_environment_variables(
     instance_type: Optional[str] = None,
     script: JumpStartScriptScope = JumpStartScriptScope.INFERENCE,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Dict[str, str]:
     """Retrieves the inference environment variables for the model matching the given arguments.
 
@@ -73,6 +75,8 @@ def _retrieve_default_environment_variables(
         script (JumpStartScriptScope): The JumpStart script for which to retrieve
             environment variables.
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         dict: the inference environment variables to use for the model.
     """
@@ -91,6 +95,7 @@ def _retrieve_default_environment_variables(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     default_environment_variables: Dict[str, str] = {}
@@ -130,6 +135,7 @@ def _retrieve_default_environment_variables(
                     sagemaker_session=sagemaker_session,
                     instance_type=instance_type,
                     config_name=config_name,
+                    model_type=model_type,
                 )
             )
 
@@ -178,6 +184,7 @@ def _retrieve_gated_model_uri_env_var_value(
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     instance_type: Optional[str] = None,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Optional[str]:
     """Retrieves the gated model env var URI matching the given arguments.
 
@@ -204,7 +211,8 @@ def _retrieve_gated_model_uri_env_var_value(
         instance_type (str): An instance type to optionally supply in order to get
             environment variables specific for the instance type.
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
-
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         Optional[str]: the s3 URI to use for the environment variable, or None if the model does not
             have gated training artifacts.
@@ -227,6 +235,7 @@ def _retrieve_gated_model_uri_env_var_value(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     s3_key: Optional[str] = (

src/sagemaker/jumpstart/artifacts/hyperparameters.py

@@ -17,6 +17,7 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.enums import (
+    JumpStartModelType,
     JumpStartScriptScope,
     VariableScope,
 )
@@ -38,6 +39,7 @@ def _retrieve_default_hyperparameters(
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     instance_type: Optional[str] = None,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ):
     """Retrieves the training hyperparameters for the model matching the given arguments.
 
@@ -71,6 +73,8 @@ def _retrieve_default_hyperparameters(
         instance_type (str): An instance type to optionally supply in order to get hyperparameters
             specific for the instance type.
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         dict: the hyperparameters to use for the model.
     """
@@ -89,6 +93,7 @@ def _retrieve_default_hyperparameters(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     default_hyperparameters: Dict[str, str] = {}

src/sagemaker/jumpstart/artifacts/image_uris.py

@@ -19,6 +19,7 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.enums import (
+    JumpStartModelType,
     JumpStartScriptScope,
     ModelFramework,
 )
@@ -48,6 +49,7 @@ def _retrieve_image_uri(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ):
     """Retrieves the container image URI for JumpStart models.
 
@@ -100,6 +102,8 @@ def _retrieve_image_uri(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         str: the ECR URI for the corresponding SageMaker Docker image.
 
@@ -123,6 +127,7 @@ def _retrieve_image_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     if image_scope == JumpStartScriptScope.INFERENCE:

src/sagemaker/jumpstart/artifacts/incremental_training.py

@@ -17,6 +17,7 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.enums import (
+    JumpStartModelType,
     JumpStartScriptScope,
 )
 from sagemaker.jumpstart.utils import (
@@ -35,6 +36,7 @@ def _model_supports_incremental_training(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> bool:
     """Returns True if the model supports incremental training.
 
@@ -59,6 +61,8 @@ def _model_supports_incremental_training(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         bool: the support status for incremental training.
     """
@@ -77,6 +81,7 @@ def _model_supports_incremental_training(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     return model_specs.supports_incremental_training()

src/sagemaker/jumpstart/artifacts/kwargs.py

@@ -167,6 +167,7 @@ def _retrieve_estimator_init_kwargs(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> dict:
     """Retrieves kwargs for `Estimator`.
 
@@ -193,6 +194,8 @@ def _retrieve_estimator_init_kwargs(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         dict: the kwargs to use for the use case.
     """
@@ -211,6 +214,7 @@ def _retrieve_estimator_init_kwargs(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     kwargs = deepcopy(model_specs.estimator_kwargs)
@@ -233,6 +237,7 @@ def _retrieve_estimator_fit_kwargs(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> dict:
     """Retrieves kwargs for `Estimator.fit`.
 
@@ -257,6 +262,8 @@ def _retrieve_estimator_fit_kwargs(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
 
     Returns:
         dict: the kwargs to use for the use case.
@@ -276,6 +283,7 @@ def _retrieve_estimator_fit_kwargs(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     return model_specs.fit_kwargs

src/sagemaker/jumpstart/artifacts/metric_definitions.py

@@ -18,6 +18,7 @@
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 )
 from sagemaker.jumpstart.enums import (
+    JumpStartModelType,
     JumpStartScriptScope,
 )
 from sagemaker.jumpstart.utils import (
@@ -37,6 +38,7 @@ def _retrieve_default_training_metric_definitions(
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     instance_type: Optional[str] = None,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Optional[List[Dict[str, str]]]:
     """Retrieves the default training metric definitions for the model.
 
@@ -63,6 +65,8 @@ def _retrieve_default_training_metric_definitions(
         instance_type (str): An instance type to optionally supply in order to get
             metric definitions specific for the instance type.
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         list: the default training metric definitions to use for the model or None.
     """
@@ -81,6 +85,7 @@ def _retrieve_default_training_metric_definitions(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         config_name=config_name,
+        model_type=model_type,
     )
 
     default_metric_definitions = (

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -130,6 +130,7 @@ def _retrieve_model_package_model_artifact_s3_uri(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     config_name: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> Optional[str]:
     """Retrieves s3 artifact uri associated with model package.
 
@@ -156,6 +157,8 @@ def _retrieve_model_package_model_artifact_s3_uri(
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+        model_type (JumpStartModelType): The type of the model, can be open weights model
+            or proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
     Returns:
         str: the model package artifact uri to use for the model or None.
 
@@ -179,6 +182,7 @@ def _retrieve_model_package_model_artifact_s3_uri(
             tolerate_deprecated_model=tolerate_deprecated_model,
             sagemaker_session=sagemaker_session,
             config_name=config_name,
+            model_type=model_type,
         )
 
         if model_specs.training_model_package_artifact_uris is None: