Merge branch 'master' into hf-inf-spec-pr

Author: bryannahm1

requirements/extras/test_requirements.txt

@@ -38,3 +38,4 @@ accelerate>=0.24.1,<=0.27.0
 schema==0.7.5
 tensorflow>=2.1,<=2.16
 mlflow>=2.12.2,<2.13
+huggingface_hub>=0.23.4

setup.py

@@ -49,7 +49,7 @@ def read_requirements(filename):
 # Declare minimal set for installation
 required_packages = [
     "attrs>=23.1.0,<24",
-    "boto3>=1.33.3,<2.0",
+    "boto3>=1.34.142,<2.0",
     "cloudpickle==2.2.1",
     "google-pasta",
     "numpy>=1.9.0,<2.0",

src/sagemaker/accept_types.py

@@ -82,6 +82,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    config_name: Optional[str] = None,
 ) -> str:
     """Retrieves the default accept type for the model matching the given arguments.
 
@@ -105,6 +106,7 @@ def retrieve_default(
             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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
     Returns:
         str: The default accept type to use for the model.
 
@@ -125,4 +127,5 @@ def retrieve_default(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         model_type=model_type,
+        config_name=config_name,
     )

src/sagemaker/content_types.py

@@ -82,6 +82,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    config_name: Optional[str] = None,
 ) -> str:
     """Retrieves the default content type for the model matching the given arguments.
 
@@ -105,6 +106,7 @@ def retrieve_default(
             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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
     Returns:
         str: The default content type to use for the model.
 
@@ -125,6 +127,7 @@ def retrieve_default(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         model_type=model_type,
+        config_name=config_name,
     )
 
 

src/sagemaker/deserializers.py

@@ -102,6 +102,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    config_name: Optional[str] = None,
 ) -> BaseDeserializer:
     """Retrieves the default deserializer for the model matching the given arguments.
 
@@ -125,6 +126,7 @@ def retrieve_default(
             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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
     Returns:
         BaseDeserializer: The default deserializer to use for the model.
 
@@ -146,4 +148,5 @@ def retrieve_default(
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         model_type=model_type,
+        config_name=config_name,
     )

src/sagemaker/enums.py

@@ -40,3 +40,12 @@ class RoutingStrategy(Enum):
     """The endpoint routes requests to the specific instances that have
     more capacity to process them.
     """
+
+
+class Tag(str, Enum):
+    """Enum class for tag keys to apply to models."""
+
+    OPTIMIZATION_JOB_NAME = "sagemaker-sdk:optimization-job-name"
+    SPECULATIVE_DRAFT_MODEL_PROVIDER = "sagemaker-sdk:speculative-draft-model-provider"
+    FINE_TUNING_MODEL_PATH = "sagemaker-sdk:fine-tuning-model-path"
+    FINE_TUNING_JOB_NAME = "sagemaker-sdk:fine-tuning-job-name"

src/sagemaker/environment_variables.py

@@ -37,6 +37,7 @@ def retrieve_default(
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     instance_type: Optional[str] = None,
     script: JumpStartScriptScope = JumpStartScriptScope.INFERENCE,
+    config_name: Optional[str] = None,
 ) -> Dict[str, str]:
     """Retrieves the default container environment variables for the model matching the arguments.
 
@@ -68,6 +69,7 @@ def retrieve_default(
             variables specific for the instance type.
         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).
     Returns:
         dict: The variables to use for the model.
 
@@ -91,4 +93,5 @@ def retrieve_default(
         sagemaker_session=sagemaker_session,
         instance_type=instance_type,
         script=script,
+        config_name=config_name,
     )

src/sagemaker/huggingface/llm_utils.py

@@ -13,7 +13,9 @@
 """Functions for generating ECR image URIs for pre-built SageMaker Docker images."""
 from __future__ import absolute_import
 
+import os
 from typing import Optional
+import importlib.util
 
 import urllib.request
 from urllib.error import HTTPError, URLError
@@ -123,3 +125,26 @@ def get_huggingface_model_metadata(model_id: str, hf_hub_token: Optional[str] =
             "Did not find model metadata for the following HuggingFace Model ID %s" % model_id
         )
     return hf_model_metadata_json
+
+
+def download_huggingface_model_metadata(
+    model_id: str, model_local_path: str, hf_hub_token: Optional[str] = None
+) -> None:
+    """Downloads the HuggingFace Model snapshot via HuggingFace API.
+
+    Args:
+        model_id (str): The HuggingFace Model ID
+        model_local_path (str): The local path to save the HuggingFace Model snapshot.
+        hf_hub_token (str): The HuggingFace Hub Token
+
+    Raises:
+        ImportError: If huggingface_hub is not installed.
+    """
+    if not importlib.util.find_spec("huggingface_hub"):
+        raise ImportError("Unable to import huggingface_hub, check if huggingface_hub is installed")
+
+    from huggingface_hub import snapshot_download
+
+    os.makedirs(model_local_path, exist_ok=True)
+    logger.info("Downloading model %s from Hugging Face Hub to %s", model_id, model_local_path)
+    snapshot_download(repo_id=model_id, local_dir=model_local_path, token=hf_hub_token)

src/sagemaker/hyperparameters.py

@@ -37,6 +37,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    config_name: Optional[str] = None,
 ) -> Dict[str, str]:
     """Retrieves the default training hyperparameters for the model matching the given arguments.
 
@@ -69,6 +70,7 @@ def retrieve_default(
             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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
     Returns:
         dict: The hyperparameters to use for the model.
 
@@ -90,6 +92,7 @@ def retrieve_default(
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        config_name=config_name,
     )
 
 

src/sagemaker/image_uris.py

@@ -71,6 +71,7 @@ def retrieve(
     inference_tool=None,
     serverless_inference_config=None,
     sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    config_name=None,
 ) -> str:
     """Retrieves the ECR URI for the Docker image matching the given arguments.
 
@@ -126,6 +127,7 @@ def retrieve(
             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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
 
     Returns:
         str: The ECR URI for the corresponding SageMaker Docker image.
@@ -166,6 +168,7 @@ def retrieve(
             tolerate_vulnerable_model,
             tolerate_deprecated_model,
             sagemaker_session=sagemaker_session,
+            config_name=config_name,
         )
 
     if training_compiler_config and (framework in [HUGGING_FACE_FRAMEWORK, "pytorch"]):