Skip to content

Commit b2b6231

Browse files
committed
Merge remote-tracking branch 'origin/master' into fix/jumpstart-neo-bucket-v3
2 parents a4245fb + 8e7485a commit b2b6231

5 files changed

Lines changed: 159 additions & 2 deletions

File tree

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ class TelemetryParamType:
192192

193193
# Calls self.<name>() and emits the return value.
194194
# Use for: computed/derived values like _is_model_customization(), _is_nova_model().
195+
# Emits nothing if the method returns None.
195196
ATTR_CALL = "attr_call"
196197

197198
# 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:
246247
method = getattr(instance, name, None)
247248
if callable(method):
248249
try:
249-
parts.append(f"&x-{key}={method()}")
250+
value = method()
250251
except Exception:
251-
pass
252+
value = None
253+
if value is not None:
254+
parts.append(f"&x-{key}={value}")
252255
elif kind == T.KWARG_VALUE:
253256
value = kwargs.get(name) if kwargs else None
254257
if value is not None:

‎sagemaker-core/tests/unit/telemetry/test_granular_telemetry.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,14 @@ def test_attr_call_skips_on_exception(self):
9898
])
9999
assert "isModelCustomization" not in result
100100

101+
def test_attr_call_skips_none(self):
102+
instance = self._make_instance()
103+
instance._jumpstart_model_id = Mock(return_value=None)
104+
result = _extract_telemetry_params(instance, {}, [
105+
("_jumpstart_model_id", TelemetryParamType.ATTR_CALL),
106+
])
107+
assert "jumpstartModelId" not in result
108+
101109
def test_kwarg_value_emits_value(self):
102110
instance = self._make_instance()
103111
result = _extract_telemetry_params(instance, {"instance_type": "ml.g5.2xlarge"}, [

‎sagemaker-serve/src/sagemaker/serve/model_builder.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4289,6 +4289,7 @@ def _reset_build_state(self):
42894289
telemetry_params=[
42904290
("mode", TelemetryParamType.ATTR_VALUE),
42914291
("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL),
4292+
("_jumpstart_model_id", TelemetryParamType.ATTR_CALL),
42924293
("network", TelemetryParamType.ATTR_EXISTS),
42934294
("source_code", TelemetryParamType.ATTR_EXISTS),
42944295
("inference_spec", TelemetryParamType.ATTR_EXISTS),
@@ -5805,6 +5806,7 @@ def _deploy_recommendation(
58055806
("instance_type", TelemetryParamType.ATTR_VALUE),
58065807
("_is_model_customization", TelemetryParamType.ATTR_CALL),
58075808
("_is_nova_model_for_telemetry", TelemetryParamType.ATTR_CALL),
5809+
("_jumpstart_model_id", TelemetryParamType.ATTR_CALL),
58085810
("network", TelemetryParamType.ATTR_EXISTS),
58095811
("compute", TelemetryParamType.ATTR_EXISTS),
58105812
("update_endpoint", TelemetryParamType.KWARG_EXISTS),

‎sagemaker-serve/src/sagemaker/serve/model_builder_utils.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2924,6 +2924,12 @@ def _is_jumpstart_model_id(self) -> bool:
29242924

29252925
return self._cached_is_jumpstart
29262926

2927+
def _jumpstart_model_id(self) -> Optional[str]:
2928+
"""Return the JumpStart model ID, or None for another model source."""
2929+
if isinstance(self.model, str) and self._is_jumpstart_model_id():
2930+
return self.model
2931+
return None
2932+
29272933
def _has_nvidia_gpu(self) -> bool:
29282934
try:
29292935
_get_available_gpus()
Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
"""Unit tests for the JumpStart model ID in ModelBuilder telemetry.
2+
3+
The model ID lets usage analytics identify the JumpStart model behind each
4+
``model_builder.build`` and ``model_builder.deploy`` event.
5+
"""
6+
7+
from __future__ import absolute_import
8+
9+
import unittest
10+
from unittest.mock import Mock, patch
11+
12+
from sagemaker.core.resources import Endpoint, Model
13+
from sagemaker.serve.model_builder import ModelBuilder
14+
from sagemaker.serve.utils.types import ModelServer
15+
16+
TELEMETRY_MODULE = "sagemaker.core.telemetry.telemetry_logging"
17+
JUMPSTART_MODEL_ID = "huggingface-llm-falcon-7b-bf16"
18+
19+
20+
def _telemetry_extra(mock_send_telemetry):
21+
"""Return the extra info string of the last telemetry request."""
22+
return mock_send_telemetry.call_args.args[5]
23+
24+
25+
@patch(f"{TELEMETRY_MODULE}.resolve_value_from_config", return_value=False)
26+
@patch(f"{TELEMETRY_MODULE}._send_telemetry_request")
27+
class TestJumpStartTelemetry(unittest.TestCase):
28+
"""Tests for the x-jumpstartModelId telemetry param."""
29+
30+
def setUp(self):
31+
"""Set up test fixtures."""
32+
self.mock_session = Mock()
33+
self.mock_session.boto_region_name = "us-west-2"
34+
self.mock_session.boto_session = Mock()
35+
self.mock_session.boto_session.region_name = "us-west-2"
36+
self.mock_session.config = {}
37+
self.mock_session.sagemaker_config = {}
38+
self.mock_session.local_mode = False
39+
self.mock_session.default_bucket.return_value = "test-bucket"
40+
self.mock_session.default_bucket_prefix = "test-prefix"
41+
42+
self.mock_client = Mock()
43+
self.mock_client._user_agent_creator = Mock()
44+
self.mock_client._user_agent_creator.to_string = Mock(return_value="test-agent")
45+
self.mock_session.sagemaker_client = self.mock_client
46+
47+
self.mock_role_arn = "arn:aws:iam::123456789012:role/TestRole"
48+
49+
def _make_builder(self, model):
50+
"""Create a ModelBuilder that reports no built model."""
51+
builder = ModelBuilder(
52+
model=model,
53+
role_arn=self.mock_role_arn,
54+
sagemaker_session=self.mock_session,
55+
model_server=ModelServer.TORCHSERVE,
56+
)
57+
builder.built_model = None
58+
return builder
59+
60+
@patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id")
61+
@patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder")
62+
@patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting")
63+
def test_build_emits_the_model_id_for_a_jumpstart_model(
64+
self,
65+
mock_serve_setting,
66+
mock_build_single,
67+
mock_is_jumpstart,
68+
mock_send_telemetry,
69+
mock_resolve_config,
70+
):
71+
"""build() emits the model ID for a JumpStart model ID."""
72+
mock_serve_setting.return_value = Mock()
73+
mock_build_single.return_value = Mock(spec=Model)
74+
mock_is_jumpstart.return_value = True
75+
76+
self._make_builder(JUMPSTART_MODEL_ID).build()
77+
78+
extra = _telemetry_extra(mock_send_telemetry)
79+
assert f"&x-jumpstartModelId={JUMPSTART_MODEL_ID}" in extra
80+
assert "&x-isJumpstartModelId=" not in extra
81+
82+
@patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id")
83+
@patch("sagemaker.serve.model_builder.ModelBuilder._build_single_modelbuilder")
84+
@patch("sagemaker.serve.model_builder.ModelBuilder._get_serve_setting")
85+
def test_build_emits_no_model_id_for_other_model(
86+
self,
87+
mock_serve_setting,
88+
mock_build_single,
89+
mock_is_jumpstart,
90+
mock_send_telemetry,
91+
mock_resolve_config,
92+
):
93+
"""build() emits no model ID for a model that is not from JumpStart."""
94+
mock_serve_setting.return_value = Mock()
95+
mock_build_single.return_value = Mock(spec=Model)
96+
mock_is_jumpstart.return_value = False
97+
98+
self._make_builder(Mock()).build()
99+
100+
assert "&x-jumpstartModelId=" not in _telemetry_extra(mock_send_telemetry)
101+
102+
@patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id")
103+
@patch("sagemaker.serve.model_builder.ModelBuilder._deploy")
104+
def test_deploy_emits_the_model_id_for_a_jumpstart_model(
105+
self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config
106+
):
107+
"""deploy() emits the model ID for a JumpStart model ID."""
108+
mock_deploy.return_value = Mock(spec=Endpoint)
109+
mock_is_jumpstart.return_value = True
110+
111+
builder = self._make_builder(JUMPSTART_MODEL_ID)
112+
builder.built_model = Mock(spec=Model)
113+
builder.instance_type = "ml.g5.2xlarge"
114+
builder.deploy(endpoint_name="test-endpoint", wait=False)
115+
116+
extra = _telemetry_extra(mock_send_telemetry)
117+
assert f"&x-jumpstartModelId={JUMPSTART_MODEL_ID}" in extra
118+
assert "&x-isJumpstartModelId=" not in extra
119+
120+
@patch("sagemaker.serve.model_builder.ModelBuilder._is_jumpstart_model_id")
121+
@patch("sagemaker.serve.model_builder.ModelBuilder._deploy")
122+
def test_deploy_emits_no_model_id_for_other_model(
123+
self, mock_deploy, mock_is_jumpstart, mock_send_telemetry, mock_resolve_config
124+
):
125+
"""deploy() emits no model ID for a model that is not from JumpStart."""
126+
mock_deploy.return_value = Mock(spec=Endpoint)
127+
mock_is_jumpstart.return_value = False
128+
129+
builder = self._make_builder(Mock())
130+
builder.built_model = Mock(spec=Model)
131+
builder.instance_type = "ml.g5.2xlarge"
132+
builder.deploy(endpoint_name="test-endpoint", wait=False)
133+
134+
assert "&x-jumpstartModelId=" not in _telemetry_extra(mock_send_telemetry)
135+
136+
137+
if __name__ == "__main__":
138+
unittest.main()

0 commit comments

Comments
 (0)