Skip to content

Commit 2c0f6cd

Browse files
knikureknikure
committed
change: Enhance telemetry logging module and feature coverage
1 parent 5e218b1 commit 2c0f6cd

File tree

4 files changed

+105
-69
lines changed

4 files changed

+105
-69
lines changed

‎src/sagemaker/remote_function/client.py

+3
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040
from sagemaker.utils import name_from_base, base_from_name
4141
from sagemaker.remote_function.spark_config import SparkConfig
4242
from sagemaker.remote_function.custom_file_filter import CustomFileFilter
43+
from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
44+
from sagemaker.telemetry.constants import Feature
4345

4446
_API_CALL_LIMIT = {
4547
"SubmittingIntervalInSecs": 1,
@@ -57,6 +59,7 @@
5759
logger = logging_config.get_logger()
5860

5961

62+
@_telemetry_emitter(feature=Feature.REMOTE_FUNCTION, func_name="remote_function.remote")
6063
def remote(
6164
_func=None,
6265
*,

‎src/sagemaker/telemetry/constants.py

+1
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ class Feature(Enum):
2424

2525
SDK_DEFAULTS = 1
2626
LOCAL_MODE = 2
27+
REMOTE_FUNCTION = 3
2728

2829
def __str__(self): # pylint: disable=E0307
2930
"""Return the feature name."""

‎src/sagemaker/telemetry/telemetry_logging.py

+89-69
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import sys
1818
from time import perf_counter
1919
from typing import List
20+
import functools
2021

2122
from sagemaker.utils import resolve_value_from_config
2223
from sagemaker.config.config_schema import TELEMETRY_OPT_OUT_PATH
@@ -47,6 +48,7 @@
4748
FEATURE_TO_CODE = {
4849
str(Feature.SDK_DEFAULTS): 1,
4950
str(Feature.LOCAL_MODE): 2,
51+
str(Feature.REMOTE_FUNCTION): 3,
5052
}
5153

5254
STATUS_TO_CODE = {
@@ -59,77 +61,95 @@ def _telemetry_emitter(feature: str, func_name: str):
5961
"""Decorator to emit telemetry logs for SageMaker Python SDK functions"""
6062

6163
def decorator(func):
62-
def wrapper(self, *args, **kwargs):
63-
logger.info(TELEMETRY_OPT_OUT_MESSAGING)
64-
response = None
65-
caught_ex = None
66-
studio_app_type = process_studio_metadata_file()
67-
68-
# Check if telemetry is opted out
69-
telemetry_opt_out_flag = resolve_value_from_config(
70-
direct_input=None,
71-
config_path=TELEMETRY_OPT_OUT_PATH,
72-
default_value=False,
73-
sagemaker_session=self.sagemaker_session,
74-
)
75-
logger.debug("TelemetryOptOut flag is set to: %s", telemetry_opt_out_flag)
76-
77-
# Construct the feature list to track feature combinations
78-
feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]]
79-
if self.sagemaker_session:
80-
if self.sagemaker_session.sagemaker_config and feature != Feature.SDK_DEFAULTS:
64+
@functools.wraps(func)
65+
def wrapper(*args, **kwargs):
66+
sagemaker_session = None
67+
# Check if sagemaker_session is passed as a function keyword argument
68+
if "sagemaker_session" in kwargs:
69+
sagemaker_session = kwargs["sagemaker_session"]
70+
# Check if sagemaker_session is passed in the instance method
71+
elif len(args) > 0 and hasattr(args[0], "sagemaker_session"):
72+
sagemaker_session = args[0].sagemaker_session
73+
74+
if sagemaker_session:
75+
logger.debug("sagemaker_session found, preparing to emit telemetry...")
76+
logger.info(TELEMETRY_OPT_OUT_MESSAGING)
77+
response = None
78+
caught_ex = None
79+
studio_app_type = process_studio_metadata_file()
80+
81+
# Check if telemetry is opted out
82+
telemetry_opt_out_flag = resolve_value_from_config(
83+
direct_input=None,
84+
config_path=TELEMETRY_OPT_OUT_PATH,
85+
default_value=False,
86+
sagemaker_session=sagemaker_session,
87+
)
88+
logger.debug("TelemetryOptOut flag is set to: %s", telemetry_opt_out_flag)
89+
90+
# Construct the feature list to track feature combinations
91+
feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]]
92+
93+
if sagemaker_session.sagemaker_config and feature != Feature.SDK_DEFAULTS:
8194
feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS)])
8295

83-
if self.sagemaker_session.local_mode and feature != Feature.LOCAL_MODE:
96+
if sagemaker_session.local_mode and feature != Feature.LOCAL_MODE:
8497
feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE)])
8598

