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

sagemaker/core/jumpstart/artifacts/model_packages.py

@@ -0,0 +1,202 @@
+# 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 functions for obtaining JumpStart model packages."""
+from __future__ import absolute_import
+from typing import Optional
+from sagemaker.core.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+)
+from sagemaker.core.jumpstart.utils import (
+    get_region_fallback,
+    verify_model_region_and_return_specs,
+)
+from sagemaker.core.jumpstart.enums import (
+    JumpStartScriptScope,
+    JumpStartModelType,
+)
+from sagemaker.core.helper.session_helper import Session
+
+
+def _retrieve_model_package_arn(
+    model_id: str,
+    model_version: str,
+    instance_type: Optional[str],
+    region: Optional[str],
+    hub_arn: Optional[str] = None,
+    scope: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    config_name: Optional[str] = None,
+) -> Optional[str]:
+    """Retrieves associated model pacakge arn for the model.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the model package arn.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            model package arn.
+        instance_type (Optional[str]): An instance type to optionally supply in order to get an arn
+            specific for the instance type.
+        region (Optional[str]): Region for which to retrieve the model package arn.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
+        scope (Optional[str]): Scope for which to retrieve the model package arn.
+        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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+
+    Returns:
+        str: the model package arn to use for the model or None.
+    """
+
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        hub_arn=hub_arn,
+        scope=scope,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+        model_type=model_type,
+        config_name=config_name,
+    )
+
+    if scope == JumpStartScriptScope.INFERENCE:
+
+        instance_specific_arn: Optional[str] = (
+            model_specs.hosting_instance_type_variants.get_model_package_arn(
+                region=region, instance_type=instance_type
+            )
+            if getattr(model_specs, "hosting_instance_type_variants", None) is not None
+            else None
+        )
+
+        if instance_specific_arn is not None:
+            return instance_specific_arn
+
+        if (
+            model_specs.hosting_model_package_arns is None
+            or model_specs.hosting_model_package_arns == {}
+        ):
+            return None
+
+        regional_arn = model_specs.hosting_model_package_arns.get(region)
+
+        if regional_arn is None:
+            raise ValueError(
+                f"Model package arn for '{model_id}' not supported in {region}. "
+                "Please try one of the following regions: "
+                f"{', '.join(model_specs.hosting_model_package_arns.keys())}."
+            )
+
+        return regional_arn
+
+    raise NotImplementedError(f"Model Package ARN not supported for scope: '{scope}'")
+
+
+def _retrieve_model_package_model_artifact_s3_uri(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    hub_arn: Optional[str] = None,
+    scope: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    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.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the model package artifact.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            model package artifact.
+        region (Optional[str]): Region for which to retrieve the model package artifact.
+            (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
+        scope (Optional[str]): Scope for which to retrieve the model package artifact.
+            (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).
+        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.
+
+    Raises:
+        NotImplementedError: If an unsupported script is used.
+    """
+
+    if scope == JumpStartScriptScope.TRAINING:
+
+        region = region or get_region_fallback(
+            sagemaker_session=sagemaker_session,
+        )
+
+        model_specs = verify_model_region_and_return_specs(
+            model_id=model_id,
+            version=model_version,
+            hub_arn=hub_arn,
+            scope=scope,
+            region=region,
+            tolerate_vulnerable_model=tolerate_vulnerable_model,
+            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:
+            return None
+
+        model_s3_uri = model_specs.training_model_package_artifact_uris.get(region)
+
+        if model_s3_uri is None:
+            raise ValueError(
+                f"Model package artifact s3 uri for '{model_id}' not supported in {region}. "
+                "Please try one of the following regions: "
+                f"{', '.join(model_specs.training_model_package_artifact_uris.keys())}."
+            )
+
+        return model_s3_uri
+
+    raise NotImplementedError(f"Model Package Artifact URI not supported for scope: '{scope}'")

sagemaker/core/jumpstart/artifacts/model_uris.py

@@ -0,0 +1,252 @@
+# 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 functions for obtaining JumpStart model uris."""
+from __future__ import absolute_import
+import os
+from typing import Optional
+
+from sagemaker.core.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE,
+)
+from sagemaker.core.jumpstart.enums import (
+    JumpStartModelType,
+    JumpStartScriptScope,
+)
+from sagemaker.core.jumpstart.utils import (
+    get_jumpstart_content_bucket,
+    get_jumpstart_gated_content_bucket,
+    get_region_fallback,
+    verify_model_region_and_return_specs,
+)
+from sagemaker.core.s3.utils import is_s3_url
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.jumpstart.types import JumpStartModelSpecs
+
+
+def _retrieve_hosting_prepacked_artifact_key(
+    model_specs: JumpStartModelSpecs, instance_type: str
+) -> str:
+    """Returns instance specific hosting prepacked artifact key or default one as fallback."""
+    instance_specific_prepacked_hosting_artifact_key: Optional[str] = (
+        model_specs.hosting_instance_type_variants.get_instance_specific_prepacked_artifact_key(
+            instance_type=instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "hosting_instance_type_variants", None) is not None
+        else None
+    )
+
+    default_prepacked_hosting_artifact_key: Optional[str] = getattr(
+        model_specs, "hosting_prepacked_artifact_key"
+    )
+
+    return (
+        instance_specific_prepacked_hosting_artifact_key or default_prepacked_hosting_artifact_key
+    )
+
+
+def _retrieve_hosting_artifact_key(model_specs: JumpStartModelSpecs, instance_type: str) -> str:
+    """Returns instance specific hosting artifact key or default one as fallback."""
+    instance_specific_hosting_artifact_key: Optional[str] = (
+        model_specs.hosting_instance_type_variants.get_instance_specific_artifact_key(
+            instance_type=instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "hosting_instance_type_variants", None) is not None
+        else None
+    )
+
+    default_hosting_artifact_key: str = model_specs.hosting_artifact_key
+
+    return instance_specific_hosting_artifact_key or default_hosting_artifact_key
+
+
+def _retrieve_training_artifact_key(model_specs: JumpStartModelSpecs, instance_type: str) -> str:
+    """Returns instance specific training artifact key or default one as fallback."""
+    instance_specific_training_artifact_key: Optional[str] = (
+        model_specs.training_instance_type_variants.get_instance_specific_training_artifact_key(
+            instance_type=instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "training_instance_type_variants", None) is not None
+        else None
+    )
+
+    default_training_artifact_key: str = model_specs.training_artifact_key
+
+    return instance_specific_training_artifact_key or default_training_artifact_key
+
+
+def _retrieve_model_uri(
+    model_id: str,
+    model_version: str,
+    hub_arn: Optional[str] = None,
+    model_scope: Optional[str] = None,
+    instance_type: Optional[str] = None,
+    region: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    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 model artifact S3 URI for the model matching the given arguments.
+
+    Optionally uses a bucket override specified by environment variable.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to retrieve
+            the model artifact S3 URI.
+        model_version (str): Version of the JumpStart model for which to retrieve the model
+            artifact S3 URI.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
+        model_scope (str): The model type, i.e. what it is used for.
+            Valid values: "training" and "inference".
+        instance_type (str): The ML compute instance type for the specified scope. (Default: None).
+        region (str): Region for which to retrieve model S3 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).
+        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 artifact S3 URI for the corresponding model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+        VulnerableJumpStartModelError: If any of the dependencies required by the script have
+            known security vulnerabilities.
+        DeprecatedJumpStartModelError: If the version of the model is deprecated.
+    """
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        hub_arn=hub_arn,
+        scope=model_scope,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+        config_name=config_name,
+        model_type=model_type,
+    )
+
+    model_artifact_key: str
+
+    if model_scope == JumpStartScriptScope.INFERENCE:
+
+        is_prepacked = not model_specs.use_inference_script_uri()
+
+        if hub_arn:
+            model_artifact_uri = model_specs.hosting_artifact_uri
+            return model_artifact_uri
+        model_artifact_key = (
+            _retrieve_hosting_prepacked_artifact_key(model_specs, instance_type)
+            if is_prepacked
+            else _retrieve_hosting_artifact_key(model_specs, instance_type)
+        )
+
+    elif model_scope == JumpStartScriptScope.TRAINING:
+
+        model_artifact_key = _retrieve_training_artifact_key(model_specs, instance_type)
+
+    default_jumpstart_bucket: str = (
+        get_jumpstart_gated_content_bucket(region)
+        if model_specs.gated_bucket
+        else get_jumpstart_content_bucket(region)
+    )
+
+    bucket = (
+        os.environ.get(ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE)
+        or default_jumpstart_bucket
+    )
+    if not is_s3_url(model_artifact_key):
+        model_s3_uri = f"s3://{bucket}/{model_artifact_key}"
+
+    return model_s3_uri
+
+
+def _model_supports_training_model_uri(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    hub_arn: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    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 training with model uri field.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the support status for model uri with training.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            support status for model uri with training.
+        region (Optional[str]): Region for which to retrieve the
+            support status for model uri with training.
+        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 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).
+        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 model uri with training.
+    """
+
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        hub_arn=hub_arn,
+        scope=JumpStartScriptScope.TRAINING,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+        config_name=config_name,
+        model_type=model_type,
+    )
+
+    return model_specs.use_training_model_artifact()

