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

sagemaker/core/__init__.py

@@ -0,0 +1,16 @@
+from sagemaker.core.utils.utils import enable_textual_rich_console_and_traceback
+
+
+enable_textual_rich_console_and_traceback()
+
+# Job management
+from sagemaker.core.job import _Job  # noqa: F401
+from sagemaker.core.processing import (  # noqa: F401
+    Processor,
+    ScriptProcessor,
+    FrameworkProcessor,
+)
+from sagemaker.core.transformer import Transformer  # noqa: F401
+
+# Note: HyperparameterTuner and WarmStartTypes are in sagemaker.train.tuner
+# They are not re-exported from core to avoid circular dependencies

sagemaker/core/_studio.py

@@ -0,0 +1,116 @@
+# 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.
+"""Provides internal tooling for studio environments."""
+from __future__ import absolute_import
+
+import json
+import logging
+
+from pathlib import Path
+
+STUDIO_PROJECT_CONFIG = ".sagemaker-code-config"
+
+logger = logging.getLogger(__name__)
+
+
+def _append_project_tags(tags=None, working_dir=None):
+    """Appends the project tag to the list of tags, if it exists.
+
+    Args:
+        working_dir: the working directory to start looking.
+        tags: the list of tags to append to.
+
+    Returns:
+        A possibly extended list of tags that includes the project id.
+    """
+    path = _find_config(working_dir)
+    if path is None:
+        return tags
+
+    config = _load_config(path)
+    if config is None:
+        return tags
+
+    additional_tags = _parse_tags(config)
+    if additional_tags is None:
+        return tags
+
+    all_tags = tags or []
+    additional_tags = [tag for tag in additional_tags if tag not in all_tags]
+    all_tags.extend(additional_tags)
+
+    return all_tags
+
+
+def _find_config(working_dir=None):
+    """Gets project config on SageMaker Studio platforms, if it exists.
+
+    Args:
+        working_dir: the working directory to start looking.
+
+    Returns:
+        The project config path, if it exists. Otherwise None.
+    """
+    try:
+        wd = Path(working_dir) if working_dir else Path.cwd()
+
+        path = None
+
+        # Get the root of the current working directory for both Windows and Unix-like systems
+        root = Path(wd.anchor)
+        while path is None and wd != root:
+            candidate = wd / STUDIO_PROJECT_CONFIG
+            if Path.exists(candidate):
+                path = candidate
+            wd = wd.parent
+
+        return path
+    except Exception as e:  # pylint: disable=W0703
+        logger.debug("Could not find the studio project config. %s", e)
+
+
+def _load_config(path):
+    """Parse out the projectId attribute if it exists at path.
+
+    Args:
+        path: path to project config
+
+    Returns:
+        Project config Json, or None if it does not exist.
+    """
+    try:
+        with open(path, "r") as f:
+            content = f.read().strip()
+        config = json.loads(content)
+
+        return config
+    except Exception as e:  # pylint: disable=W0703
+        logger.debug("Could not load project config. %s", e)
+
+
+def _parse_tags(config):
+    """Parse out appropriate attributes and formats as tags.
+
+    Args:
+        config: project config dict
+
+    Returns:
+        List of tags
+    """
+    try:
+        return [
+            {"Key": "sagemaker:project-id", "Value": config["sagemakerProjectId"]},
+            {"Key": "sagemaker:project-name", "Value": config["sagemakerProjectName"]},
+        ]
+    except Exception as e:  # pylint: disable=W0703
+        logger.debug("Could not parse project config. %s", e)

sagemaker/core/_version.py

@@ -0,0 +1,11 @@
+import os
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+
+# Get the root directory of the project
+root_dir = os.path.abspath(os.path.join(script_dir, "..", "..", ".."))
+
+version_file_path = os.path.join(root_dir, "VERSION")
+
+with open(version_file_path) as version_file:
+    __version__ = version_file.read().strip()

sagemaker/core/accept_types.py

@@ -0,0 +1,131 @@
+# 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 is for SageMaker accept types."""
+from __future__ import absolute_import
+from typing import List, Optional
+
+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_options(
+    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,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+) -> List[str]:
+    """Retrieves the supported accept types for the model matching the given arguments.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the supported accept types.
+            Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the supported accept types. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            supported accept types. (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).
+        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).
+    Returns:
+        list: The supported accept types 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 accept types."
+        )
+
+    return artifacts._retrieve_supported_accept_types(
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_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,
+    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.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the default accept type.
+            Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the default accept type. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            default accept type. (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).
+        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 default accept type 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 accept types."
+        )
+
+    return artifacts._retrieve_default_accept_type(
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        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,
+    )