Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions sagemaker-core/src/sagemaker/core/telemetry/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
21 changes: 21 additions & 0 deletions sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -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():
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isnt this creating a circular dependency ?
I think we can avoid it if possible.

Can you try calling _telemetry_emitter with eature.MODEL_CUSTOMIZATION_NOVA if that particular instance data is nova .

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No circular dependency here — telemetry_logging.py only imports from sagemaker.core internals. The NOVA/OSS detection uses duck typing (hasattr check on the instance at runtime), no new imports are added to the telemetry module.

Regarding calling _telemetry_emitter with MODEL_CUSTOMIZATION_NOVA directly: the decorator's feature parameter is evaluated at class definition time, but NOVA vs OSS is only known at runtime (depends on the model string the user passes). So we can't statically set it at the decorator call site. The runtime detection in the wrapper is the only way to do this without duplicating every decorated method into NOVA/OSS variants.

Added 4 unit tests covering NOVA, OSS, no-detection, and exception-handling cases.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The integ-tests (sagemaker-train) failure is pre-existing — recent CI runs (#553, #552, #551, #550, #549, etc.) all show the same failure regardless of code changes. The root cause appears to be missing resource cleanup in ai_registry integ tests (e.g., test_air_hub.py creates hub content without using the cleanup_list fixture). I think we can open a separate PR to add proper cleanup.

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
Expand Down
105 changes: 105 additions & 0 deletions sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
10 changes: 10 additions & 0 deletions sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
7 changes: 7 additions & 0 deletions sagemaker-serve/src/sagemaker/serve/model_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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).

Expand Down
6 changes: 6 additions & 0 deletions sagemaker-train/src/sagemaker/train/base_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down
Loading