sagemaker/core/jumpstart/artifacts/payloads.py

@@ -0,0 +1,96 @@
+# 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 functions to obtain JumpStart model payloads."""
+from __future__ import absolute_import
+from copy import deepcopy
+from typing import Dict, Optional
+from sagemaker.core.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+)
+from sagemaker.core.jumpstart.enums import (
+    JumpStartScriptScope,
+    JumpStartModelType,
+)
+from sagemaker.core.jumpstart.types import JumpStartSerializablePayload
+from sagemaker.core.jumpstart.utils import (
+    get_region_fallback,
+    verify_model_region_and_return_specs,
+)
+from sagemaker.core.helper.session_helper import Session
+
+
+def _retrieve_example_payloads(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    hub_arn: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    config_name: Optional[str] = None,
+) -> Optional[Dict[str, JumpStartSerializablePayload]]:
+    """Returns example payloads.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            get example payloads.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            example payloads.
+        region (Optional[str]): Region for which to retrieve the
+            example payloads.
+        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 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).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+    Returns:
+        Optional[Dict[str, JumpStartSerializablePayload]]: dictionary mapping payload aliases
+            to the serializable payload object.
+    """
+
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        hub_arn=hub_arn,
+        scope=JumpStartScriptScope.INFERENCE,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+        model_type=model_type,
+        config_name=config_name,
+    )
+
+    default_payloads = model_specs.default_payloads
+
+    if default_payloads:
+        for payload in default_payloads.values():
+            payload.accept = getattr(
+                payload, "accept", model_specs.predictor_specs.default_accept_type
+            )
+
+    return deepcopy(default_payloads) if default_payloads else None