Feature/js mlops telemetry (#5268)

* removed log statement

* added telemetry for js and mlops

* added for js estimator

* fixed unit tests

---------

Co-authored-by: Mohamed Zeidan <[email protected]>

Author: mohamedzeidan2021

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,

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,

src/sagemaker/jumpstart/model.py

@@ -76,13 +76,17 @@
 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.
 
     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,

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

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 = {

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,

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)