Skip to content

Commit 2514e52

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

File tree

1 file changed

+80
-69
lines changed

1 file changed

+80
-69
lines changed

‎src/sagemaker/telemetry/telemetry_logging.py

+80-69
Original file line numberDiff line numberDiff line change
@@ -59,77 +59,88 @@ def _telemetry_emitter(feature: str, func_name: str):
5959
"""Decorator to emit telemetry logs for SageMaker Python SDK functions"""
6060

6161
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:
62+
def wrapper(*args, **kwargs):
63+
# Retrieve self instance from the first argument of the function
64+
self_instance = args[0]
65+
if self_instance.sagemaker_session:
66+
logger.debug("sagemaker_session found, preparing to emit telemetry...")
67+
logger.info(TELEMETRY_OPT_OUT_MESSAGING)
68+
response = None
69+
caught_ex = None
70+
studio_app_type = process_studio_metadata_file()
71+
72+
# Check if telemetry is opted out
73+
telemetry_opt_out_flag = resolve_value_from_config(
74+
direct_input=None,
75+
config_path=TELEMETRY_OPT_OUT_PATH,
76+
default_value=False,
77+
sagemaker_session=self_instance.sagemaker_session,
78+
)
79+
logger.debug("TelemetryOptOut flag is set to: %s", telemetry_opt_out_flag)
80+
81+
# Construct the feature list to track feature combinations
82+
feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]]
83+
84+
if (
85+
self_instance.sagemaker_session.sagemaker_config
86+
and feature != Feature.SDK_DEFAULTS
87+
):
8188
feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS)])
8289

83-
if self.sagemaker_session.local_mode and feature != Feature.LOCAL_MODE:
90+
if self_instance.sagemaker_session.local_mode and feature != Feature.LOCAL_MODE:
8491
feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE)])
8592

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
93+
# Construct the extra info to track platform and environment usage metadata
94+
extra = (
95+
f"{func_name}"
96+
f"&x-sdkVersion={SDK_VERSION}"
97+
f"&x-env={PYTHON_VERSION}"
98+
f"&x-sys={OS_NAME_VERSION}"
99+
f"&x-platform={studio_app_type}"
100+
)
101+
102+
# Add endpoint ARN to the extra info if available
103+
if self_instance.sagemaker_session.endpoint_arn:
104+
extra += f"&x-endpointArn={self_instance.sagemaker_session.endpoint_arn}"
105+
106+
start_timer = perf_counter()
107+
try:
108+
# Call the original function
109+
response = func(*args, **kwargs)
110+
stop_timer = perf_counter()
111+
elapsed = stop_timer - start_timer
112+
extra += f"&x-latency={round(elapsed, 2)}"
113+
if not telemetry_opt_out_flag:
114+
_send_telemetry_request(
115+
STATUS_TO_CODE[str(Status.SUCCESS)],
116+
feature_list,
117+
self_instance.sagemaker_session,
118+
None,
119+
None,
120+
extra,
121+
)
122+
except Exception as e: # pylint: disable=W0703
123+
stop_timer = perf_counter()
124+
elapsed = stop_timer - start_timer
125+
extra += f"&x-latency={round(elapsed, 2)}"
126+
if not telemetry_opt_out_flag:
127+
_send_telemetry_request(
128+
STATUS_TO_CODE[str(Status.FAILURE)],
129+
feature_list,
130+
self_instance.sagemaker_session,
131+
str(e),
132+
e.__class__.__name__,
133+
extra,
134+
)
135+
caught_ex = e
136+
finally:
137+
if caught_ex:
138+
raise caught_ex
139+
return response # pylint: disable=W0150
140+
else:
141+
logger.debug("No sagemaker_session found, telemetry will not be emitted!")
142+
response = func(*args, **kwargs)
143+
return response
133144

134145
return wrapper
135146

@@ -165,9 +176,9 @@ def _send_telemetry_request(
165176
# Send the telemetry request
166177
logger.debug("Sending telemetry request to [%s]", url)
167178
_requests_helper(url, 2)
168-
logger.debug("SageMaker Python SDK telemetry successfully emitted!")
179+
logger.debug("SageMaker Python SDK telemetry successfully emitted.")
169180
except Exception: # pylint: disable=W0703
170-
logger.debug("SageMaker Python SDK telemetry not emitted!!")
181+
logger.debug("SageMaker Python SDK telemetry not emitted!")
171182

172183

173184
def _construct_url(

0 commit comments

Comments
 (0)