sagemaker-core 1.0.62__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/jumpstart/configs.py

@@ -0,0 +1,50 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains utilites for JumpStart model metadata."""
+from __future__ import absolute_import
+
+from pydantic import BaseModel, ConfigDict
+from typing import Optional
+
+
+class BaseConfig(BaseModel):
+    """BaseConfig"""
+
+    model_config = ConfigDict(
+        validate_assignment=True,
+        extra="forbid",
+        frozen=True,
+    )
+
+
+class JumpStartConfig(BaseConfig):
+    """Configuration Class for JumpStart.
+
+    Attributes:
+        model_id (str): The model ID of the JumpStart model.
+        model_version (Optional[str]): The version of the JumpStart model.
+            Defaults to None.
+        hub_name (Optional[str]): The name of the JumpStart hub. Defaults to None.
+        accept_eula (Optional[bool]): Whether to accept the EULA. Defaults to None.
+        training_config_name (Optional[str]): The name of the training configuration.
+            Defaults to None.
+        inference_config_name (Optional[str]): The name of the inference configuration.
+            Defaults to None.
+    """
+
+    model_id: str
+    model_version: Optional[str] = None
+    hub_name: Optional[str] = None
+    accept_eula: Optional[bool] = False
+    training_config_name: Optional[str] = None
+    inference_config_name: Optional[str] = None

sagemaker/core/jumpstart/constants.py

@@ -0,0 +1,198 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains constants for JumpStart."""
+from __future__ import absolute_import
+from __future__ import absolute_import
+import logging
+import os
+from typing import Dict, Set, Type
+import json
+import boto3
+from sagemaker.core.deserializers import BaseDeserializer, JSONDeserializer
+from sagemaker.core.jumpstart.enums import (
+    JumpStartScriptScope,
+    SerializerType,
+    DeserializerType,
+    MIMEType,
+    JumpStartModelType,
+)
+from sagemaker.core.jumpstart.types import JumpStartLaunchedRegionInfo, JumpStartS3FileType
+from sagemaker.core.serializers import (
+    BaseSerializer,
+    CSVSerializer,
+    DataSerializer,
+    IdentitySerializer,
+    JSONSerializer,
+)
+from sagemaker.core.helper.session_helper import Session
+
+
+SAGEMAKER_PUBLIC_HUB = "SageMakerPublicHub"
+DEFAULT_TRAINING_ENTRY_POINT = "transfer_learning.py"
+
+JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart")
+
+# disable logging if env var is set
+JUMPSTART_LOGGER.addHandler(
+    type(
+        "",
+        (logging.StreamHandler,),
+        {
+            "emit": lambda self, *args, **kwargs: (
+                logging.StreamHandler.emit(self, *args, **kwargs)
+                if not os.environ.get(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING)
+                else None
+            )
+        },
+    )()
+)
+
+
+_CURRENT_FILE_DIRECTORY_PATH = os.path.dirname(os.path.realpath(__file__))
+REGION_CONFIG_JSON_FILENAME = "region_config.json"
+REGION_CONFIG_JSON_FILEPATH = os.path.join(
+    _CURRENT_FILE_DIRECTORY_PATH, REGION_CONFIG_JSON_FILENAME
+)
+
+
+def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]:
+    """Load the JumpStart region config from a JSON file."""
+    debug_msg = f"Loading JumpStart region config from '{filepath}'."
+    JUMPSTART_LOGGER.debug(debug_msg)
+    try:
+        with open(filepath) as f:
+            config = json.load(f)
+
+        return {
+            JumpStartLaunchedRegionInfo(
+                region_name=region,
+                content_bucket=data["content_bucket"],
+                gated_content_bucket=data.get("gated_content_bucket"),
+                neo_content_bucket=data.get("neo_content_bucket"),
+            )
+            for region, data in config.items()
+        }
+    except Exception:  # pylint: disable=W0703
+        JUMPSTART_LOGGER.error("Unable to load JumpStart region config.", exc_info=True)
+        return set()
+
+
+ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING = "DISABLE_JUMPSTART_LOGGING"
+ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY = "DISABLE_JUMPSTART_TELEMETRY"
+
+JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = _load_region_config(
+    REGION_CONFIG_JSON_FILEPATH
+)
+
+JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT = {
+    region.region_name: region for region in JUMPSTART_LAUNCHED_REGIONS
+}
+JUMPSTART_REGION_NAME_SET = {region.region_name for region in JUMPSTART_LAUNCHED_REGIONS}
+
+JUMPSTART_BUCKET_NAME_SET = {region.content_bucket for region in JUMPSTART_LAUNCHED_REGIONS}
+JUMPSTART_GATED_BUCKET_NAME_SET = {
+    region.gated_content_bucket
+    for region in JUMPSTART_LAUNCHED_REGIONS
+    if region.gated_content_bucket is not None
+}
+
+JUMPSTART_GATED_AND_PUBLIC_BUCKET_NAME_SET = JUMPSTART_BUCKET_NAME_SET.union(
+    JUMPSTART_GATED_BUCKET_NAME_SET
+)
+
+JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
+NEO_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
+
+JUMPSTART_MODEL_HUB_NAME = "SageMakerPublicHub"
+
+JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
+JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY = "proprietary-sdk-manifest.json"
+
+HUB_CONTENT_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub-content/(.*?)/(.*?)/(.*?)/(.*?)$"
+HUB_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub/(.*?)$"
+
+INFERENCE_ENTRY_POINT_SCRIPT_NAME = "inference.py"
+TRAINING_ENTRY_POINT_SCRIPT_NAME = "transfer_learning.py"
+
+SUPPORTED_JUMPSTART_SCOPES = set(scope.value for scope in JumpStartScriptScope)
+
+ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE = "AWS_JUMPSTART_CONTENT_BUCKET_OVERRIDE"
+ENV_VARIABLE_JUMPSTART_GATED_CONTENT_BUCKET_OVERRIDE = "AWS_JUMPSTART_GATED_CONTENT_BUCKET_OVERRIDE"
+ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE = "AWS_JUMPSTART_MODEL_BUCKET_OVERRIDE"
+ENV_VARIABLE_JUMPSTART_SCRIPT_ARTIFACT_BUCKET_OVERRIDE = "AWS_JUMPSTART_SCRIPT_BUCKET_OVERRIDE"
+ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE = (
+    "AWS_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE"
+)
+ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE = "AWS_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE"
+ENV_VARIABLE_NEO_CONTENT_BUCKET_OVERRIDE = "AWS_NEO_CONTENT_BUCKET_OVERRIDE"
+
+JUMPSTART_RESOURCE_BASE_NAME = "sagemaker-jumpstart"
+
+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,
+    MIMEType.X_TEXT: SerializerType.TEXT,
+    MIMEType.JSON: SerializerType.JSON,
+    MIMEType.CSV: SerializerType.CSV,
+    MIMEType.WAV: SerializerType.RAW_BYTES,
+}
+
+
+ACCEPT_TYPE_TO_DESERIALIZER_TYPE_MAP: Dict[MIMEType, DeserializerType] = {
+    MIMEType.JSON: DeserializerType.JSON,
+}
+
+SERIALIZER_TYPE_TO_CLASS_MAP: Dict[SerializerType, Type[BaseSerializer]] = {
+    SerializerType.RAW_BYTES: DataSerializer,
+    SerializerType.JSON: JSONSerializer,
+    SerializerType.TEXT: IdentitySerializer,
+    SerializerType.CSV: CSVSerializer,
+}
+
+DESERIALIZER_TYPE_TO_CLASS_MAP: Dict[DeserializerType, Type[BaseDeserializer]] = {
+    DeserializerType.JSON: JSONDeserializer,
+}
+
+MODEL_TYPE_TO_MANIFEST_MAP: Dict[Type[JumpStartModelType], Type[JumpStartS3FileType]] = {
+    JumpStartModelType.OPEN_WEIGHTS: JumpStartS3FileType.OPEN_WEIGHT_MANIFEST,
+    JumpStartModelType.PROPRIETARY: JumpStartS3FileType.PROPRIETARY_MANIFEST,
+}
+
+MODEL_TYPE_TO_SPECS_MAP: Dict[Type[JumpStartModelType], Type[JumpStartS3FileType]] = {
+    JumpStartModelType.OPEN_WEIGHTS: JumpStartS3FileType.OPEN_WEIGHT_SPECS,
+    JumpStartModelType.PROPRIETARY: JumpStartS3FileType.PROPRIETARY_SPECS,
+}
+
+MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html"
+
+try:
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(
+        boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)
+    )
+except Exception as e:  # pylint: disable=W0703
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION = None
+    JUMPSTART_LOGGER.warning(
+        "Unable to create default JumpStart SageMaker Session due to the following error: %s.",
+        str(e),
+    )
+
+EXTRA_MODEL_ID_TAGS = ["sm-jumpstart-id", "sagemaker-studio:jumpstart-model-id"]
+EXTRA_MODEL_VERSION_TAGS = [
+    "sm-jumpstart-model-version",
+    "sagemaker-studio:jumpstart-model-version",
+]

