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

sagemaker/core/jumpstart/hub/parser_utils.py

@@ -0,0 +1,70 @@
+# 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.
+# pylint: skip-file
+"""This module contains utilities related to SageMaker JumpStart Hub."""
+from __future__ import absolute_import
+
+import re
+from typing import Any, Dict, List, Optional
+
+
+def camel_to_snake(camel_case_string: str) -> str:
+    """Converts camelCase to snake_case_string using a regex.
+
+    This regex cannot handle whitespace ("camelString TwoWords")
+    """
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", camel_case_string).lower()
+
+
+def snake_to_upper_camel(snake_case_string: str) -> str:
+    """Converts snake_case_string to UpperCamelCaseString."""
+    upper_camel_case_string = "".join(word.title() for word in snake_case_string.split("_"))
+    return upper_camel_case_string
+
+
+def walk_and_apply_json(
+    json_obj: Dict[Any, Any], apply, stop_keys: Optional[List[str]] = ["metrics"]
+) -> Dict[Any, Any]:
+    """Recursively walks a json object and applies a given function to the keys.
+
+    stop_keys (Optional[list[str]]): List of field keys that should stop the application function.
+        Any children of these keys will not have the application function applied to them.
+    """
+
+    def _walk_and_apply_json(json_obj, new):
+        if isinstance(json_obj, dict) and isinstance(new, dict):
+            for key, value in json_obj.items():
+                new_key = apply(key)
+                if (stop_keys and new_key not in stop_keys) or stop_keys is None:
+                    if isinstance(value, dict):
+                        new[new_key] = {}
+                        _walk_and_apply_json(value, new=new[new_key])
+                    elif isinstance(value, list):
+                        new[new_key] = []
+                        for item in value:
+                            _walk_and_apply_json(item, new=new[new_key])
+                    else:
+                        new[new_key] = value
+                else:
+                    new[new_key] = value
+        elif isinstance(json_obj, dict) and isinstance(new, list):
+            new.append(_walk_and_apply_json(json_obj, new={}))
+        elif isinstance(json_obj, list) and isinstance(new, dict):
+            new.update(json_obj)
+        elif isinstance(json_obj, list) and isinstance(new, list):
+            new.append(json_obj)
+        elif isinstance(json_obj, str) and isinstance(new, list):
+            new.append(json_obj)
+        return new
+
+    return _walk_and_apply_json(json_obj, new={})

sagemaker/core/jumpstart/hub/parsers.py

