change: Enhance telemetry logging module and feature coverage

knikure · knikure · commit 65f349665842 · 2024-06-13T20:54:42.000Z
diff --git a/src/sagemaker/remote_function/client.py b/src/sagemaker/remote_function/client.py
@@ -40,6 +40,8 @@
 from sagemaker.utils import name_from_base, base_from_name
 from sagemaker.remote_function.spark_config import SparkConfig
 from sagemaker.remote_function.custom_file_filter import CustomFileFilter
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 
 _API_CALL_LIMIT = {
     "SubmittingIntervalInSecs": 1,
@@ -57,6 +59,7 @@
 logger = logging_config.get_logger()
 
 
+@_telemetry_emitter(feature=Feature.REMOTE_FUNCTION, func_name="remote_function.remote")
 def remote(
     _func=None,
     *,
diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py
@@ -24,6 +24,7 @@ class Feature(Enum):
 
     SDK_DEFAULTS = 1
     LOCAL_MODE = 2
+    REMOTE_FUNCTION = 3
 
     def __str__(self):  # pylint: disable=E0307
         """Return the feature name."""
diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py
@@ -17,6 +17,7 @@
 import sys
 from time import perf_counter
 from typing import List
+import functools
 
 from sagemaker.utils import resolve_value_from_config
 from sagemaker.config.config_schema import TELEMETRY_OPT_OUT_PATH
@@ -47,6 +48,7 @@
 FEATURE_TO_CODE = {
     str(Feature.SDK_DEFAULTS): 1,
     str(Feature.LOCAL_MODE): 2,
+    str(Feature.REMOTE_FUNCTION): 3,
 }
 
 STATUS_TO_CODE = {
@@ -59,77 +61,95 @@ 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:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            sagemaker_session = None
+            # Check if sagemaker_session is passed as a function keyword argument
+            if "sagemaker_session" in kwargs:
+                sagemaker_session = kwargs["sagemaker_session"]
+            # Check if sagemaker_session is passed in the instance method
+            elif len(args) > 0 and hasattr(args[0], "sagemaker_session"):
+                sagemaker_session = args[0].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 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:
+                if 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
+                # 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 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
 
@@ -165,9 +185,9 @@ def _send_telemetry_request(
         # Send the telemetry request
         logger.debug("Sending telemetry request to [%s]", url)
         _requests_helper(url, 2)
-        logger.debug("SageMaker Python SDK telemetry successfully emitted!")
+        logger.debug("SageMaker Python SDK telemetry successfully emitted.")
     except Exception:  # pylint: disable=W0703
-        logger.debug("SageMaker Python SDK telemetry not emitted!!")
+        logger.debug("SageMaker Python SDK telemetry not emitted!")
 
 
 def _construct_url(
diff --git a/tests/unit/sagemaker/remote_function/test_client.py b/tests/unit/sagemaker/remote_function/test_client.py
@@ -15,6 +15,7 @@
 import os
 import threading
 import time
+import inspect
 
 import pytest
 from mock import MagicMock, patch, Mock, ANY, call
@@ -1498,7 +1499,6 @@ def test_consistency_between_remote_and_step_decorator():
     from sagemaker.workflow.function_step import step
 
     remote_args_to_ignore = [
-        "_remote",
         "include_local_workdir",
         "custom_file_filter",
         "s3_kms_key",
@@ -1508,7 +1508,7 @@ def test_consistency_between_remote_and_step_decorator():
 
     step_args_to_ignore = ["_step", "name", "display_name", "description", "retry_policies"]
 
-    remote_decorator_args = remote.__code__.co_varnames
+    remote_decorator_args = inspect.signature(remote).parameters.keys()
     common_remote_decorator_args = set(remote_args_to_ignore) ^ set(remote_decorator_args)
 
     step_decorator_args = step.__code__.co_varnames
@@ -1522,8 +1522,7 @@ def test_consistency_between_remote_and_executor():
     executor_arg_list.remove("self")
     executor_arg_list.remove("max_parallel_jobs")
 
-    remote_args_list = list(remote.__code__.co_varnames)
-    remote_args_list.remove("_remote")
+    remote_args_list = list(inspect.signature(remote).parameters.keys())
     remote_args_list.remove("_func")
 
     assert executor_arg_list == remote_args_list
