From 0322c4679f5e46c34ebfcf4899e7eaf5053e13a0 Mon Sep 17 00:00:00 2001 From: Gokul A Date: Fri, 21 Aug 2026 15:20:01 -0700 Subject: [PATCH] feat: track explicit debugger and profiler configuration --- src/sagemaker/estimator.py | 11 +++++++- src/sagemaker/telemetry/telemetry_logging.py | 17 ++++++++++-- .../telemetry/test_telemetry_logging.py | 26 +++++++++++++++++++ tests/unit/test_estimator.py | 13 +++++++++- 4 files changed, 63 insertions(+), 4 deletions(-) diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 2805bb73ae..011429250b 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -775,6 +775,8 @@ def __init__( self.checkpoint_local_path = checkpoint_local_path self.rules = rules + self._debugger_hook_config_explicitly_provided = debugger_hook_config is not None + self._profiler_config_explicitly_provided = profiler_config is not None # Today, we ONLY support debugger_hook_config to be provided as a boolean value # from sagemaker_config. We resolve value for this parameter as per the order @@ -1366,7 +1368,14 @@ def latest_job_profiler_artifacts_path(self): ) return None - @_telemetry_emitter(feature=Feature.ESTIMATOR_V2, func_name="estimator.fit") + @_telemetry_emitter( + feature=Feature.ESTIMATOR_V2, + func_name="estimator.fit", + telemetry_params={ + "debuggerHookConfigExplicitlyProvided": ("_debugger_hook_config_explicitly_provided"), + "profilerConfigExplicitlyProvided": "_profiler_config_explicitly_provided", + }, + ) @runnable_by_pipeline def fit( self, diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index 15dcb7be50..f8b82d0fbf 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -16,7 +16,7 @@ import platform import sys from time import perf_counter -from typing import List +from typing import Dict, List import functools import requests @@ -67,7 +67,7 @@ } -def _telemetry_emitter(feature: str, func_name: str): +def _telemetry_emitter(feature: str, func_name: str, telemetry_params: Dict[str, str] = None): """Telemetry Emitter Decorator to emit telemetry logs for SageMaker Python SDK functions. This class needs @@ -75,6 +75,11 @@ def _telemetry_emitter(feature: str, func_name: str): in this repo. When collecting telemetry for classes using sagemaker-core Session object, we should be aware of its differences, such as sagemaker_session.sagemaker_config does not exist in new Session class. + + Args: + feature: Feature enum value for the telemetry event. + func_name: Name of the instrumented function. + telemetry_params: Mapping of telemetry field names to instance attribute names. """ def decorator(func): @@ -136,6 +141,14 @@ def wrapper(*args, **kwargs): if hasattr(sagemaker_session, "endpoint_arn") and sagemaker_session.endpoint_arn: extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}" + if telemetry_params and args: + for field_name, attribute_name in telemetry_params.items(): + value = getattr(args[0], attribute_name, None) + if isinstance(value, bool): + value = str(value).lower() + if value is not None: + extra += f"&x-{field_name}={value}" + start_timer = perf_counter() try: # Call the original function diff --git a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py index 031ddfc291..21da103233 100644 --- a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py +++ b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py @@ -42,12 +42,25 @@ class LocalSagemakerClientMock: def __init__(self): self.sagemaker_session = MOCK_SESSION + self._debugger_explicitly_provided = True + self._profiler_explicitly_provided = False @_telemetry_emitter(MOCK_FEATURE, MOCK_FUNC_NAME) def mock_create_model(self, mock_exception_func=None): if mock_exception_func: mock_exception_func() + @_telemetry_emitter( + MOCK_FEATURE, + MOCK_FUNC_NAME, + telemetry_params={ + "debuggerExplicitlyProvided": "_debugger_explicitly_provided", + "profilerExplicitlyProvided": "_profiler_explicitly_provided", + }, + ) + def mock_train(self): + return None + class TestTelemetryLogging(unittest.TestCase): @patch("sagemaker.telemetry.telemetry_logging._requests_helper") @@ -147,6 +160,19 @@ def test_telemetry_emitter_decorator_success( 1, [1, 2], MOCK_SESSION, None, None, expected_extra_str ) + @patch("sagemaker.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_adds_instance_params( + self, mock_resolve_config, mock_send_telemetry_request + ): + mock_resolve_config.return_value = False + + LocalSagemakerClientMock().mock_train() + + extra = mock_send_telemetry_request.call_args.args[5] + assert "&x-debuggerExplicitlyProvided=true" in extra + assert "&x-profilerExplicitlyProvided=false" in extra + @patch("sagemaker.telemetry.telemetry_logging._send_telemetry_request") @patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config") def test_telemetry_emitter_decorator_handle_exception_success( diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index 7ea64ab2ea..7e69084f94 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -672,6 +672,8 @@ def test_estimator_initialization_with_sagemaker_config_injection(sagemaker_sess assert estimator.environment == expected_environment assert estimator.disable_profiler == expected_disable_profiler_attribute assert estimator.debugger_hook_config == expected_debugger_hook_config + assert estimator._debugger_hook_config_explicitly_provided is False + assert estimator._profiler_config_explicitly_provided is False def test_estimator_with_debugger_hook_config_provided_as_bool_from_direct_input( @@ -694,6 +696,7 @@ def test_estimator_with_debugger_hook_config_provided_as_bool_from_direct_input( debugger_hook_config=True, ) assert estimator.debugger_hook_config == {} + assert estimator._debugger_hook_config_explicitly_provided is True def test_estimator_with_debugger_hook_config_provided_as_dict_from_direct_input( @@ -715,6 +718,7 @@ def test_estimator_with_debugger_hook_config_provided_as_dict_from_direct_input( debugger_hook_config=HOOK_CONFIG, ) assert estimator.debugger_hook_config == HOOK_CONFIG + assert estimator._debugger_hook_config_explicitly_provided is True def test_estimator_initialization_with_sagemaker_config_injection_no_kms_supported( @@ -1351,8 +1355,11 @@ def test_framework_with_only_profiler_rule_specified(sagemaker_session): ] +@patch("sagemaker.telemetry.telemetry_logging._send_telemetry_request") @patch("time.time", return_value=TIME) -def test_framework_with_profiler_config_without_s3_output_path(time, sagemaker_session): +def test_framework_with_profiler_config_without_s3_output_path( + time, send_telemetry_request, sagemaker_session +): f = DummyFramework( entry_point=SCRIPT_PATH, role=ROLE, @@ -1361,6 +1368,7 @@ def test_framework_with_profiler_config_without_s3_output_path(time, sagemaker_s instance_type=INSTANCE_TYPE, profiler_config=ProfilerConfig(system_monitor_interval_millis=1000), ) + assert f._profiler_config_explicitly_provided is True f.fit("s3://mydata") sagemaker_session.train.assert_called_once() _, args = sagemaker_session.train.call_args @@ -1369,6 +1377,9 @@ def test_framework_with_profiler_config_without_s3_output_path(time, sagemaker_s "S3OutputPath": "s3://{}/".format(BUCKET_NAME), "ProfilingIntervalInMilliseconds": 1000, } + telemetry_extra = send_telemetry_request.call_args.args[5] + assert "&x-debuggerHookConfigExplicitlyProvided=false" in telemetry_extra + assert "&x-profilerConfigExplicitlyProvided=true" in telemetry_extra @pytest.mark.parametrize("region", PROFILER_UNSUPPORTED_REGIONS)