From ace5a12c22c2d67fda5345c5f0d19c8e07b95480 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Wed, 9 Sep 2026 18:47:51 +0000 Subject: [PATCH 1/5] change: emit the JumpStart model ID in v2 telemetry The v2 telemetry cannot identify the JumpStart model behind an event. JumpStartModel and JumpStartEstimator events carry no model ID. ModelBuilder events carry x-modelHub=1 but no ID, and post to the dev-exp-t bucket instead of sm-pysdk-t, so a query on the v3 bucket misses them. Append x-jumpstartModelId= to JUMPSTART_V2 events in _telemetry_emitter and to JumpStart ModelBuilder events in _capture_telemetry. Route ModelBuilder telemetry to sm-pysdk-t, the bucket the v2 emitter and the v3 SDK use. Other model sources omit the param. --- X-AI-Prompt: Can you raise a PR against master-v2 that emits the JumpStart model ID in telemetry like v3 does? X-AI-Tool: claude-code --- src/sagemaker/serve/utils/telemetry_logger.py | 6 +- src/sagemaker/telemetry/telemetry_logging.py | 12 ++++ .../serve/utils/test_telemetry_logger.py | 66 ++++++++++++++++- .../telemetry/test_telemetry_logging.py | 72 +++++++++++++++++++ 4 files changed, 151 insertions(+), 5 deletions(-) diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py index 6e7db9043b..1309799af6 100644 --- a/src/sagemaker/serve/utils/telemetry_logger.py +++ b/src/sagemaker/serve/utils/telemetry_logger.py @@ -151,6 +151,8 @@ def wrapper(self, *args, **kwargs): if getattr(self, "model_hub", False): extra += f"&x-modelHub={MODEL_HUB_TO_CODE[str(self.model_hub)]}" + if self.model_hub == ModelHub.JUMPSTART and isinstance(self.model, str): + extra += f"&x-jumpstartModelId={self.model}" if getattr(self, "is_fine_tuned", False): extra += "&x-fineTuned=1" @@ -230,10 +232,10 @@ def _construct_url( extra_info: str, region: str, ) -> str: - """Placeholder docstring""" + """Construct the URL for the telemetry request""" base_url = ( - f"https://dev-exp-t-{region}.s3.{region}.amazonaws.com/telemetry?" + f"https://sm-pysdk-t-{region}.s3.{region}.amazonaws.com/telemetry?" f"x-accountId={accountId}" f"&x-mode={mode}" f"&x-status={status}" diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index 15dcb7be50..ddca1b9483 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -67,6 +67,16 @@ } +def _jumpstart_model_id_param(feature: str, args: tuple) -> str: + """Return the JumpStart model ID query param, or an empty string for another feature.""" + if feature != Feature.JUMPSTART_V2 or len(args) == 0: + return "" + model_id = getattr(args[0], "model_id", None) + if not model_id: + return "" + return f"&x-jumpstartModelId={model_id}" + + def _telemetry_emitter(feature: str, func_name: str): """Telemetry Emitter @@ -136,6 +146,8 @@ def wrapper(*args, **kwargs): if hasattr(sagemaker_session, "endpoint_arn") and sagemaker_session.endpoint_arn: extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}" + extra += _jumpstart_model_id_param(feature, args) + start_timer = perf_counter() try: # Call the original function diff --git a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py index fc832ad02d..a2ea5e658b 100644 --- a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py +++ b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py @@ -21,7 +21,7 @@ _construct_url, ) from sagemaker.serve.utils.exceptions import ModelBuilderException, LocalModelOutOfMemoryException -from sagemaker.serve.utils.types import ImageUriOption +from sagemaker.serve.utils.types import ImageUriOption, ModelHub from sagemaker.user_agent import SDK_VERSION MOCK_SESSION = Mock() @@ -38,6 +38,7 @@ "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:2.0.1-cpu-py310" ) MOCK_HUGGINGFACE_ID = "meta-llama/Llama-2-7b-hf" +MOCK_JUMPSTART_ID = "huggingface-llm-falcon-7b-bf16" MOCK_EXCEPTION = LocalModelOutOfMemoryException("mock raise ex") MOCK_ENDPOINT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:endpoint/test" MOCK_MODEL_METADATA_FOR_MLFLOW = { @@ -69,7 +70,7 @@ def test_log_sucessfully(self, mocked_get_accountId, mocked_request_helper): mocked_get_accountId.return_value = "testAccountId" _send_telemetry("someStatus", 1, MOCK_SESSION) mocked_request_helper.assert_called_with( - "https://dev-exp-t-ap-south-1.s3.ap-south-1.amazonaws.com/" + "https://sm-pysdk-t-ap-south-1.s3.ap-south-1.amazonaws.com/" "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", 2, ) @@ -170,6 +171,65 @@ def test_capture_telemetry_decorator_tgi_success(self, mock_send_telemetry): "1", 2, MOCK_SESSION, None, None, expected_extra_str ) + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") + def test_capture_telemetry_decorator_jumpstart_emits_model_id(self, mock_send_telemetry): + mock_model_builder = ModelBuilderMock() + mock_model_builder.serve_settings.telemetry_opt_out = False + mock_model_builder.image_uri = MOCK_TGI_CONTAINER + mock_model_builder._is_custom_image_uri = False + mock_model_builder.model = MOCK_JUMPSTART_ID + mock_model_builder.model_hub = ModelHub.JUMPSTART + mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT + mock_model_builder.model_server = ModelServer.MMS + mock_model_builder.sagemaker_session.endpoint_arn = MOCK_ENDPOINT_ARN + + mock_model_builder.mock_deploy() + + args = mock_send_telemetry.call_args.args + latency = str(args[5]).split("latency=")[1] + expected_extra_str = ( + f"{MOCK_DEPLOY_FUNC_NAME}" + "&x-modelServer=2" + "&x-imageTag=huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04" + f"&x-sdkVersion={SDK_VERSION}" + f"&x-defaultImageUsage={ImageUriOption.DEFAULT_IMAGE.value}" + f"&x-endpointArn={MOCK_ENDPOINT_ARN}" + "&x-modelHub=1" + f"&x-jumpstartModelId={MOCK_JUMPSTART_ID}" + f"&x-latency={latency}" + ) + + mock_send_telemetry.assert_called_once_with( + "1", 3, MOCK_SESSION, None, None, expected_extra_str + ) + + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") + def test_capture_telemetry_decorator_huggingface_hub_omits_model_id(self, mock_send_telemetry): + mock_model_builder = ModelBuilderMock() + mock_model_builder.serve_settings.telemetry_opt_out = False + mock_model_builder.image_uri = None + mock_model_builder.model = MOCK_HUGGINGFACE_ID + mock_model_builder.model_hub = ModelHub.HUGGINGFACE + mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT + mock_model_builder.model_server = ModelServer.MMS + mock_model_builder.sagemaker_session.endpoint_arn = None + + mock_model_builder.mock_deploy() + + args = mock_send_telemetry.call_args.args + latency = str(args[5]).split("latency=")[1] + expected_extra_str = ( + f"{MOCK_DEPLOY_FUNC_NAME}" + "&x-modelServer=2" + f"&x-sdkVersion={SDK_VERSION}" + "&x-modelHub=2" + f"&x-latency={latency}" + ) + + mock_send_telemetry.assert_called_once_with( + "1", 3, MOCK_SESSION, None, None, expected_extra_str + ) + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") def test_capture_telemetry_decorator_no_call_when_disabled(self, mock_send_telemetry): mock_model_builder = ModelBuilderMock() @@ -243,7 +303,7 @@ def test_construct_url_with_failure_reason_and_extra_info(self): ) expected_base_url = ( - f"https://dev-exp-t-{mock_region}.s3.{mock_region}.amazonaws.com/telemetry?" + f"https://sm-pysdk-t-{mock_region}.s3.{mock_region}.amazonaws.com/telemetry?" f"x-accountId={mock_accountId}" f"&x-mode={mock_mode}" f"&x-status={mock_status}" diff --git a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py index 031ddfc291..b8c717ff89 100644 --- a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py +++ b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py @@ -37,6 +37,8 @@ MOCK_FEATURE = Feature.SDK_DEFAULTS_V2 MOCK_FUNC_NAME = "Mock.local_session.create_model" MOCK_ENDPOINT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:endpoint/test" +MOCK_JUMPSTART_FUNC_NAME = "Mock.jumpstart_model.deploy" +MOCK_JUMPSTART_MODEL_ID = "huggingface-llm-falcon-7b-bf16" class LocalSagemakerClientMock: @@ -49,6 +51,20 @@ def mock_create_model(self, mock_exception_func=None): mock_exception_func() +class JumpStartModelMock: + def __init__(self, model_id): + self.sagemaker_session = MOCK_SESSION + self.model_id = model_id + + @_telemetry_emitter(Feature.JUMPSTART_V2, MOCK_JUMPSTART_FUNC_NAME) + def mock_deploy(self): + pass + + @_telemetry_emitter(MOCK_FEATURE, MOCK_FUNC_NAME) + def mock_create_model(self): + pass + + class TestTelemetryLogging(unittest.TestCase): @patch("sagemaker.telemetry.telemetry_logging._requests_helper") @patch("sagemaker.telemetry.telemetry_logging._get_accountId") @@ -186,6 +202,62 @@ def test_telemetry_emitter_decorator_handle_exception_success( expected_extra_str, ) + @patch("sagemaker.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_jumpstart_feature_emits_model_id( + self, mock_resolve_config, mock_send_telemetry_request + ): + """A JUMPSTART_V2 event carries the JumpStart model ID""" + mock_resolve_config.return_value = False + mock_jumpstart_model = JumpStartModelMock(MOCK_JUMPSTART_MODEL_ID) + mock_jumpstart_model.sagemaker_session.endpoint_arn = None + mock_jumpstart_model.mock_deploy() + app_type = process_studio_metadata_file() + + args = mock_send_telemetry_request.call_args.args + latency = str(args[5]).split("latency=")[1] + expected_extra_str = ( + f"{MOCK_JUMPSTART_FUNC_NAME}" + f"&x-sdkVersion={SDK_VERSION}" + f"&x-env={PYTHON_VERSION}" + f"&x-sys={OS_NAME_VERSION}" + f"&x-platform={app_type}" + f"&x-jumpstartModelId={MOCK_JUMPSTART_MODEL_ID}" + f"&x-latency={latency}" + ) + + mock_send_telemetry_request.assert_called_once_with( + 1, [8, 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_other_feature_omits_model_id( + self, mock_resolve_config, mock_send_telemetry_request + ): + """An event for another feature omits the model ID that the instance has""" + mock_resolve_config.return_value = False + mock_jumpstart_model = JumpStartModelMock(MOCK_JUMPSTART_MODEL_ID) + mock_jumpstart_model.sagemaker_session.endpoint_arn = None + mock_jumpstart_model.mock_create_model() + + extra_str = mock_send_telemetry_request.call_args.args[5] + assert "x-jumpstartModelId" not in extra_str + + @patch("sagemaker.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_jumpstart_feature_without_model_id( + self, mock_resolve_config, mock_send_telemetry_request + ): + """A JUMPSTART_V2 event with no model ID has no model ID param""" + mock_resolve_config.return_value = False + mock_jumpstart_model = JumpStartModelMock(None) + mock_jumpstart_model.sagemaker_session.endpoint_arn = None + mock_jumpstart_model.mock_deploy() + + extra_str = mock_send_telemetry_request.call_args.args[5] + assert "x-jumpstartModelId" not in extra_str + def test_construct_url_with_failure_reason_and_extra_info(self): """Test to verify the _construct_url function with failure reason and extra info""" mock_accountId = "testAccountId" From e0e32b3fe4ca145e929d469aa28fff559c414800 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Wed, 9 Sep 2026 19:39:42 +0000 Subject: [PATCH 2/5] change: post ModelBuilder telemetry to the two buckets The previous commit moved ModelBuilder telemetry from dev-exp-t to sm-pysdk-t. Consumers of dev-exp-t lose the events. Post each ModelBuilder event to dev-exp-t and to sm-pysdk-t. The old bucket keeps its events, and the new bucket lets one query cover the v2 and v3 JumpStart events. Split the URL builder into a query builder and a per-bucket URL builder. --- X-AI-Prompt: Can you make the ModelBuilder telemetry keep the old bucket and add the new one instead of replacing it? X-AI-Tool: claude-code --- src/sagemaker/serve/utils/telemetry_logger.py | 39 ++++++++++--------- .../serve/utils/test_telemetry_logger.py | 39 ++++++++++++++----- 2 files changed, 50 insertions(+), 28 deletions(-) diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py index 1309799af6..51e79c57ac 100644 --- a/src/sagemaker/serve/utils/telemetry_logger.py +++ b/src/sagemaker/serve/utils/telemetry_logger.py @@ -50,6 +50,10 @@ "for more info." ) +# The old bucket keeps its consumers. sm-pysdk-t is the bucket that the sagemaker.telemetry +# emitter and the v3 SDK post to, so one query covers every JumpStart event. +TELEMETRY_BUCKET_PREFIXES = ("dev-exp-t", "sm-pysdk-t") + MODE_TO_CODE = { str(Mode.IN_PROCESS): 1, str(Mode.LOCAL_CONTAINER): 2, @@ -204,48 +208,47 @@ def _send_telemetry( failure_type: str = None, extra_info: str = None, ) -> None: - """Make GET request to an empty object in S3 bucket""" + """Make a GET request to an empty object in each telemetry S3 bucket""" try: accountId = _get_accountId(session) region = _get_region_or_default(session) - url = _construct_url( + query = _construct_query( accountId, str(mode), status, failure_reason, failure_type, extra_info, - region, ) - _requests_helper(url, 2) + for bucket_prefix in TELEMETRY_BUCKET_PREFIXES: + _requests_helper(_construct_url(bucket_prefix, region, query), 2) logger.debug("ModelBuilder metrics emitted.") except Exception: # pylint: disable=W0703 logger.debug("ModelBuilder metrics not emitted") -def _construct_url( +def _construct_query( accountId: str, mode: str, status: str, failure_reason: str, failure_type: str, extra_info: str, - region: str, ) -> str: - """Construct the URL for the telemetry request""" - - base_url = ( - f"https://sm-pysdk-t-{region}.s3.{region}.amazonaws.com/telemetry?" - f"x-accountId={accountId}" - f"&x-mode={mode}" - f"&x-status={status}" - ) + """Construct the query string for the telemetry request""" + + query = f"x-accountId={accountId}&x-mode={mode}&x-status={status}" if failure_reason: - base_url += f"&x-failureReason={failure_reason}" - base_url += f"&x-failureType={failure_type}" + query += f"&x-failureReason={failure_reason}" + query += f"&x-failureType={failure_type}" if extra_info: - base_url += f"&x-extra={extra_info}" - return base_url + query += f"&x-extra={extra_info}" + return query + + +def _construct_url(bucket_prefix: str, region: str, query: str) -> str: + """Construct the URL for the telemetry request to one bucket""" + return f"https://{bucket_prefix}-{region}.s3.{region}.amazonaws.com/telemetry?{query}" def _requests_helper(url, timeout): diff --git a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py index a2ea5e658b..5227bcb3f6 100644 --- a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py +++ b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py @@ -12,12 +12,14 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import import unittest -from unittest.mock import Mock, patch, MagicMock +from unittest.mock import Mock, call, patch, MagicMock from sagemaker.serve import Mode, ModelServer from sagemaker.serve.model_format.mlflow.constants import MLFLOW_MODEL_PATH, MLFLOW_TRACKING_ARN from sagemaker.serve.utils.telemetry_logger import ( + TELEMETRY_BUCKET_PREFIXES, _send_telemetry, _capture_telemetry, + _construct_query, _construct_url, ) from sagemaker.serve.utils.exceptions import ModelBuilderException, LocalModelOutOfMemoryException @@ -69,11 +71,21 @@ def test_log_sucessfully(self, mocked_get_accountId, mocked_request_helper): MOCK_SESSION.boto_session.region_name = "ap-south-1" mocked_get_accountId.return_value = "testAccountId" _send_telemetry("someStatus", 1, MOCK_SESSION) - mocked_request_helper.assert_called_with( - "https://sm-pysdk-t-ap-south-1.s3.ap-south-1.amazonaws.com/" - "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", - 2, + mocked_request_helper.assert_has_calls( + [ + call( + "https://dev-exp-t-ap-south-1.s3.ap-south-1.amazonaws.com/" + "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", + 2, + ), + call( + "https://sm-pysdk-t-ap-south-1.s3.ap-south-1.amazonaws.com/" + "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", + 2, + ), + ] ) + self.assertEqual(mocked_request_helper.call_count, 2) @patch("sagemaker.serve.utils.telemetry_logger._get_accountId") def test_log_handle_exception(self, mocked_get_accountId): @@ -292,18 +304,16 @@ def test_construct_url_with_failure_reason_and_extra_info(self): mock_extra_info = "mock_extra_info" mock_region = "us-west-2" - ret_url = _construct_url( + query = _construct_query( accountId=mock_accountId, mode=mock_mode, status=mock_status, failure_reason=mock_failure_reason, failure_type=mock_failure_type, extra_info=mock_extra_info, - region=mock_region, ) - expected_base_url = ( - f"https://sm-pysdk-t-{mock_region}.s3.{mock_region}.amazonaws.com/telemetry?" + expected_query = ( f"x-accountId={mock_accountId}" f"&x-mode={mock_mode}" f"&x-status={mock_status}" @@ -311,7 +321,16 @@ def test_construct_url_with_failure_reason_and_extra_info(self): f"&x-failureType={mock_failure_type}" f"&x-extra={mock_extra_info}" ) - self.assertEqual(ret_url, expected_base_url) + self.assertEqual(query, expected_query) + self.assertEqual(TELEMETRY_BUCKET_PREFIXES, ("dev-exp-t", "sm-pysdk-t")) + self.assertEqual( + _construct_url("dev-exp-t", mock_region, query), + f"https://dev-exp-t-us-west-2.s3.us-west-2.amazonaws.com/telemetry?{expected_query}", + ) + self.assertEqual( + _construct_url("sm-pysdk-t", mock_region, query), + f"https://sm-pysdk-t-us-west-2.s3.us-west-2.amazonaws.com/telemetry?{expected_query}", + ) @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") def test_capture_telemetry_decorator_mlflow_success(self, mock_send_telemetry): From b502e40d16d29718e7a69cd6bf4fc6d5cd91ee4a Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Wed, 9 Sep 2026 22:06:01 +0000 Subject: [PATCH 3/5] change: send JumpStart ModelBuilder events through the SDK emitter The previous commit posted every ModelBuilder event to dev-exp-t and to sm-pysdk-t. That put the ModelBuilder schema (x-mode, no x-feature) into the bucket that the SDK emitter and the v3 SDK use. Keep the ModelBuilder event on dev-exp-t as before. For a JumpStart model, ModelBuilder.build and ModelBuilder.deploy also send one JUMPSTART_V2 event through sagemaker.telemetry, with the emitter schema (x-feature=8, x-sdkVersion, x-env, x-sys, x-platform, x-mode, x-jumpstartModelId, x-latency). Another model source or another method sends nothing to sm-pysdk-t. Move the feature list and the base x-extra string of the emitter into helpers that the two paths share. --- X-AI-Prompt: Can you post to the v3 bucket only for JumpStart calls and keep the v3 schema there? X-AI-Tool: claude-code --- src/sagemaker/serve/utils/telemetry_logger.py | 106 +++++++-- src/sagemaker/telemetry/telemetry_logging.py | 65 +++--- .../serve/utils/test_telemetry_logger.py | 212 +++++++++++------- 3 files changed, 250 insertions(+), 133 deletions(-) diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py index 51e79c57ac..1f1cc9541c 100644 --- a/src/sagemaker/serve/utils/telemetry_logger.py +++ b/src/sagemaker/serve/utils/telemetry_logger.py @@ -13,7 +13,9 @@ """Placeholder docstring""" from __future__ import absolute_import import logging +from dataclasses import dataclass from time import perf_counter +from typing import Optional import requests @@ -36,6 +38,12 @@ SpeculativeDecodingDraftModelSource, ) from sagemaker.serve.validations.check_image_uri import is_1p_image_uri +from sagemaker.telemetry.constants import Feature +from sagemaker.telemetry.telemetry_logging import ( + _base_extra, + _feature_list, + _send_telemetry_request, +) from sagemaker.user_agent import SDK_VERSION logger = logging.getLogger(__name__) @@ -50,16 +58,15 @@ "for more info." ) -# The old bucket keeps its consumers. sm-pysdk-t is the bucket that the sagemaker.telemetry -# emitter and the v3 SDK post to, so one query covers every JumpStart event. -TELEMETRY_BUCKET_PREFIXES = ("dev-exp-t", "sm-pysdk-t") - MODE_TO_CODE = { str(Mode.IN_PROCESS): 1, str(Mode.LOCAL_CONTAINER): 2, str(Mode.SAGEMAKER_ENDPOINT): 3, } +# The v3 SDK emits x-jumpstartModelId on model_builder.build and model_builder.deploy only. +JUMPSTART_SDK_EVENT_FUNC_NAMES = frozenset({"ModelBuilder.build", "ModelBuilder.deploy"}) + MODEL_SERVER_TO_CODE = { str(ModelServer.TORCHSERVE): 1, str(ModelServer.MMS): 2, @@ -155,8 +162,6 @@ def wrapper(self, *args, **kwargs): if getattr(self, "model_hub", False): extra += f"&x-modelHub={MODEL_HUB_TO_CODE[str(self.model_hub)]}" - if self.model_hub == ModelHub.JUMPSTART and isinstance(self.model, str): - extra += f"&x-jumpstartModelId={self.model}" if getattr(self, "is_fine_tuned", False): extra += "&x-fineTuned=1" @@ -178,7 +183,8 @@ def wrapper(self, *args, **kwargs): config_name_code = self.deployment_config_name.lower() extra += f"&x-configName={config_name_code}" - extra += f"&x-latency={round(elapsed, 2)}" + latency = round(elapsed, 2) + extra += f"&x-latency={latency}" if hasattr(self, "serve_settings") and not self.serve_settings.telemetry_opt_out: _send_telemetry( @@ -189,6 +195,8 @@ def wrapper(self, *args, **kwargs): failure_type, extra, ) + outcome = _CallOutcome(status, latency, failure_reason, failure_type) + _send_jumpstart_telemetry(self, func_name, outcome) if caught_ex: raise caught_ex @@ -200,6 +208,57 @@ def wrapper(self, *args, **kwargs): return decorator +@dataclass +class _CallOutcome: + """The result of one decorated call: the status code, the latency, and the caught error.""" + + status: str + latency: float + failure_reason: Optional[str] + failure_type: Optional[str] + + +def _jumpstart_model_id(model_builder) -> str: + """Return the JumpStart model ID, or an empty string for another model source.""" + if getattr(model_builder, "model_hub", None) != ModelHub.JUMPSTART: + return "" + model = getattr(model_builder, "model", None) + if not isinstance(model, str): + return "" + return model + + +def _send_jumpstart_telemetry(model_builder, func_name: str, outcome: _CallOutcome) -> None: + """Send one JUMPSTART_V2 event for a JumpStart build or deploy through the SDK emitter. + + The event has the emitter schema, so it lands in the same bucket and shape as the + sagemaker.telemetry events and the v3 SDK events. Another model source or another method + sends nothing here. + """ + if func_name not in JUMPSTART_SDK_EVENT_FUNC_NAMES: + return + model_id = _jumpstart_model_id(model_builder) + if not model_id: + return + try: + session = model_builder.sagemaker_session + extra = _base_extra(func_name, session) + mode = getattr(model_builder, "mode", None) + if mode is not None: + extra += f"&x-mode={mode}" + extra += f"&x-jumpstartModelId={model_id}&x-latency={outcome.latency}" + _send_telemetry_request( + int(outcome.status), + _feature_list(Feature.JUMPSTART_V2, session), + session, + outcome.failure_reason, + outcome.failure_type, + extra, + ) + except Exception: # pylint: disable=W0703 + logger.debug("JumpStart SDK telemetry not emitted") + + def _send_telemetry( status: str, mode: int, @@ -208,47 +267,48 @@ def _send_telemetry( failure_type: str = None, extra_info: str = None, ) -> None: - """Make a GET request to an empty object in each telemetry S3 bucket""" + """Make GET request to an empty object in S3 bucket""" try: accountId = _get_accountId(session) region = _get_region_or_default(session) - query = _construct_query( + url = _construct_url( accountId, str(mode), status, failure_reason, failure_type, extra_info, + region, ) - for bucket_prefix in TELEMETRY_BUCKET_PREFIXES: - _requests_helper(_construct_url(bucket_prefix, region, query), 2) + _requests_helper(url, 2) logger.debug("ModelBuilder metrics emitted.") except Exception: # pylint: disable=W0703 logger.debug("ModelBuilder metrics not emitted") -def _construct_query( +def _construct_url( accountId: str, mode: str, status: str, failure_reason: str, failure_type: str, extra_info: str, + region: str, ) -> str: - """Construct the query string for the telemetry request""" + """Placeholder docstring""" - query = f"x-accountId={accountId}&x-mode={mode}&x-status={status}" + base_url = ( + f"https://dev-exp-t-{region}.s3.{region}.amazonaws.com/telemetry?" + f"x-accountId={accountId}" + f"&x-mode={mode}" + f"&x-status={status}" + ) if failure_reason: - query += f"&x-failureReason={failure_reason}" - query += f"&x-failureType={failure_type}" + base_url += f"&x-failureReason={failure_reason}" + base_url += f"&x-failureType={failure_type}" if extra_info: - query += f"&x-extra={extra_info}" - return query - - -def _construct_url(bucket_prefix: str, region: str, query: str) -> str: - """Construct the URL for the telemetry request to one bucket""" - return f"https://{bucket_prefix}-{region}.s3.{region}.amazonaws.com/telemetry?{query}" + base_url += f"&x-extra={extra_info}" + return base_url def _requests_helper(url, timeout): diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index ddca1b9483..9e6d8277b9 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -77,6 +77,38 @@ def _jumpstart_model_id_param(feature: str, args: tuple) -> str: return f"&x-jumpstartModelId={model_id}" +def _feature_list(feature: str, sagemaker_session) -> List[int]: + """Return the feature codes of an event: the feature, then SDK defaults and local mode.""" + feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]] + if ( + hasattr(sagemaker_session, "sagemaker_config") + and sagemaker_session.sagemaker_config + and feature != Feature.SDK_DEFAULTS_V2 + ): + feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS_V2)]) + if ( + hasattr(sagemaker_session, "local_mode") + and sagemaker_session.local_mode + and feature != Feature.LOCAL_MODE_V2 + ): + feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE_V2)]) + return feature_list + + +def _base_extra(func_name: str, sagemaker_session) -> str: + """Return the platform and environment part of the x-extra string of an event.""" + extra = ( + f"{func_name}" + f"&x-sdkVersion={SDK_VERSION}" + f"&x-env={PYTHON_VERSION}" + f"&x-sys={OS_NAME_VERSION}" + f"&x-platform={process_studio_metadata_file()}" + ) + if hasattr(sagemaker_session, "endpoint_arn") and sagemaker_session.endpoint_arn: + extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}" + return extra + + def _telemetry_emitter(feature: str, func_name: str): """Telemetry Emitter @@ -105,7 +137,6 @@ def wrapper(*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( @@ -116,36 +147,8 @@ def wrapper(*args, **kwargs): ) 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 ( - hasattr(sagemaker_session, "sagemaker_config") - and sagemaker_session.sagemaker_config - and feature != Feature.SDK_DEFAULTS_V2 - ): - feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS_V2)]) - - if ( - hasattr(sagemaker_session, "local_mode") - and sagemaker_session.local_mode - and feature != Feature.LOCAL_MODE_V2 - ): - feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE_V2)]) - - # 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 hasattr(sagemaker_session, "endpoint_arn") and sagemaker_session.endpoint_arn: - extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}" - + feature_list = _feature_list(feature, sagemaker_session) + extra = _base_extra(func_name, sagemaker_session) extra += _jumpstart_model_id_param(feature, args) start_timer = perf_counter() diff --git a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py index 5227bcb3f6..e280981fab 100644 --- a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py +++ b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py @@ -12,23 +12,23 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import import unittest -from unittest.mock import Mock, call, patch, MagicMock +from unittest.mock import Mock, patch, MagicMock from sagemaker.serve import Mode, ModelServer from sagemaker.serve.model_format.mlflow.constants import MLFLOW_MODEL_PATH, MLFLOW_TRACKING_ARN from sagemaker.serve.utils.telemetry_logger import ( - TELEMETRY_BUCKET_PREFIXES, _send_telemetry, _capture_telemetry, - _construct_query, _construct_url, ) from sagemaker.serve.utils.exceptions import ModelBuilderException, LocalModelOutOfMemoryException from sagemaker.serve.utils.types import ImageUriOption, ModelHub +from sagemaker.telemetry.telemetry_logging import OS_NAME_VERSION, PYTHON_VERSION from sagemaker.user_agent import SDK_VERSION MOCK_SESSION = Mock() MOCK_DEPLOY_FUNC_NAME = "Mock.deploy" MOCK_OPTIMIZE_FUNC_NAME = "Mock.optimize" +MODEL_BUILDER_DEPLOY_FUNC_NAME = "ModelBuilder.deploy" MOCK_DJL_CONTAINER = ( "763104351884.dkr.ecr.us-west-2.amazonaws.com/" "djl-inference:0.25.0-deepspeed0.11.0-cu118" ) @@ -41,6 +41,7 @@ ) MOCK_HUGGINGFACE_ID = "meta-llama/Llama-2-7b-hf" MOCK_JUMPSTART_ID = "huggingface-llm-falcon-7b-bf16" +MOCK_STUDIO_APP_TYPE = "KernelGateway" MOCK_EXCEPTION = LocalModelOutOfMemoryException("mock raise ex") MOCK_ENDPOINT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:endpoint/test" MOCK_MODEL_METADATA_FOR_MLFLOW = { @@ -63,29 +64,51 @@ def mock_deploy(self, mock_exception_func=None): def mock_optimize(self, *args, **kwargs): pass + @_capture_telemetry(MODEL_BUILDER_DEPLOY_FUNC_NAME) + def deploy(self, mock_exception_func=None): + if mock_exception_func: + mock_exception_func() + class TestTelemetryLogger(unittest.TestCase): + def setUp(self): + MOCK_SESSION.sagemaker_config = None + MOCK_SESSION.local_mode = False + + def _jumpstart_model_builder(self): + mock_model_builder = ModelBuilderMock() + mock_model_builder.serve_settings.telemetry_opt_out = False + mock_model_builder.image_uri = None + mock_model_builder.model = MOCK_JUMPSTART_ID + mock_model_builder.model_hub = ModelHub.JUMPSTART + mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT + mock_model_builder.model_server = ModelServer.MMS + mock_model_builder.sagemaker_session.endpoint_arn = None + return mock_model_builder + + def _jumpstart_extra(self, latency): + return ( + f"{MODEL_BUILDER_DEPLOY_FUNC_NAME}" + f"&x-sdkVersion={SDK_VERSION}" + f"&x-env={PYTHON_VERSION}" + f"&x-sys={OS_NAME_VERSION}" + f"&x-platform={MOCK_STUDIO_APP_TYPE}" + f"&x-mode={Mode.SAGEMAKER_ENDPOINT}" + f"&x-jumpstartModelId={MOCK_JUMPSTART_ID}" + f"&x-latency={latency}" + ) + @patch("sagemaker.serve.utils.telemetry_logger._requests_helper") @patch("sagemaker.serve.utils.telemetry_logger._get_accountId") def test_log_sucessfully(self, mocked_get_accountId, mocked_request_helper): MOCK_SESSION.boto_session.region_name = "ap-south-1" mocked_get_accountId.return_value = "testAccountId" _send_telemetry("someStatus", 1, MOCK_SESSION) - mocked_request_helper.assert_has_calls( - [ - call( - "https://dev-exp-t-ap-south-1.s3.ap-south-1.amazonaws.com/" - "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", - 2, - ), - call( - "https://sm-pysdk-t-ap-south-1.s3.ap-south-1.amazonaws.com/" - "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", - 2, - ), - ] + mocked_request_helper.assert_called_with( + "https://dev-exp-t-ap-south-1.s3.ap-south-1.amazonaws.com/" + "telemetry?x-accountId=testAccountId&x-mode=1&x-status=someStatus", + 2, ) - self.assertEqual(mocked_request_helper.call_count, 2) @patch("sagemaker.serve.utils.telemetry_logger._get_accountId") def test_log_handle_exception(self, mocked_get_accountId): @@ -184,76 +207,114 @@ def test_capture_telemetry_decorator_tgi_success(self, mock_send_telemetry): ) @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") - def test_capture_telemetry_decorator_jumpstart_emits_model_id(self, mock_send_telemetry): + def test_capture_telemetry_decorator_no_call_when_disabled(self, mock_send_telemetry): mock_model_builder = ModelBuilderMock() - mock_model_builder.serve_settings.telemetry_opt_out = False - mock_model_builder.image_uri = MOCK_TGI_CONTAINER + mock_model_builder.serve_settings.telemetry_opt_out = True + mock_model_builder.image_uri = MOCK_DJL_CONTAINER mock_model_builder._is_custom_image_uri = False - mock_model_builder.model = MOCK_JUMPSTART_ID - mock_model_builder.model_hub = ModelHub.JUMPSTART - mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT - mock_model_builder.model_server = ModelServer.MMS - mock_model_builder.sagemaker_session.endpoint_arn = MOCK_ENDPOINT_ARN + mock_model_builder.model = MOCK_HUGGINGFACE_ID + mock_model_builder.model_server = ModelServer.DJL_SERVING mock_model_builder.mock_deploy() - args = mock_send_telemetry.call_args.args - latency = str(args[5]).split("latency=")[1] - expected_extra_str = ( - f"{MOCK_DEPLOY_FUNC_NAME}" - "&x-modelServer=2" - "&x-imageTag=huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04" - f"&x-sdkVersion={SDK_VERSION}" - f"&x-defaultImageUsage={ImageUriOption.DEFAULT_IMAGE.value}" - f"&x-endpointArn={MOCK_ENDPOINT_ARN}" - "&x-modelHub=1" - f"&x-jumpstartModelId={MOCK_JUMPSTART_ID}" - f"&x-latency={latency}" - ) + assert not mock_send_telemetry.called + + @patch("sagemaker.telemetry.telemetry_logging.process_studio_metadata_file") + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") + def test_capture_telemetry_decorator_jumpstart_emits_sdk_event( + self, mock_send_telemetry, mock_send_telemetry_request, mock_studio_metadata + ): + mock_studio_metadata.return_value = MOCK_STUDIO_APP_TYPE + mock_model_builder = self._jumpstart_model_builder() + mock_model_builder.deploy() + + serve_extra = mock_send_telemetry.call_args.args[5] + latency = serve_extra.split("latency=")[1] mock_send_telemetry.assert_called_once_with( - "1", 3, MOCK_SESSION, None, None, expected_extra_str + "1", + 3, + MOCK_SESSION, + None, + None, + f"{MODEL_BUILDER_DEPLOY_FUNC_NAME}&x-modelServer=2&x-sdkVersion={SDK_VERSION}" + f"&x-modelHub=1&x-latency={latency}", + ) + mock_send_telemetry_request.assert_called_once_with( + 1, + [8], + MOCK_SESSION, + None, + None, + self._jumpstart_extra(latency), ) + @patch("sagemaker.telemetry.telemetry_logging.process_studio_metadata_file") + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") - def test_capture_telemetry_decorator_huggingface_hub_omits_model_id(self, mock_send_telemetry): - mock_model_builder = ModelBuilderMock() - mock_model_builder.serve_settings.telemetry_opt_out = False - mock_model_builder.image_uri = None - mock_model_builder.model = MOCK_HUGGINGFACE_ID - mock_model_builder.model_hub = ModelHub.HUGGINGFACE - mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT - mock_model_builder.model_server = ModelServer.MMS - mock_model_builder.sagemaker_session.endpoint_arn = None + def test_capture_telemetry_decorator_jumpstart_sdk_event_on_failure( + self, mock_send_telemetry, mock_send_telemetry_request, mock_studio_metadata + ): + mock_studio_metadata.return_value = MOCK_STUDIO_APP_TYPE + mock_model_builder = self._jumpstart_model_builder() + mock_model_builder.sagemaker_session.local_mode = True + mock_exception = Mock() + mock_exception.side_effect = MOCK_EXCEPTION - mock_model_builder.mock_deploy() + with self.assertRaises(ModelBuilderException): + mock_model_builder.deploy(mock_exception) - args = mock_send_telemetry.call_args.args - latency = str(args[5]).split("latency=")[1] - expected_extra_str = ( - f"{MOCK_DEPLOY_FUNC_NAME}" - "&x-modelServer=2" - f"&x-sdkVersion={SDK_VERSION}" - "&x-modelHub=2" - f"&x-latency={latency}" + latency = mock_send_telemetry.call_args.args[5].split("latency=")[1] + mock_send_telemetry_request.assert_called_once_with( + 0, + [8, 2], + MOCK_SESSION, + str(MOCK_EXCEPTION), + MOCK_EXCEPTION.__class__.__name__, + self._jumpstart_extra(latency), ) - mock_send_telemetry.assert_called_once_with( - "1", 3, MOCK_SESSION, None, None, expected_extra_str - ) + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") + def test_capture_telemetry_decorator_jumpstart_other_method_sends_no_sdk_event( + self, mock_send_telemetry, mock_send_telemetry_request + ): + mock_model_builder = self._jumpstart_model_builder() + + mock_model_builder.mock_optimize() + assert "&x-modelHub=1" in mock_send_telemetry.call_args.args[5] + assert not mock_send_telemetry_request.called + + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") - def test_capture_telemetry_decorator_no_call_when_disabled(self, mock_send_telemetry): - mock_model_builder = ModelBuilderMock() - mock_model_builder.serve_settings.telemetry_opt_out = True - mock_model_builder.image_uri = MOCK_DJL_CONTAINER - mock_model_builder._is_custom_image_uri = False + def test_capture_telemetry_decorator_huggingface_hub_sends_no_sdk_event( + self, mock_send_telemetry, mock_send_telemetry_request + ): + mock_model_builder = self._jumpstart_model_builder() mock_model_builder.model = MOCK_HUGGINGFACE_ID - mock_model_builder.model_server = ModelServer.DJL_SERVING + mock_model_builder.model_hub = ModelHub.HUGGINGFACE - mock_model_builder.mock_deploy() + mock_model_builder.deploy() + + serve_extra = mock_send_telemetry.call_args.args[5] + assert "&x-modelHub=2" in serve_extra + assert "jumpstartModelId" not in serve_extra + assert not mock_send_telemetry_request.called + + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") + def test_capture_telemetry_decorator_jumpstart_opt_out_sends_no_sdk_event( + self, mock_send_telemetry, mock_send_telemetry_request + ): + mock_model_builder = self._jumpstart_model_builder() + mock_model_builder.serve_settings.telemetry_opt_out = True + + mock_model_builder.deploy() assert not mock_send_telemetry.called + assert not mock_send_telemetry_request.called @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") def test_capture_telemetry_decorator_handle_exception_success(self, mock_send_telemetry): @@ -304,16 +365,18 @@ def test_construct_url_with_failure_reason_and_extra_info(self): mock_extra_info = "mock_extra_info" mock_region = "us-west-2" - query = _construct_query( + ret_url = _construct_url( accountId=mock_accountId, mode=mock_mode, status=mock_status, failure_reason=mock_failure_reason, failure_type=mock_failure_type, extra_info=mock_extra_info, + region=mock_region, ) - expected_query = ( + expected_base_url = ( + f"https://dev-exp-t-{mock_region}.s3.{mock_region}.amazonaws.com/telemetry?" f"x-accountId={mock_accountId}" f"&x-mode={mock_mode}" f"&x-status={mock_status}" @@ -321,16 +384,7 @@ def test_construct_url_with_failure_reason_and_extra_info(self): f"&x-failureType={mock_failure_type}" f"&x-extra={mock_extra_info}" ) - self.assertEqual(query, expected_query) - self.assertEqual(TELEMETRY_BUCKET_PREFIXES, ("dev-exp-t", "sm-pysdk-t")) - self.assertEqual( - _construct_url("dev-exp-t", mock_region, query), - f"https://dev-exp-t-us-west-2.s3.us-west-2.amazonaws.com/telemetry?{expected_query}", - ) - self.assertEqual( - _construct_url("sm-pysdk-t", mock_region, query), - f"https://sm-pysdk-t-us-west-2.s3.us-west-2.amazonaws.com/telemetry?{expected_query}", - ) + self.assertEqual(ret_url, expected_base_url) @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") def test_capture_telemetry_decorator_mlflow_success(self, mock_send_telemetry): From d304998481ca6f26ff8c3922573b95c4509e0224 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Thu, 10 Sep 2026 01:46:03 +0000 Subject: [PATCH 4/5] test(serve): prove the JumpStart SDK event is additive Assert the unchanged ModelBuilder event to dev-exp-t next to the new JumpStart event on the failure path, and assert that the ModelBuilder event survives a failure of the SDK emitter. State the additive contract in the sender docstring. --- X-AI-Prompt: Can you make sure the old telemetry still sends and the new v3-compatible JumpStart event is additive? X-AI-Tool: claude-code --- src/sagemaker/serve/utils/telemetry_logger.py | 3 ++- .../serve/utils/test_telemetry_logger.py | 22 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py index 1f1cc9541c..9613721b01 100644 --- a/src/sagemaker/serve/utils/telemetry_logger.py +++ b/src/sagemaker/serve/utils/telemetry_logger.py @@ -233,7 +233,8 @@ def _send_jumpstart_telemetry(model_builder, func_name: str, outcome: _CallOutco The event has the emitter schema, so it lands in the same bucket and shape as the sagemaker.telemetry events and the v3 SDK events. Another model source or another method - sends nothing here. + sends nothing here. The event is additive: the ModelBuilder event to the old bucket is + sent first and stays as it was, and a failure here never gets to the caller. """ if func_name not in JUMPSTART_SDK_EVENT_FUNC_NAMES: return diff --git a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py index e280981fab..ad3bc69c19 100644 --- a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py +++ b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py @@ -266,6 +266,15 @@ def test_capture_telemetry_decorator_jumpstart_sdk_event_on_failure( mock_model_builder.deploy(mock_exception) latency = mock_send_telemetry.call_args.args[5].split("latency=")[1] + mock_send_telemetry.assert_called_once_with( + "0", + 3, + MOCK_SESSION, + str(MOCK_EXCEPTION), + MOCK_EXCEPTION.__class__.__name__, + f"{MODEL_BUILDER_DEPLOY_FUNC_NAME}&x-modelServer=2&x-sdkVersion={SDK_VERSION}" + f"&x-modelHub=1&x-latency={latency}", + ) mock_send_telemetry_request.assert_called_once_with( 0, [8, 2], @@ -275,6 +284,19 @@ def test_capture_telemetry_decorator_jumpstart_sdk_event_on_failure( self._jumpstart_extra(latency), ) + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") + def test_capture_telemetry_decorator_jumpstart_sdk_event_failure_keeps_legacy_event( + self, mock_send_telemetry, mock_send_telemetry_request + ): + mock_send_telemetry_request.side_effect = RuntimeError("sdk emitter down") + mock_model_builder = self._jumpstart_model_builder() + + mock_model_builder.deploy() + + assert mock_send_telemetry.call_count == 1 + assert "&x-modelHub=1" in mock_send_telemetry.call_args.args[5] + @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry_request") @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry") def test_capture_telemetry_decorator_jumpstart_other_method_sends_no_sdk_event( From 8ea9f274e91c86f521f6d47e7ede1950d25d2d7d Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Thu, 10 Sep 2026 02:47:15 +0000 Subject: [PATCH 5/5] test(serve): add an integ test for JumpStart ModelBuilder telemetry Build and deploy a real JumpStart model through ModelBuilder with the SageMaker writes and the telemetry GET requests intercepted. Assert the unchanged dev-exp-t event, the added sm-pysdk-t event with the emitter schema and the model ID, and that every sm-pysdk-t event carries x-feature=8 and no ModelBuilder key. --- X-AI-Prompt: Can you add an integ test that checks the old telemetry still sends and the new JumpStart event is additive? X-AI-Tool: claude-code --- .../serve/test_serve_js_telemetry.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 tests/integ/sagemaker/serve/test_serve_js_telemetry.py diff --git a/tests/integ/sagemaker/serve/test_serve_js_telemetry.py b/tests/integ/sagemaker/serve/test_serve_js_telemetry.py new file mode 100644 index 0000000000..2fb2aded27 --- /dev/null +++ b/tests/integ/sagemaker/serve/test_serve_js_telemetry.py @@ -0,0 +1,113 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Telemetry emitted by a real JumpStart ModelBuilder flow. + +The JumpStart metadata, the image URI, and the model artifact location come from the live +JumpStart catalog. SageMaker resource creation and the telemetry GET requests are intercepted, +so the test creates no endpoint and sends no telemetry. +""" +from __future__ import absolute_import + +from unittest.mock import Mock, patch +from urllib.parse import parse_qs, urlsplit + +import requests + +from sagemaker.session import Session +from sagemaker.serve.builder.model_builder import ModelBuilder +from sagemaker.serve.builder.schema_builder import SchemaBuilder +from sagemaker.user_agent import SDK_VERSION + +ROLE_NAME = "SageMakerRole" +JS_MMS_MODEL_ID = "huggingface-sentencesimilarity-bge-m3" +SAMPLE_MMS_PROMPT = ["How cute your dog is!"] +SAMPLE_MMS_RESPONSE = {"embedding": []} +LEGACY_BUCKET_PREFIX = "dev-exp-t-" +SDK_BUCKET_PREFIX = "sm-pysdk-t-" +JUMPSTART_FEATURE_CODE = "8" + + +def _telemetry_query(url): + return {key: values[0] for key, values in parse_qs(urlsplit(url).query).items()} + + +def _events(telemetry_requests, bucket_prefix, func_name): + return [ + _telemetry_query(url) + for url in telemetry_requests + if urlsplit(url).netloc.startswith(bucket_prefix) + and _telemetry_query(url).get("x-extra") == func_name + ] + + +def test_jumpstart_model_builder_emits_legacy_and_sdk_telemetry(sagemaker_session, account): + telemetry_requests = [] + + def record_telemetry_request(url, *args, **kwargs): + if "/telemetry?" in url: + telemetry_requests.append(url) + return Mock(status_code=200) + + with ( + patch.object(requests, "get", side_effect=record_telemetry_request), + patch.object(Session, "create_model", return_value="mock_model"), + patch.object(Session, "endpoint_from_production_variants", return_value="mock_endpoint"), + ): + model_builder = ModelBuilder( + model=JS_MMS_MODEL_ID, + schema_builder=SchemaBuilder(SAMPLE_MMS_PROMPT, SAMPLE_MMS_RESPONSE), + role_arn=f"arn:aws:iam::{account}:role/{ROLE_NAME}", + sagemaker_session=sagemaker_session, + instance_type="ml.m5.xlarge", + ) + model = model_builder.build() + model.deploy(instance_type="ml.m5.xlarge", endpoint_logging=False) + + legacy_build_events = _events(telemetry_requests, LEGACY_BUCKET_PREFIX, "ModelBuilder.build") + assert len(legacy_build_events) == 1 + legacy_build_event = legacy_build_events[0] + assert legacy_build_event["x-accountId"] == account + assert legacy_build_event["x-mode"] == "3" + assert legacy_build_event["x-modelHub"] == "1" + assert legacy_build_event["x-sdkVersion"] == SDK_VERSION + assert "x-feature" not in legacy_build_event + assert "x-jumpstartModelId" not in legacy_build_event + + sdk_build_events = _events(telemetry_requests, SDK_BUCKET_PREFIX, "ModelBuilder.build") + assert len(sdk_build_events) == 1 + sdk_build_event = sdk_build_events[0] + assert sdk_build_event["x-accountId"] == account + assert sdk_build_event["x-feature"] == JUMPSTART_FEATURE_CODE + assert sdk_build_event["x-jumpstartModelId"] == JS_MMS_MODEL_ID + assert sdk_build_event["x-mode"] == "SAGEMAKER_ENDPOINT" + assert sdk_build_event["x-sdkVersion"] == SDK_VERSION + assert {"x-env", "x-sys", "x-platform", "x-latency"} <= set(sdk_build_event) + assert "x-modelHub" not in sdk_build_event + + sdk_deploy_events = _events(telemetry_requests, SDK_BUCKET_PREFIX, "jumpstart_model.deploy") + assert len(sdk_deploy_events) == 1 + assert sdk_deploy_events[0]["x-feature"] == JUMPSTART_FEATURE_CODE + assert sdk_deploy_events[0]["x-jumpstartModelId"] == JS_MMS_MODEL_ID + + legacy_deploy_events = _events(telemetry_requests, LEGACY_BUCKET_PREFIX, "jumpstart.deploy") + assert len(legacy_deploy_events) == 1 + assert legacy_deploy_events[0]["x-modelHub"] == "1" + assert "x-jumpstartModelId" not in legacy_deploy_events[0] + + sdk_bucket_events = [ + _telemetry_query(url) + for url in telemetry_requests + if urlsplit(url).netloc.startswith(SDK_BUCKET_PREFIX) + ] + assert all(event["x-feature"] == JUMPSTART_FEATURE_CODE for event in sdk_bucket_events) + assert all("x-modelHub" not in event for event in sdk_bucket_events)