From 76d460a2a41f0186fba3aab3c226718c1765f224 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Tue, 1 Sep 2026 18:36:24 +0000 Subject: [PATCH 1/4] change: emit a JumpStart flag in ModelBuilder telemetry Record whether ModelBuilder build and deploy calls use a JumpStart model ID. This flag lets downstream analytics identify JumpStart usage. --- X-AI-Prompt: Can you add a JumpStart identity flag to ModelBuilder telemetry and raise a pull request? X-AI-Tool: claude-code --- .../src/sagemaker/serve/model_builder.py | 2 + .../unit/test_jumpstart_telemetry_flag.py | 149 ++++++++++++++++++ 2 files changed, 151 insertions(+) create mode 100644 sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 0d1e6e4746..998d4dfefd 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4289,6 +4289,7 @@ def _reset_build_state(self): telemetry_params=[ ("mode", TelemetryParamType.ATTR_VALUE), ("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL), + ("_is_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("network", TelemetryParamType.ATTR_EXISTS), ("source_code", TelemetryParamType.ATTR_EXISTS), ("inference_spec", TelemetryParamType.ATTR_EXISTS), @@ -5805,6 +5806,7 @@ def _deploy_recommendation( ("instance_type", TelemetryParamType.ATTR_VALUE), ("_is_model_customization", TelemetryParamType.ATTR_CALL), ("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL), + ("_is_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("network", TelemetryParamType.ATTR_EXISTS), ("compute", TelemetryParamType.ATTR_EXISTS), ("update_endpoint", TelemetryParamType.KWARG_EXISTS), diff --git a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py new file mode 100644 index 0000000000..b5b0d8c837 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py @@ -0,0 +1,149 @@ +"""Unit tests for the JumpStart flag on ModelBuilder build and deploy telemetry. + +The flag lets usage analytics separate JumpStart deployments from other +``model_builder.build`` and ``model_builder.deploy`` events. +""" + +import unittest +from unittest.mock import Mock, patch + +from sagemaker.core.resources import Endpoint, Model +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.serve.utils.types import ModelServer + +TELEMETRY_MODULE = "sagemaker.core.telemetry.telemetry_logging" +JUMPSTART_MODEL_ID = "huggingface-llm-falcon-7b-bf16" + + +def _telemetry_extra(mock_send_telemetry): + """Return the extra info string of the last telemetry request.""" + return mock_send_telemetry.call_args.args[5] + + +@patch(f"{TELEMETRY_MODULE}.resolve_value_from_config", return_value=False) +@patch(f"{TELEMETRY_MODULE}._send_telemetry_request") +class TestJumpStartTelemetryFlag(unittest.TestCase): + """Tests for the x-isJumpstartModelId telemetry param.""" + + def setUp(self): + """Set up test fixtures.""" + self.mock_session = Mock() + self.mock_session.boto_region_name = "us-west-2" + self.mock_session.boto_session = Mock() + self.mock_session.boto_session.region_name = "us-west-2" + self.mock_session.config = {} + self.mock_session.sagemaker_config = {} + self.mock_session.local_mode = False + self.mock_session.default_bucket.return_value = "test-bucket" + self.mock_session.default_bucket_prefix = "test-prefix" + + self.mock_client = Mock() + self.mock_client._user_agent_creator = Mock() + self.mock_client._user_agent_creator.to_string = Mock(return_value="test-agent") + self.mock_session.sagemaker_client = self.mock_client + + self.mock_role_arn = "arn:aws:iam::123456789012:role/TestRole" + + def _make_builder(self, model): + """Create a ModelBuilder that reports no built model.""" + builder = ModelBuilder( + model=model, + role_arn=self.mock_role_arn, + sagemaker_session=self.mock_session, + model_server=ModelServer.TORCHSERVE, + ) + builder.built_model = None + return builder + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") + @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") + def test_build_emits_true_for_jumpstart_model_id( + self, + mock_serve_setting, + mock_build_single, + mock_is_jumpstart, + mock_send_telemetry, + mock_resolve_config, + ): + """build() emits the flag as True for a JumpStart model ID.""" + mock_serve_setting.return_value = Mock() + mock_build_single.return_value = Mock(spec=Model) + mock_is_jumpstart.return_value = True + + self._make_builder(JUMPSTART_MODEL_ID).build() + + assert "&x-isJumpstartModelId=True" in _telemetry_extra(mock_send_telemetry) + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") + @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") + def test_build_emits_false_for_other_model( + self, + mock_serve_setting, + mock_build_single, + mock_is_jumpstart, + mock_send_telemetry, + mock_resolve_config, + ): + """build() emits the flag as False for a model that is not from JumpStart.""" + mock_serve_setting.return_value = Mock() + mock_build_single.return_value = Mock(spec=Model) + mock_is_jumpstart.return_value = False + + self._make_builder(Mock()).build() + + assert "&x-isJumpstartModelId=False" in _telemetry_extra(mock_send_telemetry) + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") + def test_deploy_emits_true_for_jumpstart_model_id( + self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config + ): + """deploy() emits the flag as True for a JumpStart model ID.""" + mock_deploy.return_value = Mock(spec=Endpoint) + mock_is_jumpstart.return_value = True + + builder = self._make_builder(JUMPSTART_MODEL_ID) + builder.built_model = Mock(spec=Model) + builder.instance_type = "ml.g5.2xlarge" + builder.deploy(endpoint_name="test-endpoint", wait=False) + + assert "&x-isJumpstartModelId=True" in _telemetry_extra(mock_send_telemetry) + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") + def test_deploy_emits_false_for_other_model( + self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config + ): + """deploy() emits the flag as False for a model that is not from JumpStart.""" + mock_deploy.return_value = Mock(spec=Endpoint) + mock_is_jumpstart.return_value = False + + builder = self._make_builder(Mock()) + builder.built_model = Mock(spec=Model) + builder.instance_type = "ml.g5.2xlarge" + builder.deploy(endpoint_name="test-endpoint", wait=False) + + assert "&x-isJumpstartModelId=False" in _telemetry_extra(mock_send_telemetry) + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") + def test_deploy_emits_the_flag_before_the_latency_param( + self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config + ): + """The flag arrives before x-latency, so a greedy field parser reads it.""" + mock_deploy.return_value = Mock(spec=Endpoint) + mock_is_jumpstart.return_value = True + + builder = self._make_builder(JUMPSTART_MODEL_ID) + builder.built_model = Mock(spec=Model) + builder.instance_type = "ml.g5.2xlarge" + builder.deploy(endpoint_name="test-endpoint", wait=False) + + extra = _telemetry_extra(mock_send_telemetry) + assert extra.index("&x-isJumpstartModelId=") < extra.index("&x-latency=") + + +if __name__ == "__main__": + unittest.main() From 430242122bf4961aafa31226f0a744a9c5bbb7ea Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Tue, 1 Sep 2026 20:51:33 +0000 Subject: [PATCH 2/4] change: emit the JumpStart model ID in ModelBuilder telemetry The isJumpstartModelId flag separates JumpStart traffic from other traffic, but it does not name the model. Analytics cannot rank JumpStart models by build count or deployment count from a boolean. Add a jumpstartModelId param to the build and deploy telemetry param lists. The param emits the model ID string when the model source is a JumpStart model ID, and emits nothing for another model source. The ATTR_CALL branch of the param extractor drops a None return, so a param that has no value stays out of the beacon. --- X-AI-Prompt: Can we track the JumpStart model ID itself in telemetry instead of only a boolean flag? X-AI-Tool: claude-code --- .../core/telemetry/telemetry_logging.py | 7 +- .../unit/telemetry/test_granular_telemetry.py | 8 ++ .../src/sagemaker/serve/model_builder.py | 2 + .../sagemaker/serve/model_builder_utils.py | 6 ++ .../unit/test_jumpstart_telemetry_flag.py | 81 ++++++++++++++++++- 5 files changed, 99 insertions(+), 5 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index d55a390056..49ba16c543 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -192,6 +192,7 @@ class TelemetryParamType: # Calls self.() and emits the return value. # Use for: computed/derived values like _is_model_customization(), _is_nova_model(). + # Emits nothing if the method returns None. ATTR_CALL = "attr_call" # Reads kwargs[] from the decorated method's keyword arguments and emits the value. @@ -246,9 +247,11 @@ def _extract_telemetry_params(instance, kwargs, telemetry_params=None) -> str: method = getattr(instance, name, None) if callable(method): try: - parts.append(f"&x-{key}={method()}") + value = method() except Exception: - pass + value = None + if value is not None: + parts.append(f"&x-{key}={value}") elif kind == T.KWARG_VALUE: value = kwargs.get(name) if kwargs else None if value is not None: diff --git a/sagemaker-core/tests/unit/telemetry/test_granular_telemetry.py b/sagemaker-core/tests/unit/telemetry/test_granular_telemetry.py index d216d323da..4c6fb2fea6 100644 --- a/sagemaker-core/tests/unit/telemetry/test_granular_telemetry.py +++ b/sagemaker-core/tests/unit/telemetry/test_granular_telemetry.py @@ -98,6 +98,14 @@ def test_attr_call_skips_on_exception(self): ]) assert "isModelCustomization" not in result + def test_attr_call_skips_none(self): + instance = self._make_instance() + instance._jumpstart_model_id = Mock(return_value=None) + result = _extract_telemetry_params(instance, {}, [ + ("_jumpstart_model_id", TelemetryParamType.ATTR_CALL), + ]) + assert "jumpstartModelId" not in result + def test_kwarg_value_emits_value(self): instance = self._make_instance() result = _extract_telemetry_params(instance, {"instance_type": "ml.g5.2xlarge"}, [ diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 998d4dfefd..ccf2abe40c 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4290,6 +4290,7 @@ def _reset_build_state(self): ("mode", TelemetryParamType.ATTR_VALUE), ("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL), ("_is_jumpstart_model_id", TelemetryParamType.ATTR_CALL), + ("_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("network", TelemetryParamType.ATTR_EXISTS), ("source_code", TelemetryParamType.ATTR_EXISTS), ("inference_spec", TelemetryParamType.ATTR_EXISTS), @@ -5807,6 +5808,7 @@ def _deploy_recommendation( ("_is_model_customization", TelemetryParamType.ATTR_CALL), ("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL), ("_is_jumpstart_model_id", TelemetryParamType.ATTR_CALL), + ("_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("network", TelemetryParamType.ATTR_EXISTS), ("compute", TelemetryParamType.ATTR_EXISTS), ("update_endpoint", TelemetryParamType.KWARG_EXISTS), diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index 1ea1265f53..68217dfde5 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py @@ -2924,6 +2924,12 @@ def _is_jumpstart_model_id(self) -> bool: return self._cached_is_jumpstart + def _jumpstart_model_id(self) -> Optional[str]: + """Return the JumpStart model ID, or None for another model source.""" + if isinstance(self.model, str) and self._is_jumpstart_model_id(): + return self.model + return None + def _has_nvidia_gpu(self) -> bool: try: _get_available_gpus() diff --git a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py index b5b0d8c837..958a11df2d 100644 --- a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py +++ b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py @@ -1,7 +1,8 @@ -"""Unit tests for the JumpStart flag on ModelBuilder build and deploy telemetry. +"""Unit tests for the JumpStart params on ModelBuilder build and deploy telemetry. The flag lets usage analytics separate JumpStart deployments from other -``model_builder.build`` and ``model_builder.deploy`` events. +``model_builder.build`` and ``model_builder.deploy`` events. The model ID +names the JumpStart model behind each of those events. """ import unittest @@ -23,7 +24,7 @@ def _telemetry_extra(mock_send_telemetry): @patch(f"{TELEMETRY_MODULE}.resolve_value_from_config", return_value=False) @patch(f"{TELEMETRY_MODULE}._send_telemetry_request") class TestJumpStartTelemetryFlag(unittest.TestCase): - """Tests for the x-isJumpstartModelId telemetry param.""" + """Tests for the x-isJumpstartModelId and x-jumpstartModelId telemetry params.""" def setUp(self): """Set up test fixtures.""" @@ -127,6 +128,80 @@ def test_deploy_emits_false_for_other_model( assert "&x-isJumpstartModelId=False" in _telemetry_extra(mock_send_telemetry) + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") + @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") + def test_build_emits_the_model_id_for_a_jumpstart_model( + self, + mock_serve_setting, + mock_build_single, + mock_is_jumpstart, + mock_send_telemetry, + mock_resolve_config, + ): + """build() emits the model ID for a JumpStart model ID.""" + mock_serve_setting.return_value = Mock() + mock_build_single.return_value = Mock(spec=Model) + mock_is_jumpstart.return_value = True + + self._make_builder(JUMPSTART_MODEL_ID).build() + + extra = _telemetry_extra(mock_send_telemetry) + assert f"&x-jumpstartModelId={JUMPSTART_MODEL_ID}" in extra + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") + @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") + def test_build_emits_no_model_id_for_other_model( + self, + mock_serve_setting, + mock_build_single, + mock_is_jumpstart, + mock_send_telemetry, + mock_resolve_config, + ): + """build() emits no model ID for a model that is not from JumpStart.""" + mock_serve_setting.return_value = Mock() + mock_build_single.return_value = Mock(spec=Model) + mock_is_jumpstart.return_value = False + + self._make_builder(Mock()).build() + + assert "&x-jumpstartModelId=" not in _telemetry_extra(mock_send_telemetry) + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") + def test_deploy_emits_the_model_id_for_a_jumpstart_model( + self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config + ): + """deploy() emits the model ID for a JumpStart model ID.""" + mock_deploy.return_value = Mock(spec=Endpoint) + mock_is_jumpstart.return_value = True + + builder = self._make_builder(JUMPSTART_MODEL_ID) + builder.built_model = Mock(spec=Model) + builder.instance_type = "ml.g5.2xlarge" + builder.deploy(endpoint_name="test-endpoint", wait=False) + + extra = _telemetry_extra(mock_send_telemetry) + assert f"&x-jumpstartModelId={JUMPSTART_MODEL_ID}" in extra + + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") + @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") + def test_deploy_emits_no_model_id_for_other_model( + self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config + ): + """deploy() emits no model ID for a model that is not from JumpStart.""" + mock_deploy.return_value = Mock(spec=Endpoint) + mock_is_jumpstart.return_value = False + + builder = self._make_builder(Mock()) + builder.built_model = Mock(spec=Model) + builder.instance_type = "ml.g5.2xlarge" + builder.deploy(endpoint_name="test-endpoint", wait=False) + + assert "&x-jumpstartModelId=" not in _telemetry_extra(mock_send_telemetry) + @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") def test_deploy_emits_the_flag_before_the_latency_param( From fd5329c86c6bd8c7c02703eb25c891b58283b258 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Tue, 1 Sep 2026 21:42:10 +0000 Subject: [PATCH 3/4] change: remove the JumpStart boolean telemetry param The model ID identifies JumpStart traffic and the exact model. The boolean param duplicates this information and adds a second field for consumers. Remove the boolean param from ModelBuilder telemetry and keep only jumpstartModelId. --- X-AI-Prompt: Can you remove the boolean flag and keep only the JumpStart model ID telemetry param? X-AI-Tool: claude-code --- .../src/sagemaker/serve/model_builder.py | 2 - .../unit/test_jumpstart_telemetry_flag.py | 103 ++---------------- 2 files changed, 7 insertions(+), 98 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index ccf2abe40c..cff687849e 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4289,7 +4289,6 @@ def _reset_build_state(self): telemetry_params=[ ("mode", TelemetryParamType.ATTR_VALUE), ("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL), - ("_is_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("network", TelemetryParamType.ATTR_EXISTS), ("source_code", TelemetryParamType.ATTR_EXISTS), @@ -5807,7 +5806,6 @@ def _deploy_recommendation( ("instance_type", TelemetryParamType.ATTR_VALUE), ("_is_model_customization", TelemetryParamType.ATTR_CALL), ("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL), - ("_is_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("_jumpstart_model_id", TelemetryParamType.ATTR_CALL), ("network", TelemetryParamType.ATTR_EXISTS), ("compute", TelemetryParamType.ATTR_EXISTS), diff --git a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py index 958a11df2d..21b321849f 100644 --- a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py +++ b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py @@ -1,8 +1,7 @@ -"""Unit tests for the JumpStart params on ModelBuilder build and deploy telemetry. +"""Unit tests for the JumpStart model ID in ModelBuilder telemetry. -The flag lets usage analytics separate JumpStart deployments from other -``model_builder.build`` and ``model_builder.deploy`` events. The model ID -names the JumpStart model behind each of those events. +The model ID lets usage analytics identify the JumpStart model behind each +``model_builder.build`` and ``model_builder.deploy`` event. """ import unittest @@ -23,8 +22,8 @@ def _telemetry_extra(mock_send_telemetry): @patch(f"{TELEMETRY_MODULE}.resolve_value_from_config", return_value=False) @patch(f"{TELEMETRY_MODULE}._send_telemetry_request") -class TestJumpStartTelemetryFlag(unittest.TestCase): - """Tests for the x-isJumpstartModelId and x-jumpstartModelId telemetry params.""" +class TestJumpStartTelemetry(unittest.TestCase): + """Tests for the x-jumpstartModelId telemetry param.""" def setUp(self): """Set up test fixtures.""" @@ -56,78 +55,6 @@ def _make_builder(self, model): builder.built_model = None return builder - @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") - @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") - @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") - def test_build_emits_true_for_jumpstart_model_id( - self, - mock_serve_setting, - mock_build_single, - mock_is_jumpstart, - mock_send_telemetry, - mock_resolve_config, - ): - """build() emits the flag as True for a JumpStart model ID.""" - mock_serve_setting.return_value = Mock() - mock_build_single.return_value = Mock(spec=Model) - mock_is_jumpstart.return_value = True - - self._make_builder(JUMPSTART_MODEL_ID).build() - - assert "&x-isJumpstartModelId=True" in _telemetry_extra(mock_send_telemetry) - - @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") - @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") - @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") - def test_build_emits_false_for_other_model( - self, - mock_serve_setting, - mock_build_single, - mock_is_jumpstart, - mock_send_telemetry, - mock_resolve_config, - ): - """build() emits the flag as False for a model that is not from JumpStart.""" - mock_serve_setting.return_value = Mock() - mock_build_single.return_value = Mock(spec=Model) - mock_is_jumpstart.return_value = False - - self._make_builder(Mock()).build() - - assert "&x-isJumpstartModelId=False" in _telemetry_extra(mock_send_telemetry) - - @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") - @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") - def test_deploy_emits_true_for_jumpstart_model_id( - self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config - ): - """deploy() emits the flag as True for a JumpStart model ID.""" - mock_deploy.return_value = Mock(spec=Endpoint) - mock_is_jumpstart.return_value = True - - builder = self._make_builder(JUMPSTART_MODEL_ID) - builder.built_model = Mock(spec=Model) - builder.instance_type = "ml.g5.2xlarge" - builder.deploy(endpoint_name="test-endpoint", wait=False) - - assert "&x-isJumpstartModelId=True" in _telemetry_extra(mock_send_telemetry) - - @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") - @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") - def test_deploy_emits_false_for_other_model( - self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config - ): - """deploy() emits the flag as False for a model that is not from JumpStart.""" - mock_deploy.return_value = Mock(spec=Endpoint) - mock_is_jumpstart.return_value = False - - builder = self._make_builder(Mock()) - builder.built_model = Mock(spec=Model) - builder.instance_type = "ml.g5.2xlarge" - builder.deploy(endpoint_name="test-endpoint", wait=False) - - assert "&x-isJumpstartModelId=False" in _telemetry_extra(mock_send_telemetry) - @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") @patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting") @@ -148,6 +75,7 @@ def test_build_emits_the_model_id_for_a_jumpstart_model( extra = _telemetry_extra(mock_send_telemetry) assert f"&x-jumpstartModelId={JUMPSTART_MODEL_ID}" in extra + assert "&x-isJumpstartModelId=" not in extra @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") @patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder") @@ -185,6 +113,7 @@ def test_deploy_emits_the_model_id_for_a_jumpstart_model( extra = _telemetry_extra(mock_send_telemetry) assert f"&x-jumpstartModelId={JUMPSTART_MODEL_ID}" in extra + assert "&x-isJumpstartModelId=" not in extra @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") @@ -202,23 +131,5 @@ def test_deploy_emits_no_model_id_for_other_model( assert "&x-jumpstartModelId=" not in _telemetry_extra(mock_send_telemetry) - @patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id") - @patch("sagemaker.serve.model_builder.ModelBuilder._deploy") - def test_deploy_emits_the_flag_before_the_latency_param( - self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config - ): - """The flag arrives before x-latency, so a greedy field parser reads it.""" - mock_deploy.return_value = Mock(spec=Endpoint) - mock_is_jumpstart.return_value = True - - builder = self._make_builder(JUMPSTART_MODEL_ID) - builder.built_model = Mock(spec=Model) - builder.instance_type = "ml.g5.2xlarge" - builder.deploy(endpoint_name="test-endpoint", wait=False) - - extra = _telemetry_extra(mock_send_telemetry) - assert extra.index("&x-isJumpstartModelId=") < extra.index("&x-latency=") - - if __name__ == "__main__": unittest.main() From 4417e185703eca336e7460fa55dd80766c8ca878 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Wed, 2 Sep 2026 18:51:31 +0000 Subject: [PATCH 4/4] fix(serve): obey telemetry test lint rules Add the necessary future import and the blank line for the lint gate. --- X-AI-Prompt: Can you fix the lint errors in my pull request and push the fix? X-AI-Tool: claude-code --- sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py index 21b321849f..93146c24ab 100644 --- a/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py +++ b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py @@ -4,6 +4,8 @@ ``model_builder.build`` and ``model_builder.deploy`` event. """ +from __future__ import absolute_import + import unittest from unittest.mock import Mock, patch @@ -131,5 +133,6 @@ def test_deploy_emits_no_model_id_for_other_model( assert "&x-jumpstartModelId=" not in _telemetry_extra(mock_send_telemetry) + if __name__ == "__main__": unittest.main()