fix: gated training environment variables for jumpstart (#1338)

Author: evakravi

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -16,11 +16,13 @@
 from sagemaker.jumpstart.constants import (
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     JUMPSTART_DEFAULT_REGION_NAME,
+    SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import (
     JumpStartScriptScope,
 )
 from sagemaker.jumpstart.utils import (
+    get_jumpstart_gated_content_bucket,
     verify_model_region_and_return_specs,
 )
 from sagemaker.session import Session
@@ -102,10 +104,89 @@ def _retrieve_default_environment_variables(
         elif script == JumpStartScriptScope.TRAINING and getattr(
             model_specs, "training_instance_type_variants", None
         ):
-            default_environment_variables.update(
-                model_specs.training_instance_type_variants.get_instance_specific_environment_variables(  # noqa E501  # pylint: disable=c0301
-                    instance_type
-                )
+            instance_specific_environment_variables = model_specs.training_instance_type_variants.get_instance_specific_environment_variables(  # noqa E501  # pylint: disable=c0301
+                instance_type
             )
 
+            default_environment_variables.update(instance_specific_environment_variables)
+
+            gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
+                model_id=model_id,
+                model_version=model_version,
+                region=region,
+                tolerate_vulnerable_model=tolerate_vulnerable_model,
+                tolerate_deprecated_model=tolerate_deprecated_model,
+                sagemaker_session=sagemaker_session,
+                instance_type=instance_type,
+            )
+
+            if gated_model_env_var is not None:
+                default_environment_variables.update(
+                    {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: gated_model_env_var}
+                )
+
     return default_environment_variables
+
+
+def _retrieve_gated_model_uri_env_var_value(
+    model_id: str,
+    model_version: str,
+    region: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    instance_type: Optional[str] = None,
+) -> Optional[str]:
+    """Retrieves the gated model env var URI matching the given arguments.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the gated model env var URI.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            gated model env var URI.
+        region (Optional[str]): Region for which to retrieve the gated model env var URI.
+            (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        instance_type (str): An instance type to optionally supply in order to get
+            environment variables specific for the instance type.
+
+    Returns:
+        Optional[str]: the s3 URI to use for the environment variable, or None if the model does not
+            have gated training artifacts.
+
+    Raises:
+        ValueError: If the model specs specified are invalid.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        scope=JumpStartScriptScope.TRAINING,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+    )
+
+    s3_key: Optional[
+        str
+    ] = model_specs.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value(  # noqa E501  # pylint: disable=c0301
+        instance_type
+    )
+    if s3_key is None:
+        return None
+
+    return f"s3://{get_jumpstart_gated_content_bucket(region)}/{s3_key}"

src/sagemaker/jumpstart/factory/estimator.py

@@ -16,6 +16,7 @@
 
 from typing import Dict, List, Optional, Union
 from sagemaker import (
+    environment_variables,
     hyperparameters as hyperparameters_utils,
     image_uris,
     instance_types,
@@ -557,6 +558,18 @@ def _add_env_to_kwargs(
 ) -> JumpStartEstimatorInitKwargs:
     """Sets environment in kwargs based on default or override, returns full kwargs."""
 
+    extra_env_vars = environment_variables.retrieve_default(
+        model_id=kwargs.model_id,
+        model_version=kwargs.model_version,
+        region=kwargs.region,
+        include_aws_sdk_env_vars=False,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+        sagemaker_session=kwargs.sagemaker_session,
+        script=JumpStartScriptScope.TRAINING,
+        instance_type=kwargs.instance_type,
+    )
+
     model_package_artifact_uri = _retrieve_model_package_model_artifact_s3_uri(
         model_id=kwargs.model_id,
         model_version=kwargs.model_version,
@@ -568,12 +581,16 @@ def _add_env_to_kwargs(
     )
 
     if model_package_artifact_uri:
-        if kwargs.environment is None:
-            kwargs.environment = {}
-        kwargs.environment = {
-            **{SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: model_package_artifact_uri},
-            **kwargs.environment,
-        }
+        extra_env_vars.update(
+            {SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: model_package_artifact_uri}
+        )
+
+    for key, value in extra_env_vars.items():
+        update_dict_if_key_not_present(
+            kwargs.environment,
+            key,
+            value,
+        )
 
     return kwargs
 

src/sagemaker/jumpstart/types.py

@@ -581,6 +581,16 @@ def get_instance_specific_environment_variables(self, instance_type: str) -> Dic
 
         return instance_family_environment_variables
 
+    def get_instance_specific_gated_model_key_env_var_value(
+        self, instance_type: str
+    ) -> Optional[str]:
+        """Returns instance specific gated model env var s3 key.
+
+        Returns None if a model, instance type tuple does not have instance
+        specific property.
+        """
+        return self._get_instance_specific_property(instance_type, "gated_model_key_env_var_value")
+
     def get_instance_specific_default_inference_instance_type(
         self, instance_type: str
     ) -> Optional[str]:
@@ -901,10 +911,12 @@ def use_inference_script_uri(self) -> bool:
 
     def use_training_model_artifact(self) -> bool:
         """Returns True if the model should use a model uri when kicking off training job."""
-        return (
-            self.training_model_package_artifact_uris is None
-            or len(self.training_model_package_artifact_uris) == 0
-        )
+        # gated model never use training model artifact
+        if self.gated_bucket:
+            return False
+
+        # otherwise, return true is a training model package is not set
+        return len(self.training_model_package_artifact_uris or {}) == 0
 
     def supports_incremental_training(self) -> bool:
         """Returns True if the model supports incremental training."""

src/sagemaker/jumpstart/utils.py

@@ -611,11 +611,18 @@ def verify_model_region_and_return_specs(
 
 
 def update_dict_if_key_not_present(
-    dict_to_update: dict, key_to_add: Any, value_to_add: Any
-) -> dict:
-    """If a key is not present in the dict, add the new (key, value) pair, and return dict."""
+    dict_to_update: Optional[dict], key_to_add: Any, value_to_add: Any
+) -> Optional[dict]:
+    """If a key is not present in the dict, add the new (key, value) pair, and return dict.
+
+    If dict is empty, return None.
+    """
+    if dict_to_update is None:
+        dict_to_update = {}
     if key_to_add not in dict_to_update:
         dict_to_update[key_to_add] = value_to_add
+    if dict_to_update == {}:
+        dict_to_update = None
 
     return dict_to_update