change: emit the JumpStart model ID in ModelBuilder telemetry

sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py

@@ -192,6 +192,7 @@ class TelemetryParamType:
 
     # Calls self.<name>() 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[<name>] 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:

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"}, [

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),

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()

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()