sagemaker/core/jumpstart/deserializers.py

@@ -0,0 +1,81 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""JumpStart deserializers module - provides retrieve_default function for backward compatibility."""
+from __future__ import absolute_import
+
+from typing import Optional
+
+from sagemaker.core.deserializers import BaseDeserializer
+from sagemaker.core.jumpstart import artifacts, utils as jumpstart_utils
+from sagemaker.core.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.core.jumpstart.enums import JumpStartModelType
+from sagemaker.core.helper.session_helper import Session
+
+
+def retrieve_default(
+    region: Optional[str] = None,
+    model_id: Optional[str] = None,
+    model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    config_name: Optional[str] = None,
+) -> BaseDeserializer:
+    """Retrieves the default deserializer for the model matching the given arguments.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the default deserializer.
+            Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the default deserializer. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            default deserializer. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (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 models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
+        model_type (JumpStartModelType): The model type. (Default: JumpStartModelType.OPEN_WEIGHTS).
+        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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+    Returns:
+        BaseDeserializer: The default deserializer to use for the model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError(
+            "Must specify JumpStart `model_id` and `model_version` when retrieving deserializers."
+        )
+
+    return artifacts._retrieve_default_deserializer(
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        sagemaker_session=sagemaker_session,
+        model_type=model_type,
+        config_name=config_name,
+    )

sagemaker/core/jumpstart/document.py

@@ -0,0 +1,76 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains utilites for JumpStart model metadata."""
+from __future__ import absolute_import
+
+import json
+from typing import Optional, Tuple
+from functools import lru_cache
+from botocore.exceptions import ClientError
+
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.utils.utils import logger
+from sagemaker.core.resources import HubContent
+from sagemaker.core.jumpstart.configs import JumpStartConfig
+from sagemaker.core.jumpstart.models import HubContentDocument
+from sagemaker.core.jumpstart.constants import SAGEMAKER_PUBLIC_HUB
+
+
+@lru_cache(maxsize=128)
+def get_hub_content_and_document(
+    jumpstart_config: JumpStartConfig,
+    sagemaker_session: Optional[Session] = None,
+) -> Tuple[HubContent, HubContentDocument]:
+    """Get model metadata for JumpStart.
+
+
+    Args:
+        jumpstart_config (JumpStartConfig): JumpStart configuration.
+        sagemaker_session (Session, optional): SageMaker session.
+            Defaults to None.
+
+    Returns:
+        HubContentDocument: Model metadata.
+    """
+    if sagemaker_session is None:
+        sagemaker_session = Session()
+        logger.debug("No sagemaker session provided. Using default session.")
+
+    hub_name = jumpstart_config.hub_name if jumpstart_config.hub_name else SAGEMAKER_PUBLIC_HUB
+    hub_content_type = "Model" if hub_name == SAGEMAKER_PUBLIC_HUB else "ModelReference"
+
+    region = sagemaker_session.boto_region_name
+
+    try:
+        hub_content = HubContent.get(
+            hub_name=hub_name,
+            hub_content_name=jumpstart_config.model_id,
+            hub_content_version=jumpstart_config.model_version,
+            hub_content_type=hub_content_type,
+            session=sagemaker_session.boto_session,
+            region=region,
+        )
+    except ClientError as e:
+        if e.response["Error"]["Code"] == "ResourceNotFound":
+            logger.error(
+                f"Hub content {jumpstart_config.model_id} not found in {hub_name}.\n"
+                "Please check that the Model ID is availble in the specified hub."
+            )
+        raise e
+
+    logger.info(
+        f"hub_content_name: {hub_content.hub_content_name}, "
+        f"hub_content_version: {hub_content.hub_content_version}"
+    )
+    document_json = json.loads(hub_content.hub_content_document)
+    return (hub_content, HubContentDocument(**document_json))

sagemaker/core/jumpstart/enums.py

@@ -0,0 +1,168 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module stores enums related to SageMaker JumpStart."""
+
+from __future__ import absolute_import
+
+from enum import Enum
+from typing import List
+
+
+class ModelFramework(str, Enum):
+    """Enum class for JumpStart model framework.
+
+    The ML framework as referenced in the prefix of the model ID.
+    This value does not necessarily correspond to the container name.
+    """
+
+    PYTORCH = "pytorch"
+    TENSORFLOW = "tensorflow"
+    MXNET = "mxnet"
+    HUGGINGFACE = "huggingface"
+    LIGHTGBM = "lightgbm"
+    CATBOOST = "catboost"
+    XGBOOST = "xgboost"
+    SKLEARN = "sklearn"
+
+
+class JumpStartModelType(str, Enum):
+    """Enum class for JumpStart model type.
+
+    OPEN_WEIGHTS: Publicly available models have open weights
+    and are onboarded and maintained by JumpStart.
+    PROPRIETARY: Proprietary models from third-party providers do not have open weights.
+    You must subscribe to proprietary models in AWS Marketplace before use.
+    """
+
+    OPEN_WEIGHTS = "open_weights"
+    PROPRIETARY = "proprietary"
+
+
+class VariableScope(str, Enum):
+    """Possible value of the ``scope`` attribute for a hyperparameter or environment variable.
+
+    Used for hosting environment variables and training hyperparameters.
+    """
+
+    CONTAINER = "container"
+    ALGORITHM = "algorithm"
+
+
+class JumpStartScriptScope(str, Enum):
+    """Enum class for JumpStart script scopes."""
+
+    INFERENCE = "inference"
+    TRAINING = "training"
+
+
+class HyperparameterValidationMode(str, Enum):
+    """Possible modes for validating hyperparameters."""
+
+    VALIDATE_PROVIDED = "validate_provided"
+    VALIDATE_ALGORITHM = "validate_algorithm"
+    VALIDATE_ALL = "validate_all"
+
+
+class VariableTypes(str, Enum):
+    """Possible types for hyperparameters and environment variables."""
+
+    TEXT = "text"
+    INT = "int"
+    FLOAT = "float"
+    BOOL = "bool"
+
+
+class HubContentCapability(str, Enum):
+    """Enum class for HubContent capabilities."""
+
+    BEDROCK_CONSOLE = "BEDROCK_CONSOLE"
+
+
+class JumpStartTag(str, Enum):
+    """Enum class for tag keys to apply to JumpStart models."""
+
+    INFERENCE_MODEL_URI = "aws-jumpstart-inference-model-uri"
+    INFERENCE_SCRIPT_URI = "aws-jumpstart-inference-script-uri"
+    TRAINING_MODEL_URI = "aws-jumpstart-training-model-uri"
+    TRAINING_SCRIPT_URI = "aws-jumpstart-training-script-uri"
+
+    MODEL_ID = "sagemaker-sdk:jumpstart-model-id"
+    MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
+    MODEL_TYPE = "sagemaker-sdk:jumpstart-model-type"
+
+    INFERENCE_CONFIG_NAME = "sagemaker-sdk:jumpstart-inference-config-name"
+    TRAINING_CONFIG_NAME = "sagemaker-sdk:jumpstart-training-config-name"
+
+    HUB_CONTENT_ARN = "sagemaker-sdk:hub-content-arn"
+
+    BEDROCK = "sagemaker-sdk:bedrock"
+
+
+class SerializerType(str, Enum):
+    """Enum class for serializers associated with JumpStart models."""
+
+    TEXT = "text"
+    JSON = "json"
+    CSV = "csv"
+    RAW_BYTES = "raw_bytes"
+
+
+class DeserializerType(str, Enum):
+    """Enum class for deserializers associated with JumpStart models."""
+
+    JSON = "json"
+
+
+class MIMEType(str, Enum):
+    """Enum class for MIME types associated with JumpStart models."""
+
+    X_IMAGE = "application/x-image"
+    LIST_TEXT = "application/list-text"
+    X_TEXT = "application/x-text"
+    JSON = "application/json"
+    CSV = "text/csv"
+    WAV = "audio/wav"
+
+    @staticmethod
+    def from_suffixed_type(mime_type_with_suffix: str) -> "MIMEType":
+        """Removes suffix from type and instantiates enum."""
+        base_type, _, _ = mime_type_with_suffix.partition(";")
+        return MIMEType(base_type)
+
+
+class NamingConventionType(str, Enum):
+    """Enum class for naming conventions."""
+
+    SNAKE_CASE = "snake_case"
+    UPPER_CAMEL_CASE = "upper_camel_case"
+    DEFAULT = UPPER_CAMEL_CASE
+
+
+class ModelSpecKwargType(str, Enum):
+    """Enum class for types of kwargs for model hub content document and model specs."""
+
+    FIT = "fit_kwargs"
+    MODEL = "model_kwargs"
+    ESTIMATOR = "estimator_kwargs"
+    DEPLOY = "deploy_kwargs"
+
+    @classmethod
+    def arg_keys(cls) -> List[str]:
+        """Returns a list of kwargs keys that each type can have"""
+        return [member.value for member in cls]
+
+
+class JumpStartConfigRankingName(str, Enum):
+    """Enum class for ranking of JumpStart config."""
+
+    DEFAULT = "overall"