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

sagemaker/core/telemetry/__init__.py

@@ -0,0 +1,23 @@
+# 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 telemetry utilities for the SageMaker Python SDK."""
+from __future__ import absolute_import
+
+from sagemaker.core.telemetry.constants import Feature, Status  # noqa: F401
+from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter  # noqa: F401
+
+__all__ = [
+    "Feature",
+    "Status",
+    "_telemetry_emitter",
+]

sagemaker/core/telemetry/constants.py

@@ -0,0 +1,84 @@
+# 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.
+"""Constants used in SageMaker Python SDK telemetry."""
+
+from __future__ import absolute_import
+from enum import Enum
+
+# Default AWS region used by SageMaker
+DEFAULT_AWS_REGION = "us-west-2"
+
+
+class Feature(Enum):
+    """Enumeration of feature names used in telemetry."""
+
+    SDK_DEFAULTS = 1
+    LOCAL_MODE = 2
+    REMOTE_FUNCTION = 3
+    MODEL_TRAINER = 4
+    ESTIMATOR = 5
+    HYPERPOD = 6  # Added to support telemetry in sagemaker-hyperpod-cli
+    HYPERPOD_CLI = 7  # Added to support telemetry in sagemaker-hyperpod-cli
+    MODEL_CUSTOMIZATION = 8
+
+    def __str__(self):  # pylint: disable=E0307
+        """Return the feature name."""
+        return self.name
+
+
+class Status(Enum):
+    """Enumeration of status values used in telemetry."""
+
+    SUCCESS = 1
+    FAILURE = 0
+
+    def __str__(self):  # pylint: disable=E0307
+        """Return the status name."""
+        return self.name
+
+
+class Region(str, Enum):
+    """Telemetry: List of all supported AWS regions."""
+
+    # Classic
+    US_EAST_1 = "us-east-1"  # IAD
+    US_EAST_2 = "us-east-2"  # CMH
+    US_WEST_1 = "us-west-1"  # SFO
+    US_WEST_2 = "us-west-2"  # PDX
+    AP_NORTHEAST_1 = "ap-northeast-1"  # NRT
+    AP_NORTHEAST_2 = "ap-northeast-2"  # ICN
+    AP_NORTHEAST_3 = "ap-northeast-3"  # KIX
+    AP_SOUTH_1 = "ap-south-1"  # BOM
+    AP_SOUTHEAST_1 = "ap-southeast-1"  # SIN
+    AP_SOUTHEAST_2 = "ap-southeast-2"  # SYD
+    CA_CENTRAL_1 = "ca-central-1"  # YUL
+    EU_CENTRAL_1 = "eu-central-1"  # FRA
+    EU_NORTH_1 = "eu-north-1"  # ARN
+    EU_WEST_1 = "eu-west-1"  # DUB
+    EU_WEST_2 = "eu-west-2"  # LHR
+    EU_WEST_3 = "eu-west-3"  # CDG
+    SA_EAST_1 = "sa-east-1"  # GRU
+    # Opt-in
+    AP_EAST_1 = "ap-east-1"  # HKG
+    AP_SOUTHEAST_3 = "ap-southeast-3"  # CGK
+    AF_SOUTH_1 = "af-south-1"  # CPT
+    EU_SOUTH_1 = "eu-south-1"  # MXP
+    ME_SOUTH_1 = "me-south-1"  # BAH
+    MX_CENTRAL_1 = "mx-central-1"  # QRO
+    AP_SOUTHEAST_7 = "ap-southeast-7"  # BKK
+    AP_SOUTH_2 = "ap-south-2"  # HYD
+    AP_SOUTHEAST_4 = "ap-southeast-4"  # MEL
+    EU_CENTRAL_2 = "eu-central-2"  # ZRH
+    EU_SOUTH_2 = "eu-south-2"  # ZAZ
+    IL_CENTRAL_1 = "il-central-1"  # TLV
+    ME_CENTRAL_1 = "me-central-1"  # DXB

