From d2fa3af87f501c25e196a81a275d8f3e0a641c33 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Mon, 6 Apr 2026 16:47:11 -0700 Subject: [PATCH 1/2] feat: split MODEL_CUSTOMIZATION telemetry into NOVA and OSS sub-features --- .../src/sagemaker/core/telemetry/constants.py | 2 ++ .../core/telemetry/telemetry_logging.py | 21 +++++++++++++++++++ .../sagemaker/serve/bedrock_model_builder.py | 10 +++++++++ .../src/sagemaker/serve/model_builder.py | 7 +++++++ .../src/sagemaker/train/base_trainer.py | 6 ++++++ .../train/evaluate/base_evaluator.py | 6 ++++++ 6 files changed, 52 insertions(+) diff --git a/sagemaker-core/src/sagemaker/core/telemetry/constants.py b/sagemaker-core/src/sagemaker/core/telemetry/constants.py index 2cd1fd44c0..04231ef8a0 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/constants.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/constants.py @@ -30,6 +30,8 @@ class Feature(Enum): MLOPS = 16 FEATURE_STORE = 17 PROCESSING = 18 + MODEL_CUSTOMIZATION_NOVA = 19 + MODEL_CUSTOMIZATION_OSS = 20 def __str__(self): # pylint: disable=E0307 """Return the feature name.""" diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index 1041502064..738b47e309 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -62,6 +62,8 @@ str(Feature.MLOPS): 16, str(Feature.FEATURE_STORE): 17, str(Feature.PROCESSING): 18, + str(Feature.MODEL_CUSTOMIZATION_NOVA): 19, + str(Feature.MODEL_CUSTOMIZATION_OSS): 20, } STATUS_TO_CODE = { @@ -115,6 +117,25 @@ def wrapper(*args, **kwargs): # Construct the feature list to track feature combinations feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]] + # For MODEL_CUSTOMIZATION, append NOVA or OSS sub-feature + # based on the instance's _is_nova_model_for_telemetry() method + if feature == Feature.MODEL_CUSTOMIZATION and len(args) > 0: + instance = args[0] + try: + if hasattr(instance, "_is_nova_model_for_telemetry"): + if instance._is_nova_model_for_telemetry(): + feature_list.append( + FEATURE_TO_CODE[str(Feature.MODEL_CUSTOMIZATION_NOVA)] + ) + else: + feature_list.append( + FEATURE_TO_CODE[str(Feature.MODEL_CUSTOMIZATION_OSS)] + ) + except Exception: # pylint: disable=W0703 + logger.debug( + "Unable to determine NOVA/OSS model type for telemetry." + ) + if ( hasattr(sagemaker_session, "sagemaker_config") and sagemaker_session.sagemaker_config diff --git a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py index 5f9ccd14f8..38cbba09c2 100644 --- a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py @@ -95,6 +95,16 @@ def _get_sagemaker_client(self): self._sagemaker_client = self.boto_session.client("sagemaker") return self._sagemaker_client + def _is_nova_model_for_telemetry(self) -> bool: + """Check if the model is a Nova model for telemetry tracking.""" + try: + if not self.model_package: + return False + container = self.model_package.inference_specification.containers[0] + return _is_nova_model(container) + except Exception: + return False + @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="BedrockModelBuilder.deploy") def deploy( self, diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index ea4ec1eed7..7c7af2defc 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -1063,6 +1063,13 @@ def _is_nova_model(self) -> bool: hub_content_name = getattr(base_model, "hub_content_name", "") or "" return "nova" in recipe_name.lower() or "nova" in hub_content_name.lower() + def _is_nova_model_for_telemetry(self) -> bool: + """Check if the model is a Nova model for telemetry tracking.""" + try: + return self._is_nova_model() + except Exception: + return False + def _get_nova_hosting_config(self, instance_type=None): """Get Nova hosting config (image URI, env vars, instance type). diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index 873b42f81b..a422dc3240 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -69,6 +69,12 @@ def __init__( self.input_data_config = input_data_config self.environment = environment or {} + def _is_nova_model_for_telemetry(self) -> bool: + """Check if the model is a Nova model for telemetry tracking.""" + from sagemaker.train.common_utils.recipe_utils import _is_nova_model + model_name = getattr(self, "_model_name", None) + return _is_nova_model(model_name) if model_name else False + @abstractmethod def train(self, input_data_config: List[InputData], wait: bool = True, logs: bool = True): """Common training method that calls the specific implementation.""" diff --git a/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py b/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py index 3a5e7472d6..d53981ca43 100644 --- a/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py +++ b/sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py @@ -414,6 +414,12 @@ def _source_model_package_arn(self) -> Optional[str]: info = self._get_resolved_model_info() return info.source_model_package_arn if info else None + def _is_nova_model_for_telemetry(self) -> bool: + """Check if the model is a Nova model for telemetry tracking.""" + from ..common_utils.recipe_utils import _is_nova_model + base_model_name = self._base_model_name + return _is_nova_model(base_model_name) if base_model_name else False + @property def _is_jumpstart_model(self) -> bool: """Determine if model is a JumpStart model""" From ff31e873a161e4b1a996ee12f8ffd37b79c8bbc3 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 7 Apr 2026 10:53:26 -0700 Subject: [PATCH 2/2] test: add unit tests for NOVA/OSS telemetry sub-feature detection --- .../unit/telemetry/test_telemetry_logging.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py index 6ce1cb3269..04154e6ad7 100644 --- a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py +++ b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py @@ -542,3 +542,108 @@ def test_telemetry_emitter_without_resource_arn( args = mock_send_telemetry_request.call_args.args extra_str = str(args[5]) self.assertNotIn("x-resourceArn", extra_str) + + @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_appends_nova_sub_feature( + self, mock_resolve_config, mock_send_telemetry_request + ): + """Test that MODEL_CUSTOMIZATION_NOVA (19) is appended when instance reports Nova model.""" + mock_resolve_config.return_value = False + + class NovaModelMock: + def __init__(self): + self.sagemaker_session = MOCK_SESSION + + def _is_nova_model_for_telemetry(self): + return True + + @_telemetry_emitter(Feature.MODEL_CUSTOMIZATION, "NovaModelMock.train") + def train(self): + pass + + NovaModelMock().train() + + args = mock_send_telemetry_request.call_args.args + feature_list = args[1] + self.assertIn(15, feature_list) # MODEL_CUSTOMIZATION + self.assertIn(19, feature_list) # MODEL_CUSTOMIZATION_NOVA + self.assertNotIn(20, feature_list) # MODEL_CUSTOMIZATION_OSS should NOT be present + + @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_appends_oss_sub_feature( + self, mock_resolve_config, mock_send_telemetry_request + ): + """Test that MODEL_CUSTOMIZATION_OSS (20) is appended when instance reports non-Nova model.""" + mock_resolve_config.return_value = False + + class OssModelMock: + def __init__(self): + self.sagemaker_session = MOCK_SESSION + + def _is_nova_model_for_telemetry(self): + return False + + @_telemetry_emitter(Feature.MODEL_CUSTOMIZATION, "OssModelMock.train") + def train(self): + pass + + OssModelMock().train() + + args = mock_send_telemetry_request.call_args.args + feature_list = args[1] + self.assertIn(15, feature_list) # MODEL_CUSTOMIZATION + self.assertIn(20, feature_list) # MODEL_CUSTOMIZATION_OSS + self.assertNotIn(19, feature_list) # MODEL_CUSTOMIZATION_NOVA should NOT be present + + @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_no_sub_feature_without_detection_method( + self, mock_resolve_config, mock_send_telemetry_request + ): + """Test that no NOVA/OSS sub-feature is appended when instance lacks detection method.""" + mock_resolve_config.return_value = False + + class NoDetectionMock: + def __init__(self): + self.sagemaker_session = MOCK_SESSION + + @_telemetry_emitter(Feature.MODEL_CUSTOMIZATION, "NoDetectionMock.do_work") + def do_work(self): + pass + + NoDetectionMock().do_work() + + args = mock_send_telemetry_request.call_args.args + feature_list = args[1] + self.assertIn(15, feature_list) # MODEL_CUSTOMIZATION + self.assertNotIn(19, feature_list) # No NOVA + self.assertNotIn(20, feature_list) # No OSS + + @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_telemetry_emitter_handles_detection_method_exception( + self, mock_resolve_config, mock_send_telemetry_request + ): + """Test that telemetry still works when _is_nova_model_for_telemetry raises an exception.""" + mock_resolve_config.return_value = False + + class BrokenDetectionMock: + def __init__(self): + self.sagemaker_session = MOCK_SESSION + + def _is_nova_model_for_telemetry(self): + raise RuntimeError("detection failed") + + @_telemetry_emitter(Feature.MODEL_CUSTOMIZATION, "BrokenDetectionMock.train") + def train(self): + pass + + BrokenDetectionMock().train() + + args = mock_send_telemetry_request.call_args.args + feature_list = args[1] + self.assertIn(15, feature_list) # MODEL_CUSTOMIZATION still present + self.assertNotIn(19, feature_list) # No NOVA (detection failed gracefully) + self.assertNotIn(20, feature_list) # No OSS