Skip to content

Commit d1ae993

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

File tree

3 files changed

+85
-69
lines changed

3 files changed

+85
-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

+81-69
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
FEATURE_TO_CODE = {
4848
str(Feature.SDK_DEFAULTS): 1,
4949
str(Feature.LOCAL_MODE): 2,
50+
str(Feature.REMOTE_FUNCTION): 3,
5051
}
5152

5253
STATUS_TO_CODE = {
@@ -59,77 +60,88 @@ def _telemetry_emitter(feature: str, func_name: str):
5960
"""Decorator to emit telemetry logs for SageMaker Python SDK functions"""
6061

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

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

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

134146
return wrapper
135147

@@ -165,9 +177,9 @@ def _send_telemetry_request(
165177
# Send the telemetry request
166178
logger.debug("Sending telemetry request to [%s]", url)
167179
_requests_helper(url, 2)
168-
logger.debug("SageMaker Python SDK telemetry successfully emitted!")
180+
logger.debug("SageMaker Python SDK telemetry successfully emitted.")
169181
except Exception: # pylint: disable=W0703
170-
logger.debug("SageMaker Python SDK telemetry not emitted!!")
182+
logger.debug("SageMaker Python SDK telemetry not emitted!")
171183

172184

173185
def _construct_url(

0 commit comments

Comments
 (0)