sagemaker/core/telemetry/telemetry_logging.py

@@ -0,0 +1,284 @@
+# 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.
+"""Telemetry module for SageMaker Python SDK to collect usage data and metrics."""
+from __future__ import absolute_import
+import logging
+import platform
+import sys
+from time import perf_counter
+from typing import List
+import functools
+import requests
+
+import boto3
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.common_utils import resolve_value_from_config
+from sagemaker.core.config.config_schema import TELEMETRY_OPT_OUT_PATH
+from sagemaker.core.telemetry.constants import (
+    Feature,
+    Status,
+    Region,
+    DEFAULT_AWS_REGION,
+)
+from sagemaker.core.user_agent import SDK_VERSION, process_studio_metadata_file
+
+logger = logging.getLogger(__name__)
+
+OS_NAME = platform.system() or "UnresolvedOS"
+OS_VERSION = platform.release() or "UnresolvedOSVersion"
+OS_NAME_VERSION = "{}/{}".format(OS_NAME, OS_VERSION)
+PYTHON_VERSION = "{}.{}.{}".format(
+    sys.version_info.major, sys.version_info.minor, sys.version_info.micro
+)
+
+TELEMETRY_OPT_OUT_MESSAGING = (
+    "SageMaker Python SDK will collect telemetry to help us better understand our user's needs, "
+    "diagnose issues, and deliver additional features.\n"
+    "To opt out of telemetry, please disable via TelemetryOptOut parameter in SDK defaults config. "
+    "For more information, refer to https://sagemaker.readthedocs.io/en/stable/overview.html"
+    "#configuring-and-using-defaults-with-the-sagemaker-python-sdk."
+)
+
+FEATURE_TO_CODE = {
+    str(Feature.SDK_DEFAULTS): 1,
+    str(Feature.LOCAL_MODE): 2,
+    str(Feature.REMOTE_FUNCTION): 3,
+    str(Feature.MODEL_TRAINER): 4,
+    str(Feature.MODEL_CUSTOMIZATION): 8,
+}
+
+STATUS_TO_CODE = {
+    str(Status.SUCCESS): 1,
+    str(Status.FAILURE): 0,
+}
+
+
+def _telemetry_emitter(feature: str, func_name: str):
+    """Telemetry Emitter
+
+    Decorator to emit telemetry logs for SageMaker Python SDK functions. This class needs
+    sagemaker_session object as a member. Default session object is a pysdk v2 Session object
+    in this repo. When collecting telemetry for classes using sagemaker-core Session object,
+    we should be aware of its differences, such as sagemaker_session.sagemaker_config does not
+    exist in new Session class.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            sagemaker_session = None
+            if len(args) > 0 and hasattr(args[0], "sagemaker_session"):
+                # Get the sagemaker_session from the instance method args
+                sagemaker_session = args[0].sagemaker_session
+            elif feature == Feature.REMOTE_FUNCTION:
+                # Get the sagemaker_session from the function keyword arguments for remote function
+                sagemaker_session = kwargs.get(
+                    "sagemaker_session", _get_default_sagemaker_session()
+                )
+
+            if sagemaker_session:
+                logger.debug("sagemaker_session found, preparing to emit telemetry...")
+                logger.info(TELEMETRY_OPT_OUT_MESSAGING)
+                response = None
+                caught_ex = None
+                studio_app_type = process_studio_metadata_file()
+
+                # Check if telemetry is opted out
+                telemetry_opt_out_flag = resolve_value_from_config(
+                    direct_input=None,
+                    config_path=TELEMETRY_OPT_OUT_PATH,
+                    default_value=False,
+                    sagemaker_session=sagemaker_session,
+                )
+                logger.debug("TelemetryOptOut flag is set to: %s", telemetry_opt_out_flag)
+
+                # Construct the feature list to track feature combinations
+                feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]]
+
+                if (
+                    hasattr(sagemaker_session, "sagemaker_config")
+                    and sagemaker_session.sagemaker_config
+                    and feature != Feature.SDK_DEFAULTS
+                ):
+                    feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS)])
+
+                if (
+                    hasattr(sagemaker_session, "local_mode")
+                    and sagemaker_session.local_mode
+                    and feature != Feature.LOCAL_MODE
+                ):
+                    feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE)])
+
+                # Construct the extra info to track platform and environment usage metadata
+                extra = (
+                    f"{func_name}"
+                    f"&x-sdkVersion={SDK_VERSION}"
+                    f"&x-env={PYTHON_VERSION}"
+                    f"&x-sys={OS_NAME_VERSION}"
+                    f"&x-platform={studio_app_type}"
+                )
+
+                # Add endpoint ARN to the extra info if available
+                if hasattr(sagemaker_session, "endpoint_arn") and sagemaker_session.endpoint_arn:
+                    extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}"
+
+                start_timer = perf_counter()
+                try:
+                    # Call the original function
+                    response = func(*args, **kwargs)
+                    stop_timer = perf_counter()
+                    elapsed = stop_timer - start_timer
+                    extra += f"&x-latency={round(elapsed, 2)}"
+                    if not telemetry_opt_out_flag:
+                        _send_telemetry_request(
+                            STATUS_TO_CODE[str(Status.SUCCESS)],
+                            feature_list,
+                            sagemaker_session,
+                            None,
+                            None,
+                            extra,
+                        )
+                except Exception as e:  # pylint: disable=W0703
+                    stop_timer = perf_counter()
+                    elapsed = stop_timer - start_timer
+                    extra += f"&x-latency={round(elapsed, 2)}"
+                    if not telemetry_opt_out_flag:
+                        _send_telemetry_request(
+                            STATUS_TO_CODE[str(Status.FAILURE)],
+                            feature_list,
+                            sagemaker_session,
+                            str(e),
+                            e.__class__.__name__,
+                            extra,
+                        )
+                    caught_ex = e
+                finally:
+                    if caught_ex:
+                        raise caught_ex
+                    return response  # pylint: disable=W0150
+            else:
+                logger.debug(
+                    "Unable to send telemetry for function %s. "
+                    "sagemaker_session is not provided or not valid.",
+                    func_name,
+                )
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def _send_telemetry_request(
+    status: int,
+    feature_list: List[int],
+    session: Session,
+    failure_reason: str = None,
+    failure_type: str = None,
+    extra_info: str = None,
+) -> None:
+    """Make GET request to an empty object in S3 bucket"""
+    try:
+        accountId = _get_accountId(session) if session else "NotAvailable"
+        region = _get_region_or_default(session)
+
+        try:
+            Region(region)  # Validate the region
+        except ValueError:
+            logger.warning(
+                "Region not found in supported regions. Telemetry request will not be emitted."
+            )
+            return
+
+        url = _construct_url(
+            accountId,
+            region,
+            str(status),
+            str(
+                ",".join(map(str, feature_list))
+            ),  # Remove brackets and quotes to cut down on length
+            failure_reason,
+            failure_type,
+            extra_info,
+        )
+        # Send the telemetry request
+        logger.debug("Sending telemetry request to [%s]", url)
+        _requests_helper(url, 2)
+        logger.debug("SageMaker Python SDK telemetry successfully emitted.")
+    except Exception:  # pylint: disable=W0703
+        logger.debug("SageMaker Python SDK telemetry not emitted!")
+
+
+def _construct_url(
+    accountId: str,
+    region: str,
+    status: str,
+    feature: str,
+    failure_reason: str,
+    failure_type: str,
+    extra_info: str,
+) -> str:
+    """Construct the URL for the telemetry request"""
+
+    base_url = (
+        f"https://sm-pysdk-t-{region}.s3.{region}.amazonaws.com/telemetry?"
+        f"x-accountId={accountId}"
+        f"&x-status={status}"
+        f"&x-feature={feature}"
+    )
+    logger.debug("Failure reason: %s", failure_reason)
+    if failure_reason:
+        base_url += f"&x-failureReason={failure_reason}"
+        base_url += f"&x-failureType={failure_type}"
+    if extra_info:
+        base_url += f"&x-extra={extra_info}"
+    return base_url
+
+
+def _requests_helper(url, timeout):
+    """Make a GET request to the given URL"""
+
+    response = None
+    try:
+        response = requests.get(url, timeout)
+    except requests.exceptions.RequestException as e:
+        logger.exception("Request exception: %s", str(e))
+    return response
+
+
+def _get_accountId(session):
+    """Return the account ID from the boto session"""
+
+    try:
+        sts = session.boto_session.client("sts")
+        return sts.get_caller_identity()["Account"]
+    except Exception:  # pylint: disable=W0703
+        return None
+
+
+def _get_region_or_default(session):
+    """Return the region name from the boto session or default to us-west-2"""
+
+    try:
+        return session.boto_session.region_name
+    except Exception:  # pylint: disable=W0703
+        return DEFAULT_AWS_REGION
+
+
+def _get_default_sagemaker_session():
+    """Return the default sagemaker session"""
+
+    boto_session = boto3.Session(region_name=DEFAULT_AWS_REGION)
+    sagemaker_session = Session(boto_session=boto_session)
+
+    return sagemaker_session

sagemaker/core/tools/__init__.py

@@ -0,0 +1 @@
+from sagemaker.core.utils.code_injection.codec import pascal_to_snake

{sagemaker_core → sagemaker/core}/tools/codegen.py

@@ -11,12 +11,12 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Generates the code for the service model."""
-from sagemaker_core.main.utils import reformat_file_with_black
-from sagemaker_core.tools.shapes_codegen import ShapesCodeGen
-from sagemaker_core.tools.resources_codegen import ResourcesCodeGen
+from sagemaker.core.utils.utils import reformat_file_with_black
+from sagemaker.core.tools.shapes_codegen import ShapesCodeGen
+from sagemaker.core.tools.resources_codegen import ResourcesCodeGen
 from typing import Optional
 
