diff --git a/src/sagemaker/experiments/experiment.py b/src/sagemaker/experiments/experiment.py index 6f33fafb0f..5ee31a7934 100644 --- a/src/sagemaker/experiments/experiment.py +++ b/src/sagemaker/experiments/experiment.py @@ -21,6 +21,8 @@ from sagemaker.experiments.trial import _Trial from sagemaker.experiments.trial_component import _TrialComponent from sagemaker.utils import format_tags +from sagemaker.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.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/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py index e61e1c49a5..989f520f42 100644 --- a/src/sagemaker/jumpstart/estimator.py +++ b/src/sagemaker/jumpstart/estimator.py @@ -52,6 +52,8 @@ from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig from sagemaker.workflow.entities import PipelineVariable +from sagemaker.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.telemetry.constants import Feature class JumpStartEstimator(Estimator): @@ -60,6 +62,7 @@ class JumpStartEstimator(Estimator): This class sets defaults based on the model ID and version. """ + @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_estimator.create") def __init__( self, model_id: Optional[str] = None, @@ -646,6 +649,7 @@ def _validate_model_id_and_get_type_hook(): super(JumpStartEstimator, self).__init__(**estimator_init_kwargs.to_kwargs_dict()) + @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_estimator.fit") def fit( self, inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None, @@ -833,6 +837,7 @@ def attach( additional_kwargs=additional_kwargs, ) + @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_estimator.deploy") def deploy( self, initial_instance_count: Optional[int] = None, diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index 7dec3d78f9..4e5d059c2c 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/src/sagemaker/jumpstart/model.py @@ -76,6 +76,9 @@ from sagemaker.drift_check_baselines import DriftCheckBaselines from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements +from sagemaker.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.telemetry.constants import Feature + class JumpStartModel(Model): """JumpStartModel class. @@ -83,6 +86,7 @@ class JumpStartModel(Model): This class sets defaults based on the model ID and version. """ + @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_model.create") def __init__( self, model_id: Optional[str] = None, @@ -639,6 +643,7 @@ def _create_sagemaker_model( **kwargs, ) + @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_model.deploy") def deploy( self, initial_instance_count: Optional[int] = None, diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index 6766d45b4e..e860e5ced3 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/src/sagemaker/telemetry/constants.py @@ -28,6 +28,9 @@ class Feature(Enum): MODEL_TRAINER = 4 ESTIMATOR = 5 HYPERPOD = 6 # Added to support telemetry in sagemaker-hyperpod-cli + # Note: HyperPod CLI uses codes 6 and 7 + JUMPSTART = 8 # Added to support JumpStart telemetry + MLOPS = 9 # Added to support MLOps telemetry def __str__(self): # pylint: disable=E0307 """Return the feature name.""" diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index 990e12124f..f8261a8c2d 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -56,6 +56,9 @@ str(Feature.MODEL_TRAINER): 4, str(Feature.ESTIMATOR): 5, str(Feature.HYPERPOD): 6, # Added to support telemetry in sagemaker-hyperpod-cli + # Note: HyperPod CLI uses codes 6 and 7 + str(Feature.JUMPSTART): 8, + str(Feature.MLOPS): 9, } STATUS_TO_CODE = { diff --git a/src/sagemaker/workflow/pipeline.py b/src/sagemaker/workflow/pipeline.py index f1a62fa637..f111f5e40b 100644 --- a/src/sagemaker/workflow/pipeline.py +++ b/src/sagemaker/workflow/pipeline.py @@ -64,6 +64,8 @@ ) from sagemaker.workflow.utilities import list_to_request from sagemaker.workflow._steps_compiler import StepsCompiler +from sagemaker.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.telemetry.constants import Feature logger = logging.getLogger(__name__) @@ -134,6 +136,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, @@ -342,6 +345,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/tests/unit/sagemaker/jumpstart/estimator/test_sagemaker_config.py b/tests/unit/sagemaker/jumpstart/estimator/test_sagemaker_config.py index 39eca166ee..0e4b9d8201 100644 --- a/tests/unit/sagemaker/jumpstart/estimator/test_sagemaker_config.py +++ b/tests/unit/sagemaker/jumpstart/estimator/test_sagemaker_config.py @@ -18,6 +18,7 @@ from sagemaker.config.config_schema import ( MODEL_ENABLE_NETWORK_ISOLATION_PATH, MODEL_EXECUTION_ROLE_ARN_PATH, + TELEMETRY_OPT_OUT_PATH, TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH, TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH, TRAINING_JOB_ROLE_ARN_PATH, @@ -75,6 +76,9 @@ def config_value_impl(sagemaker_session: Session, config_path: str, sagemaker_co if config_path == MODEL_ENABLE_NETWORK_ISOLATION_PATH: return config_inference_enable_network_isolation + if config_path == TELEMETRY_OPT_OUT_PATH: + return False # Default to telemetry enabled for tests + raise AssertionError(f"Bad config path: {config_path}") @@ -130,7 +134,7 @@ def test_without_arg_overwrites_without_kwarg_collisions_with_config( estimator.deploy() - self.assertEqual(mock_get_sagemaker_config_value.call_count, 3) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 4) self.assertEqual(mock_estimator_deploy.call_args[1].get("role"), config_inference_role) @@ -200,7 +204,7 @@ def test_without_arg_overwrites_with_kwarg_collisions_with_config( estimator.deploy() - self.assertEqual(mock_get_sagemaker_config_value.call_count, 6) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 7) self.assertEqual(mock_estimator_deploy.call_args[1].get("role"), config_inference_role) @@ -280,7 +284,7 @@ def test_with_arg_overwrites_with_kwarg_collisions_with_config( enable_network_isolation=override_inference_enable_network_isolation, ) - self.assertEqual(mock_get_sagemaker_config_value.call_count, 3) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 4) self.assertEqual( mock_estimator_deploy.call_args[1].get("role"), mock_inference_override_role @@ -355,7 +359,7 @@ def test_with_arg_overwrites_without_kwarg_collisions_with_config( enable_network_isolation=override_inference_enable_network_isolation, ) - self.assertEqual(mock_get_sagemaker_config_value.call_count, 3) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 4) self.assertEqual( mock_estimator_deploy.call_args[1].get("role"), mock_inference_override_role @@ -421,7 +425,7 @@ def test_without_arg_overwrites_without_kwarg_collisions_without_config( mock_retrieve_model_init_kwargs.return_value = {} - self.assertEqual(mock_get_sagemaker_config_value.call_count, 3) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 4) self.assertEqual(mock_estimator_deploy.call_args[1].get("role"), execution_role) @@ -492,7 +496,7 @@ def test_without_arg_overwrites_with_kwarg_collisions_without_config( estimator.deploy() - self.assertEqual(mock_get_sagemaker_config_value.call_count, 6) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 7) self.assertEqual(mock_estimator_deploy.call_args[1].get("role"), execution_role) @@ -568,7 +572,7 @@ def test_with_arg_overwrites_with_kwarg_collisions_without_config( enable_network_isolation=override_inference_enable_network_isolation, ) - self.assertEqual(mock_get_sagemaker_config_value.call_count, 3) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 4) self.assertEqual(mock_estimator_deploy.call_args[1].get("role"), override_inference_role) @@ -634,7 +638,7 @@ def test_with_arg_overwrites_without_kwarg_collisions_without_config( enable_network_isolation=override_enable_network_isolation, ) - self.assertEqual(mock_get_sagemaker_config_value.call_count, 3) + self.assertEqual(mock_get_sagemaker_config_value.call_count, 4) self.assertEqual(mock_estimator_deploy.call_args[1].get("role"), override_inference_role)