Implement custom telemetry logging in SDK (#4721)

knikure · web-flow · commit 5e218b198eb0 · 2024-06-12T14:07:38.000-07:00
diff --git a/README.rst b/README.rst
@@ -95,6 +95,13 @@ SageMaker Python SDK is tested on:
 - Python 3.10
 - Python 3.11
 
+Telemetry
+~~~~~~~~~~~~~~~
+
+The ``sagemaker`` library has telemetry enabled to help us better understand user needs, diagnose issues, and deliver new features. This telemetry tracks the usage of various SageMaker functions.
+
+If you prefer to opt out of telemetry, you can easily do so by setting the ``TelemetryOptOut`` parameter to ``true`` in the SDK defaults configuration. For detailed instructions, please visit `Configuring and using defaults with the SageMaker Python SDK <https://sagemaker.readthedocs.io/en/stable/overview.html#configuring-and-using-defaults-with-the-sagemaker-python-sdk>`__.
+
 AWS Permissions
 ~~~~~~~~~~~~~~~
 
diff --git a/src/sagemaker/local/local_session.py b/src/sagemaker/local/local_session.py
@@ -42,6 +42,8 @@
     _LocalPipeline,
 )
 from sagemaker.session import Session
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 from sagemaker.utils import (
     get_config_value,
     _module_import_error,
@@ -83,6 +85,7 @@ def __init__(self, sagemaker_session=None):
         """
         self.sagemaker_session = sagemaker_session or LocalSession()
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_processing_job")
     def create_processing_job(
         self,
         ProcessingJobName,
@@ -165,6 +168,7 @@ def describe_processing_job(self, ProcessingJobName):
             raise ClientError(error_response, "describe_processing_job")
         return LocalSagemakerClient._processing_jobs[ProcessingJobName].describe()
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_training_job")
     def create_training_job(
         self,
         TrainingJobName,
@@ -235,6 +239,7 @@ def describe_training_job(self, TrainingJobName):
             raise ClientError(error_response, "describe_training_job")
         return LocalSagemakerClient._training_jobs[TrainingJobName].describe()
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_transform_job")
     def create_transform_job(
         self,
         TransformJobName,
@@ -280,6 +285,7 @@ def describe_transform_job(self, TransformJobName):
             raise ClientError(error_response, "describe_transform_job")
         return LocalSagemakerClient._transform_jobs[TransformJobName].describe()
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_model")
     def create_model(
         self, ModelName, PrimaryContainer, *args, **kwargs
     ):  # pylint: disable=unused-argument
@@ -329,6 +335,7 @@ def describe_endpoint_config(self, EndpointConfigName):
             raise ClientError(error_response, "describe_endpoint_config")
         return LocalSagemakerClient._endpoint_configs[EndpointConfigName].describe()
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_endpoint_config")
     def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
         """Create the endpoint configuration.
 
@@ -360,6 +367,7 @@ def describe_endpoint(self, EndpointName):
             raise ClientError(error_response, "describe_endpoint")
         return LocalSagemakerClient._endpoints[EndpointName].describe()
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_endpoint")
     def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
         """Create the endpoint.
 
@@ -428,6 +436,7 @@ def delete_model(self, ModelName):
         if ModelName in LocalSagemakerClient._models:
             del LocalSagemakerClient._models[ModelName]
 
+    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_pipeline")
     def create_pipeline(
         self, pipeline, pipeline_description, **kwargs  # pylint: disable=unused-argument
     ):
diff --git a/src/sagemaker/telemetry/__init__.py b/src/sagemaker/telemetry/__init__.py
@@ -0,0 +1,16 @@
+# 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.
+"""Placeholder docstring"""
+from __future__ import absolute_import
+
+from .telemetry_logging import _telemetry_emitter  # noqa: F401
diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py
@@ -0,0 +1,41 @@
+# 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
+
+    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
diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py
@@ -0,0 +1,229 @@
+# 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
+
+from sagemaker.utils import resolve_value_from_config
+from sagemaker.config.config_schema import TELEMETRY_OPT_OUT_PATH
+from sagemaker.telemetry.constants import (
+    Feature,
+    Status,
+    DEFAULT_AWS_REGION,
+)
+from sagemaker.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,
+}
+
+STATUS_TO_CODE = {
+    str(Status.SUCCESS): 1,
+    str(Status.FAILURE): 0,
+}
+
+
+def _telemetry_emitter(feature: str, func_name: str):
+    """Decorator to emit telemetry logs for SageMaker Python SDK functions"""
+
+    def decorator(func):
+        def wrapper(self, *args, **kwargs):
+            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=self.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 self.sagemaker_session:
+                if self.sagemaker_session.sagemaker_config and feature != Feature.SDK_DEFAULTS:
+                    feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS)])
+
+                if self.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 self.sagemaker_session and self.sagemaker_session.endpoint_arn:
+                extra += f"&x-endpointArn={self.sagemaker_session.endpoint_arn}"
+
+            start_timer = perf_counter()
+            try:
+                # Call the original function
+                response = func(self, *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,
+                        self.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,
+                        self.sagemaker_session,
+                        str(e),
+                        e.__class__.__name__,
+                        extra,
+                    )
+                caught_ex = e
+            finally:
+                if caught_ex:
+                    raise caught_ex
+                return response  # pylint: disable=W0150
+
+        return wrapper
+
+    return decorator
+
+
+from sagemaker.session import Session  # noqa: E402 pylint: disable=C0413
+
+
+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)
+        region = _get_region_or_default(session)
+        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
+
+
+import requests  # noqa: E402 pylint: disable=C0413,C0411
+
+
+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
diff --git a/tests/integ/sagemaker/conftest.py b/tests/integ/sagemaker/conftest.py
@@ -102,7 +102,9 @@
     "channels:\n"
     "  - defaults\n"
     "dependencies:\n"
-    "  - scipy=1.10.1\n"
+    "  - requests=2.32.3\n"
+    "  - charset-normalizer=3.3.2\n"
+    "  - scipy=1.13.1\n"
     "  - pip:\n"
     "    - /sagemaker-{sagemaker_version}.tar.gz\n"
     "prefix: /opt/conda/bin/conda\n"
diff --git a/tests/unit/sagemaker/local/test_local_session.py b/tests/unit/sagemaker/local/test_local_session.py
diff --git a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py
