Feature/js mlops telemetry

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,