@@ -0,0 +1,288 @@
+# 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.
+# pylint: skip-file
+"""This module stores Hub converter utilities for JumpStart."""
+from __future__ import absolute_import
+
+from typing import Any, Dict, List
+from sagemaker.core.jumpstart.enums import ModelSpecKwargType, NamingConventionType
+from sagemaker.core.s3 import parse_s3_url
+from sagemaker.core.jumpstart.types import (
+    JumpStartModelSpecs,
+    HubContentType,
+    JumpStartDataHolderType,
+)
+from sagemaker.core.jumpstart.hub.interfaces import (
+    DescribeHubContentResponse,
+    HubModelDocument,
+)
+from sagemaker.core.jumpstart.hub.parser_utils import (
+    camel_to_snake,
+    snake_to_upper_camel,
+    walk_and_apply_json,
+)
+
+
+def _to_json(dictionary: Dict[Any, Any]) -> Dict[Any, Any]:
+    """Convert a nested dictionary of JumpStartDataHolderType into json with UpperCamelCase keys"""
+    for key, value in dictionary.items():
+        if issubclass(type(value), JumpStartDataHolderType):
+            dictionary[key] = walk_and_apply_json(value.to_json(), snake_to_upper_camel)
+        elif isinstance(value, list):
+            new_value = []
+            for value_in_list in value:
+                new_value_in_list = value_in_list
+                if issubclass(type(value_in_list), JumpStartDataHolderType):
+                    new_value_in_list = walk_and_apply_json(
+                        value_in_list.to_json(), snake_to_upper_camel
+                    )
+                new_value.append(new_value_in_list)
+            dictionary[key] = new_value
+        elif isinstance(value, dict):
+            for key_in_dict, value_in_dict in value.items():
+                if issubclass(type(value_in_dict), JumpStartDataHolderType):
+                    value[key_in_dict] = walk_and_apply_json(
+                        value_in_dict.to_json(), snake_to_upper_camel
+                    )
+    return dictionary
+
+
+def get_model_spec_arg_keys(
+    arg_type: ModelSpecKwargType,
+    naming_convention: NamingConventionType = NamingConventionType.DEFAULT,
+) -> List[str]:
+    """Returns a list of arg keys for a specific model spec arg type.
+
+    Args:
+        arg_type (ModelSpecKwargType): Type of the model spec's kwarg.
+        naming_convention (NamingConventionType): Type of naming convention to return.
+
+    Raises:
+        ValueError: If the naming convention is not valid.
+    """
+    arg_keys: List[str] = []
+    if arg_type == ModelSpecKwargType.DEPLOY:
+        arg_keys = [
+            "ModelDataDownloadTimeout",
+            "ContainerStartupHealthCheckTimeout",
+            "InferenceAmiVersion",
+        ]
+    elif arg_type == ModelSpecKwargType.ESTIMATOR:
+        arg_keys = [
+            "EncryptInterContainerTraffic",
+            "MaxRuntimeInSeconds",
+            "DisableOutputCompression",
+            "ModelDir",
+        ]
+    elif arg_type == ModelSpecKwargType.MODEL:
+        arg_keys = []
+    elif arg_type == ModelSpecKwargType.FIT:
+        arg_keys = []
+
+    if naming_convention == NamingConventionType.SNAKE_CASE:
+        arg_keys = [camel_to_snake(key) for key in arg_keys]
+    elif naming_convention == NamingConventionType.UPPER_CAMEL_CASE:
+        return arg_keys
+    else:
+        raise ValueError("Please provide a valid naming convention.")
+    return arg_keys
+
+
+def get_model_spec_kwargs_from_hub_model_document(
+    arg_type: ModelSpecKwargType,
+    hub_content_document: Dict[str, Any],
+    naming_convention: NamingConventionType = NamingConventionType.UPPER_CAMEL_CASE,
+) -> Dict[str, Any]:
+    """Returns a map of arg type to arg keys for a given hub content document.
+
+    Args:
+        arg_type (ModelSpecKwargType): Type of the model spec's kwarg.
+        hub_content_document: A dictionary representation of hub content document.
+        naming_convention (NamingConventionType): Type of naming convention to return.
+
+    """
+    kwargs = dict()
+    keys = get_model_spec_arg_keys(arg_type, naming_convention=naming_convention)
+    for k in keys:
+        kwarg_value = hub_content_document.get(k)
+        if kwarg_value is not None:
+            kwargs[k] = kwarg_value
+    return kwargs
+
+
+def make_model_specs_from_describe_hub_content_response(
+    response: DescribeHubContentResponse,
+) -> JumpStartModelSpecs:
+    """Sets fields in JumpStartModelSpecs based on values in DescribeHubContentResponse
+
+    Args:
+        response (Dict[str, any]): parsed DescribeHubContentResponse returned
+            from SageMaker:DescribeHubContent
+    """
+    if response.hub_content_type not in {HubContentType.MODEL, HubContentType.MODEL_REFERENCE}:
+        raise AttributeError(
+            "Invalid content type, use either HubContentType.MODEL or HubContentType.MODEL_REFERENCE."
+        )
+    region = response.get_hub_region()
+    specs = {}
+    model_id = response.hub_content_name
+    specs["model_id"] = model_id
+    specs["version"] = response.hub_content_version
+    hub_model_document: HubModelDocument = response.hub_content_document
+    specs["url"] = hub_model_document.url
+    specs["min_sdk_version"] = hub_model_document.min_sdk_version
+    specs["model_types"] = hub_model_document.model_types
+    specs["capabilities"] = hub_model_document.capabilities
+    specs["training_supported"] = bool(hub_model_document.training_supported)
+    specs["incremental_training_supported"] = bool(
+        hub_model_document.incremental_training_supported
+    )
+    specs["hosting_ecr_uri"] = hub_model_document.hosting_ecr_uri
+    specs["inference_configs"] = hub_model_document.inference_configs
+    specs["inference_config_components"] = hub_model_document.inference_config_components
+    specs["inference_config_rankings"] = hub_model_document.inference_config_rankings
+
+    if hub_model_document.hosting_artifact_uri:
+        _, hosting_artifact_key = parse_s3_url(  # pylint: disable=unused-variable
+            hub_model_document.hosting_artifact_uri
+        )
+        specs["hosting_artifact_key"] = hosting_artifact_key
+        specs["hosting_artifact_uri"] = hub_model_document.hosting_artifact_uri
+
+    if hub_model_document.hosting_script_uri:
+        _, hosting_script_key = parse_s3_url(  # pylint: disable=unused-variable
+            hub_model_document.hosting_script_uri
+        )
+        specs["hosting_script_key"] = hosting_script_key
+
+    specs["inference_environment_variables"] = hub_model_document.inference_environment_variables
+    specs["inference_vulnerable"] = False
+    specs["inference_dependencies"] = hub_model_document.inference_dependencies
+    specs["inference_vulnerabilities"] = []
+    specs["training_vulnerable"] = False
+    specs["training_vulnerabilities"] = []
+    specs["deprecated"] = False
+    specs["deprecated_message"] = None
+    specs["deprecate_warn_message"] = None
+    specs["usage_info_message"] = None
+    specs["default_inference_instance_type"] = hub_model_document.default_inference_instance_type
+    specs["supported_inference_instance_types"] = (
+        hub_model_document.supported_inference_instance_types
+    )
+    specs["dynamic_container_deployment_supported"] = (
+        hub_model_document.dynamic_container_deployment_supported
+    )
+    specs["hosting_resource_requirements"] = hub_model_document.hosting_resource_requirements
+
+    specs["hosting_prepacked_artifact_key"] = None
+    if hub_model_document.hosting_prepacked_artifact_uri is not None:
+        (
+            hosting_prepacked_artifact_bucket,  # pylint: disable=unused-variable
+            hosting_prepacked_artifact_key,
+        ) = parse_s3_url(hub_model_document.hosting_prepacked_artifact_uri)
+        specs["hosting_prepacked_artifact_key"] = hosting_prepacked_artifact_key
+
+    hub_content_document_dict: Dict[str, Any] = hub_model_document.to_json()
+
+    specs["fit_kwargs"] = get_model_spec_kwargs_from_hub_model_document(
+        ModelSpecKwargType.FIT, hub_content_document_dict
+    )
+    specs["model_kwargs"] = get_model_spec_kwargs_from_hub_model_document(
+        ModelSpecKwargType.MODEL, hub_content_document_dict
+    )
+    specs["deploy_kwargs"] = get_model_spec_kwargs_from_hub_model_document(
+        ModelSpecKwargType.DEPLOY, hub_content_document_dict
+    )
+    specs["estimator_kwargs"] = get_model_spec_kwargs_from_hub_model_document(
+        ModelSpecKwargType.ESTIMATOR, hub_content_document_dict
+    )
+
+    specs["predictor_specs"] = hub_model_document.sage_maker_sdk_predictor_specifications
+    default_payloads: Dict[str, Any] = {}
+    if hub_model_document.default_payloads is not None:
+        for alias, payload in hub_model_document.default_payloads.items():
+            default_payloads[alias] = walk_and_apply_json(payload.to_json(), camel_to_snake)
+        specs["default_payloads"] = default_payloads
+    specs["gated_bucket"] = hub_model_document.gated_bucket
+    specs["inference_volume_size"] = hub_model_document.inference_volume_size
+    specs["inference_enable_network_isolation"] = (
+        hub_model_document.inference_enable_network_isolation
+    )
+    specs["resource_name_base"] = hub_model_document.resource_name_base
+
+    specs["hosting_eula_key"] = None
+    if hub_model_document.hosting_eula_uri is not None:
+        hosting_eula_bucket, hosting_eula_key = parse_s3_url(  # pylint: disable=unused-variable
+            hub_model_document.hosting_eula_uri
+        )
+        specs["hosting_eula_key"] = hosting_eula_key
+
+    if hub_model_document.hosting_model_package_arn:
+        specs["hosting_model_package_arns"] = {region: hub_model_document.hosting_model_package_arn}
+
+    specs["model_subscription_link"] = hub_model_document.model_subscription_link
+
+    specs["hosting_use_script_uri"] = hub_model_document.hosting_use_script_uri
+
+    specs["hosting_instance_type_variants"] = hub_model_document.hosting_instance_type_variants
+
+    if specs["training_supported"]:
+        specs["training_ecr_uri"] = hub_model_document.training_ecr_uri
+        (
+            training_artifact_bucket,  # pylint: disable=unused-variable
+            training_artifact_key,
+        ) = parse_s3_url(hub_model_document.training_artifact_uri)
+        specs["training_artifact_key"] = training_artifact_key
+        (
+            training_script_bucket,  # pylint: disable=unused-variable
+            training_script_key,
+        ) = parse_s3_url(hub_model_document.training_script_uri)
+        specs["training_script_key"] = training_script_key
+
+        specs["training_configs"] = hub_model_document.training_configs
+        specs["training_config_components"] = hub_model_document.training_config_components
+        specs["training_config_rankings"] = hub_model_document.training_config_rankings
+
+        specs["training_dependencies"] = hub_model_document.training_dependencies
+        specs["default_training_instance_type"] = hub_model_document.default_training_instance_type
+        specs["supported_training_instance_types"] = (
+            hub_model_document.supported_training_instance_types
+        )
+        specs["metrics"] = hub_model_document.training_metrics
+        specs["training_prepacked_script_key"] = None
+        if hub_model_document.training_prepacked_script_uri is not None:
+            (
+                training_prepacked_script_bucket,  # pylint: disable=unused-variable
+                training_prepacked_script_key,
+            ) = parse_s3_url(hub_model_document.training_prepacked_script_uri)
+            specs["training_prepacked_script_key"] = training_prepacked_script_key
+
+        specs["hyperparameters"] = hub_model_document.hyperparameters
+        specs["training_volume_size"] = hub_model_document.training_volume_size
+        specs["training_enable_network_isolation"] = (
+            hub_model_document.training_enable_network_isolation
+        )
+        if hub_model_document.training_model_package_artifact_uri:
+            specs["training_model_package_artifact_uris"] = {
+                region: hub_model_document.training_model_package_artifact_uri
+            }
+        specs["training_instance_type_variants"] = (
+            hub_model_document.training_instance_type_variants
+        )
+        if hub_model_document.default_training_dataset_uri:
+            _, default_training_dataset_key = parse_s3_url(  # pylint: disable=unused-variable
+                hub_model_document.default_training_dataset_uri
+            )
+            specs["default_training_dataset_key"] = default_training_dataset_key
+            specs["default_training_dataset_uri"] = hub_model_document.default_training_dataset_uri
+    return JumpStartModelSpecs(_to_json(specs), is_hub_content=True)