-from sagemaker_core.tools.data_extractor import ServiceJsonData, load_service_jsons
+from sagemaker.core.tools.data_extractor import ServiceJsonData, load_service_jsons
 
 
 def generate_code(

{sagemaker_core → sagemaker/core}/tools/constants.py

@@ -13,9 +13,9 @@
 """Constants used in the code_generator modules."""
 import os
 
-CLASS_METHODS = set(["create", "add", "start", "register", "import", "list", "get"])
+CLASS_METHODS = set(["create", "add", "register", "import", "list", "get"])
 OBJECT_METHODS = set(
-    ["refresh", "delete", "update", "stop", "deregister", "wait", "wait_for_status"]
+    ["refresh", "delete", "update", "start", "stop", "deregister", "wait", "wait_for_status"]
 )
 
 TERMINAL_STATES = set(["Completed", "Stopped", "Deleted", "Failed", "Succeeded", "Cancelled"])
@@ -32,7 +32,7 @@ CONFIGURABLE_ATTRIBUTE_SUBSTRINGS = [
 ]
 
 BASIC_JSON_TYPES_TO_PYTHON_TYPES = {
-    "string": "str",
+    "string": "StrPipeVar",
     "integer": "int",
     "boolean": "bool",
     "long": "int",
@@ -46,9 +46,10 @@ BASIC_JSON_TYPES_TO_PYTHON_TYPES = {
 
 BASIC_RETURN_TYPES = {"str", "int", "bool", "float", "datetime.datetime"}
 
-SHAPE_DAG_FILE_PATH = os.getcwd() + "/src/sagemaker_core/main/code_injection/shape_dag.py"
+SHAPE_DAG_FILE_PATH = os.getcwd() + "/src/sagemaker/core/utils/code_injection/shape_dag.py"
 PYTHON_TYPES_TO_BASIC_JSON_TYPES = {
     "str": "string",
+    "StrPipeVar": "string",
     "int": "integer",
     "bool": "boolean",
     "float": "double",
@@ -77,18 +78,27 @@ logger = get_textual_rich_logger(__name__)
 
 # TODO: The file name should be injected, we should update it to be more generic
 ADDITIONAL_OPERATION_FILE_PATH = (
-    os.getcwd() + "/src/sagemaker_core/tools/additional_operations.json"
+    os.getcwd() + "/src/sagemaker/core/tools/additional_operations.json"
 )
-SERVICE_JSON_FILE_PATH = os.getcwd() + "/sample/sagemaker/2017-07-24/service-2.json"
-RUNTIME_SERVICE_JSON_FILE_PATH = os.getcwd() + "/sample/sagemaker-runtime/2017-05-13/service-2.json"
-FEATURE_STORE_SERVICE_JSON_FILE_PATH = (
-    os.getcwd() + "/sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json"
+
+# Get the package root directory (sagemaker-core/)
+# When installed, __file__ points to site-packages, so we need to find the actual source/sample location
+_PACKAGE_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
+
+SERVICE_JSON_FILE_PATH = os.path.join(_PACKAGE_ROOT, "sample/sagemaker/2017-07-24/service-2.json")
+RUNTIME_SERVICE_JSON_FILE_PATH = os.path.join(
+    _PACKAGE_ROOT, "sample/sagemaker-runtime/2017-05-13/service-2.json"
+)
+FEATURE_STORE_SERVICE_JSON_FILE_PATH = os.path.join(
+    _PACKAGE_ROOT, "sample/sagemaker-featurestore-runtime/2020-07-01/service-2.json"
+)
+METRICS_SERVICE_JSON_FILE_PATH = os.path.join(
+    _PACKAGE_ROOT, "sample/sagemaker-metrics/2022-09-30/service-2.json"
 )
-METRICS_SERVICE_JSON_FILE_PATH = os.getcwd() + "/sample/sagemaker-metrics/2022-09-30/service-2.json"
 
-GENERATED_CLASSES_LOCATION = os.getcwd() + "/src/sagemaker_core/main"
+GENERATED_CLASSES_LOCATION = os.getcwd() + "/src/sagemaker/core"
 UTILS_CODEGEN_FILE_NAME = "utils.py"
-DEFAULT_CONFIGS_CODEGEN_FILE_NAME = "default_configs_helper.py"
+INTELLIGENT_DEFAULTS_HELPER_CODEGEN_FILE_NAME = "intelligent_defaults_helper.py"
 
 RESOURCES_CODEGEN_FILE_NAME = "resources.py"
 
@@ -96,6 +106,4 @@ SHAPES_CODEGEN_FILE_NAME = "shapes.py"
 
 CONFIG_SCHEMA_FILE_NAME = "config_schema.py"
 
-API_COVERAGE_JSON_FILE_PATH = os.getcwd() + "/src/sagemaker_core/tools/api_coverage.json"
-
-SHAPES_WITH_JSON_FIELD_ALIAS = ["MonitoringDatasetFormat"]  # Shapes with field name with "json"
+API_COVERAGE_JSON_FILE_PATH = os.getcwd() + "/src/sagemaker/core/tools/api_coverage.json"

{sagemaker_core → sagemaker/core}/tools/data_extractor.py

@@ -3,7 +3,7 @@ from functools import lru_cache
 
 from pydantic import BaseModel
 
-from sagemaker_core.tools.constants import (
+from sagemaker.core.tools.constants import (
     ADDITIONAL_OPERATION_FILE_PATH,
     FEATURE_STORE_SERVICE_JSON_FILE_PATH,
     METRICS_SERVICE_JSON_FILE_PATH,

{sagemaker_core → sagemaker/core}/tools/method.py

@@ -1,6 +1,6 @@
 from enum import Enum
 
-from sagemaker_core.main.utils import remove_html_tags
+from sagemaker.core.utils.utils import remove_html_tags
 
 
 class MethodType(Enum):