86-
# Construct the extra info to track platform and environment usage metadata
87-
extra = (
88-
f"{func_name}"
89-
f"&x-sdkVersion={SDK_VERSION}"
90-
f"&x-env={PYTHON_VERSION}"
91-
f"&x-sys={OS_NAME_VERSION}"
92-
f"&x-platform={studio_app_type}"
93-
)
94-
95-
# Add endpoint ARN to the extra info if available
96-
if self.sagemaker_session and self.sagemaker_session.endpoint_arn:
97-
extra += f"&x-endpointArn={self.sagemaker_session.endpoint_arn}"
98-
99-
start_timer = perf_counter()
100-
try:
101-
# Call the original function
102-
response = func(self, *args, **kwargs)
103-
stop_timer = perf_counter()
104-
elapsed = stop_timer - start_timer
105-
extra += f"&x-latency={round(elapsed, 2)}"
106-
if not telemetry_opt_out_flag:
107-
_send_telemetry_request(
108-
STATUS_TO_CODE[str(Status.SUCCESS)],
109-
feature_list,
110-
self.sagemaker_session,
111-
None,
112-
None,
113-
extra,
114-
)
115-
except Exception as e: # pylint: disable=W0703
116-
stop_timer = perf_counter()
117-
elapsed = stop_timer - start_timer
118-
extra += f"&x-latency={round(elapsed, 2)}"
119-
if not telemetry_opt_out_flag:
120-
_send_telemetry_request(
121-
STATUS_TO_CODE[str(Status.FAILURE)],
122-
feature_list,
123-
self.sagemaker_session,
124-
str(e),
125-
e.__class__.__name__,
126-
extra,
127-
)
128-
caught_ex = e
129-
finally:
130-
if caught_ex:
131-
raise caught_ex
132-
return response # pylint: disable=W0150
99+
# Construct the extra info to track platform and environment usage metadata
100+
extra = (
101+
f"{func_name}"
102+
f"&x-sdkVersion={SDK_VERSION}"
103+
f"&x-env={PYTHON_VERSION}"
104+
f"&x-sys={OS_NAME_VERSION}"
105+
f"&x-platform={studio_app_type}"
106+
)
107+
108+
# Add endpoint ARN to the extra info if available
109+
if sagemaker_session.endpoint_arn:
110+
extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}"
111+
112+
start_timer = perf_counter()
113+
try:
114+
# Call the original function
115+
response = func(*args, **kwargs)
116+
stop_timer = perf_counter()
117+
elapsed = stop_timer - start_timer
118+
extra += f"&x-latency={round(elapsed, 2)}"
119+
if not telemetry_opt_out_flag:
120+
_send_telemetry_request(
121+
STATUS_TO_CODE[str(Status.SUCCESS)],
122+
feature_list,
123+
sagemaker_session,
124+
None,
125+
None,
126+
extra,
127+
)
128+
except Exception as e: # pylint: disable=W0703
129+
stop_timer = perf_counter()
130+
elapsed = stop_timer - start_timer
131+
extra += f"&x-latency={round(elapsed, 2)}"
132+
if not telemetry_opt_out_flag:
133+
_send_telemetry_request(
134+
STATUS_TO_CODE[str(Status.FAILURE)],
135+
feature_list,
136+
sagemaker_session,
137+
str(e),
138+
e.__class__.__name__,
139+
extra,
140+
)
141+
caught_ex = e
142+
finally:
143+
if caught_ex:
144+
raise caught_ex
145+
return response # pylint: disable=W0150
146+
else:
147+
logger.debug(
148+
"Unable to send telemetry for function %s. "
149+
"sagemaker_session is not provided or not valid.",
150+
func_name,
151+
)
152+
return func(*args, **kwargs)
133153

134154
return wrapper
135155

@@ -165,9 +185,9 @@ def _send_telemetry_request(
165185
# Send the telemetry request
166186
logger.debug("Sending telemetry request to [%s]", url)
167187
_requests_helper(url, 2)
168-
logger.debug("SageMaker Python SDK telemetry successfully emitted!")
188+
logger.debug("SageMaker Python SDK telemetry successfully emitted.")
169189
except Exception: # pylint: disable=W0703
170-
logger.debug("SageMaker Python SDK telemetry not emitted!!")
190+
logger.debug("SageMaker Python SDK telemetry not emitted!")
171191

172192

173193
def _construct_url(

‎tests/unit/sagemaker/remote_function/test_client.py

+12
Original file line numberDiff line numberDiff line change
@@ -1504,6 +1504,18 @@ def test_consistency_between_remote_and_step_decorator():
15041504
"s3_kms_key",
15051505
"s3_root_uri",
15061506
"sagemaker_session",
1507+
"args",
1508+
"caught_ex",
1509+
"elapsed",
1510+
"extra",
1511+
"feature_list",
1512+
"kwargs",
1513+
"response",
1514+
"start_timer",
1515+
"stop_timer",
1516+
"studio_app_type",
1517+
"telemetry_opt_out_flag",
1518+
"e",
15071519
]
15081520

15091521
step_args_to_ignore = ["_step", "name", "display_name", "description", "retry_policies"]

0 commit comments

Comments
 (0)