diff --git a/sagemaker-core/src/sagemaker/core/experiments/experiment.py b/sagemaker-core/src/sagemaker/core/experiments/experiment.py index e0cf35599e..f555c44642 100644 --- a/sagemaker-core/src/sagemaker/core/experiments/experiment.py +++ b/sagemaker-core/src/sagemaker/core/experiments/experiment.py @@ -21,6 +21,8 @@ from sagemaker.core.experiments.trial import _Trial from sagemaker.core.experiments.trial_component import _TrialComponent from sagemaker.core.common_utils import format_tags +from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.core.telemetry.constants import Feature class Experiment(_base_types.Record): @@ -93,6 +95,7 @@ def load(cls, experiment_name, sagemaker_session=None): ) @classmethod + @_telemetry_emitter(feature=Feature.MLOPS, func_name="experiment.create") def create( cls, experiment_name, diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/models.py b/sagemaker-core/src/sagemaker/core/jumpstart/models.py index 05021ed9d2..1e87627821 100644 --- a/sagemaker-core/src/sagemaker/core/jumpstart/models.py +++ b/sagemaker-core/src/sagemaker/core/jumpstart/models.py @@ -398,6 +398,7 @@ class CapabilityEnum(StrEnum): TRAINING = "TRAINING" FINE_TUNING = "FINE_TUNING" VALIDATION = "VALIDATION" + INCREMENTAL_TRAINING = "INCREMENTAL_TRAINING" class DemoNotebookModel(BaseConfig): diff --git a/sagemaker-core/src/sagemaker/core/telemetry/constants.py b/sagemaker-core/src/sagemaker/core/telemetry/constants.py index 20f05706f2..5f6b6ff607 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/constants.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/constants.py @@ -22,14 +22,12 @@ class Feature(Enum): """Enumeration of feature names used in telemetry.""" - SDK_DEFAULTS = 1 - LOCAL_MODE = 2 - REMOTE_FUNCTION = 3 - MODEL_TRAINER = 4 - ESTIMATOR = 5 - HYPERPOD = 6 # Added to support telemetry in sagemaker-hyperpod-cli - HYPERPOD_CLI = 7 # Added to support telemetry in sagemaker-hyperpod-cli - MODEL_CUSTOMIZATION = 8 + SDK_DEFAULTS = 11 + LOCAL_MODE = 12 + REMOTE_FUNCTION = 13 + MODEL_TRAINER = 14 + MODEL_CUSTOMIZATION = 15 + MLOPS = 16 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 65d3a83900..728c8c2d5b 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -50,11 +50,12 @@ ) FEATURE_TO_CODE = { - str(Feature.SDK_DEFAULTS): 1, - str(Feature.LOCAL_MODE): 2, - str(Feature.REMOTE_FUNCTION): 3, - str(Feature.MODEL_TRAINER): 4, - str(Feature.MODEL_CUSTOMIZATION): 8, + str(Feature.SDK_DEFAULTS): 11, + str(Feature.LOCAL_MODE): 12, + str(Feature.REMOTE_FUNCTION): 13, + str(Feature.MODEL_TRAINER): 14, + str(Feature.MODEL_CUSTOMIZATION): 15, + str(Feature.MLOPS): 16, } STATUS_TO_CODE = { diff --git a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py index 8c0244a5c3..41af998d73 100644 --- a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py +++ b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py @@ -155,7 +155,7 @@ def test_telemetry_emitter_decorator_success( ) mock_send_telemetry_request.assert_called_once_with( - 1, [1, 2], MOCK_SESSION, None, None, expected_extra_str + 1, [11, 12], MOCK_SESSION, None, None, expected_extra_str ) @pytest.mark.skipif( diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py index 144726f690..3a389344ca 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py @@ -69,6 +69,8 @@ ) from sagemaker.core.workflow.utilities import list_to_request from sagemaker.mlops.workflow._steps_compiler import StepsCompiler +from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.core.telemetry.constants import Feature logger = logging.getLogger(__name__) @@ -139,6 +141,7 @@ def latest_pipeline_version_id(self): else: return summaries[0].get("PipelineVersionId") + @_telemetry_emitter(feature=Feature.MLOPS, func_name="pipeline.create") def create( self, role_arn: str = None, @@ -348,6 +351,7 @@ def delete(self) -> Dict[str, Any]: ) return self.sagemaker_session.sagemaker_client.delete_pipeline(PipelineName=self.name) + @_telemetry_emitter(feature=Feature.MLOPS, func_name="pipeline.start") def start( self, parameters: Dict[str, Union[str, bool, int, float]] = None, diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py index 1550a95c36..7f227f53b6 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline.py @@ -26,6 +26,7 @@ def mock_session(): session.boto_session.client.return_value = Mock() session.sagemaker_client = Mock() session.local_mode = False + session.sagemaker_config = {} return session diff --git a/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py b/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py index e45fccab42..59502bed3f 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py +++ b/sagemaker-mlops/tests/unit/workflow/test_pipeline_class.py @@ -34,6 +34,7 @@ def mock_session(): session.boto_session = Mock() session.boto_session.client = Mock(return_value=Mock()) session.local_mode = False + session.sagemaker_config = {} session._append_sagemaker_config_tags = Mock(return_value=[]) return session