sagemaker/core/jumpstart/hub/types.py

@@ -0,0 +1,35 @@
+# 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 types related to SageMaker JumpStart Hub."""
+from __future__ import absolute_import
+from typing import Dict
+from dataclasses import dataclass
+
+
+@dataclass
+class S3ObjectLocation:
+    """Helper class for S3 object references."""
+
+    bucket: str
+    key: str
+
+    def format_for_s3_copy(self) -> Dict[str, str]:
+        """Returns a dict formatted for S3 copy calls"""
+        return {
+            "Bucket": self.bucket,
+            "Key": self.key,
+        }
+
+    def get_uri(self) -> str:
+        """Returns the s3 URI"""
+        return f"s3://{self.bucket}/{self.key}"

sagemaker/core/jumpstart/hub/utils.py

@@ -0,0 +1,260 @@
+# 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.
+# pylint: skip-file
+"""This module contains utilities related to SageMaker JumpStart Hub."""
+from __future__ import absolute_import
+import re
+from typing import Optional, List, Any
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.common_utils import aws_partition
+from sagemaker.core.jumpstart.types import HubContentType, HubArnExtractedInfo
+from sagemaker.core.jumpstart import constants
+from packaging.specifiers import SpecifierSet, InvalidSpecifier
+from packaging import version
+
+PROPRIETARY_VERSION_KEYWORD = "@marketplace-version:"
+
+
+def _convert_str_to_optional(string: str) -> Optional[str]:
+    if string == "None":
+        string = None
+    return string
+
+
+def get_info_from_hub_resource_arn(
+    arn: str,
+) -> HubArnExtractedInfo:
+    """Extracts descriptive information from a Hub or HubContent Arn."""
+
+    match = re.match(constants.HUB_CONTENT_ARN_REGEX, arn)
+    if match:
+        partition = match.group(1)
+        hub_region = match.group(2)
+        account_id = match.group(3)
+        hub_name = match.group(4)
+        hub_content_type = match.group(5)
+        hub_content_name = match.group(6)
+        hub_content_version = _convert_str_to_optional(match.group(7))
+
+        return HubArnExtractedInfo(
+            partition=partition,
+            region=hub_region,
+            account_id=account_id,
+            hub_name=hub_name,
+            hub_content_type=hub_content_type,
+            hub_content_name=hub_content_name,
+            hub_content_version=hub_content_version,
+        )
+
+    match = re.match(constants.HUB_ARN_REGEX, arn)
+    if match:
+        partition = match.group(1)
+        hub_region = match.group(2)
+        account_id = match.group(3)
+        hub_name = match.group(4)
+        return HubArnExtractedInfo(
+            partition=partition,
+            region=hub_region,
+            account_id=account_id,
+            hub_name=hub_name,
+        )
+
+
+def construct_hub_arn_from_name(
+    hub_name: str,
+    region: Optional[str] = None,
+    session: Optional[Session] = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    account_id: Optional[str] = None,
+) -> str:
+    """Constructs a Hub arn from the Hub name using default Session values."""
+    if session is None:
+        # session is overridden to none by some callers
+        session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+
+    account_id = account_id or session.account_id()
+    region = region or session.boto_region_name
+    partition = aws_partition(region)
+
+    return f"arn:{partition}:sagemaker:{region}:{account_id}:hub/{hub_name}"
+
+
+def construct_hub_model_arn_from_inputs(hub_arn: str, model_name: str, version: str) -> str:
+    """Constructs a HubContent model arn from the Hub name, model name, and model version."""
+
+    info = get_info_from_hub_resource_arn(hub_arn)
+    arn = (
+        f"arn:{info.partition}:sagemaker:{info.region}:{info.account_id}:hub-content/"
+        f"{info.hub_name}/{HubContentType.MODEL.value}/{model_name}/{version}"
+    )
+
+    return arn
+
+
+def construct_hub_model_reference_arn_from_inputs(
+    hub_arn: str, model_name: str, version: str
+) -> str:
+    """Constructs a HubContent model arn from the Hub name, model name, and model version."""
+
+    info = get_info_from_hub_resource_arn(hub_arn)
+    arn = (
+        f"arn:{info.partition}:sagemaker:{info.region}:{info.account_id}:hub-content/"
+        f"{info.hub_name}/{HubContentType.MODEL_REFERENCE.value}/{model_name}/{version}"
+    )
+
+    return arn
+
+
+def generate_hub_arn_for_init_kwargs(
+    hub_name: str, region: Optional[str] = None, session: Optional[Session] = None
+):
+    """Generates the Hub Arn for JumpStart class args from a HubName or Arn.
+
+    Args:
+        hub_name (str): HubName or HubArn from JumpStart class args
+        region (str): Region from JumpStart class args
+        session (Session): Custom SageMaker Session from JumpStart class args
+    """
+
+    hub_arn = None
+    if hub_name:
+        if hub_name == constants.JUMPSTART_MODEL_HUB_NAME:
+            return None
+        match = re.match(constants.HUB_ARN_REGEX, hub_name)
+        if match:
+            hub_arn = hub_name
+        else:
+            hub_arn = construct_hub_arn_from_name(hub_name=hub_name, region=region, session=session)
+    return hub_arn
+
+
+def is_gated_bucket(bucket_name: str) -> bool:
+    """Returns true if the bucket name is the JumpStart gated bucket."""
+    return bucket_name in constants.JUMPSTART_GATED_BUCKET_NAME_SET
+
+
+def get_hub_model_version(
+    hub_name: str,
+    hub_model_name: str,
+    hub_model_type: str,
+    hub_model_version: Optional[str] = None,
+    sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+) -> str:
+    """Returns available Jumpstart hub model version.
+
+    It will attempt both a semantic HubContent version search and Marketplace version search.
+    If the Marketplace version is also semantic, this function will default to HubContent version.
+
+    Raises:
+        ClientError: If the specified model is not found in the hub.
+        KeyError: If the specified model version is not found.
+    """
+    if sagemaker_session is None:
+        # sagemaker_session is overridden to none by some callers
+        sagemaker_session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+
+    try:
+        hub_content_summaries = _list_hub_content_versions_helper(
+            hub_name=hub_name,
+            hub_content_name=hub_model_name,
+            hub_content_type=hub_model_type,
+            sagemaker_session=sagemaker_session,
+        )
+    except Exception as ex:
+        raise Exception(f"Failed calling list_hub_content_versions: {str(ex)}")
+
+    try:
+        return _get_hub_model_version_for_open_weight_version(
+            hub_content_summaries, hub_model_version
+        )
+    except KeyError:
+        marketplace_hub_content_version = _get_hub_model_version_for_marketplace_version(
+            hub_content_summaries, hub_model_version
+        )
+        if marketplace_hub_content_version:
+            return marketplace_hub_content_version
+        raise
+
+
+def _list_hub_content_versions_helper(
+    hub_name, hub_content_name, hub_content_type, sagemaker_session
+):
+    all_hub_content_summaries = []
+    list_hub_content_versions_response = sagemaker_session.list_hub_content_versions(
+        hub_name=hub_name, hub_content_name=hub_content_name, hub_content_type=hub_content_type
+    )
+    all_hub_content_summaries.extend(list_hub_content_versions_response.get("HubContentSummaries"))
+    while "NextToken" in list_hub_content_versions_response:
+        list_hub_content_versions_response = sagemaker_session.list_hub_content_versions(
+            hub_name=hub_name,
+            hub_content_name=hub_content_name,
+            hub_content_type=hub_content_type,
+            next_token=list_hub_content_versions_response["NextToken"],
+        )
+        all_hub_content_summaries.extend(
+            list_hub_content_versions_response.get("HubContentSummaries")
+        )
+    return all_hub_content_summaries
+
+
+def _get_hub_model_version_for_open_weight_version(
+    hub_content_summaries: List[Any], hub_model_version: Optional[str] = None
+) -> str:
+    available_model_versions = [model.get("HubContentVersion") for model in hub_content_summaries]
+
+    if hub_model_version == "*" or hub_model_version is None:
+        return str(max(version.parse(v) for v in available_model_versions))
+
+    try:
+        spec = SpecifierSet(f"=={hub_model_version}")
+    except InvalidSpecifier:
+        raise KeyError(f"Bad semantic version: {hub_model_version}")
+    available_versions_filtered = list(spec.filter(available_model_versions))
+    if not available_versions_filtered:
+        raise KeyError("Model version not available in the Hub")
+    hub_model_version = str(max(available_versions_filtered))
+
+    return hub_model_version
+
+
+def _get_hub_model_version_for_marketplace_version(
+    hub_content_summaries: List[Any], marketplace_version: str
+) -> Optional[str]:
+    """Returns the HubContent version associated with the Marketplace version.
+
+    This function will check within the HubContentSearchKeywords for the proprietary version.
+    """
+    for model in hub_content_summaries:
+        model_search_keywords = model.get("HubContentSearchKeywords", [])
+        if _hub_search_keywords_contains_marketplace_version(
+            model_search_keywords, marketplace_version
+        ):
+            return model.get("HubContentVersion")
+
+    return None
+
+
+def _hub_search_keywords_contains_marketplace_version(
+    model_search_keywords: List[str], marketplace_version: str
+) -> bool:
+    proprietary_version_keyword = next(
+        filter(lambda s: s.startswith(PROPRIETARY_VERSION_KEYWORD), model_search_keywords), None
+    )
+
+    if not proprietary_version_keyword:
+        return False
+
+    proprietary_version = proprietary_version_keyword.lstrip(PROPRIETARY_VERSION_KEYWORD)
+    if proprietary_version == marketplace_version:
+        return True
+
+    return False