feat: split MODEL_CUSTOMIZATION telemetry into NOVA and OSS sub-features

Author: lucasjia-aws

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."""

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

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,

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

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."""

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"""