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 0d1e6e4746..cff687849e 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), + ("_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), + ("_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 new file mode 100644 index 0000000000..93146c24ab --- /dev/null +++ b/sagemaker-serve/tests/unit/test_jumpstart_telemetry_flag.py @@ -0,0 +1,138 @@ +"""Unit tests for the JumpStart model ID in ModelBuilder telemetry. + +The model ID lets usage analytics identify the JumpStart model behind each +``model_builder.build`` and ``model_builder.deploy`` event. +""" + +from __future__ import absolute_import + +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 TestJumpStartTelemetry(unittest.TestCase): + """Tests for the x-jumpstartModelId 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_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 + 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") + @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 + assert "&x-isJumpstartModelId=" not 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) + + +if __name__ == "__main__": + unittest.main()