fix: Update telemetry constants to track v2 feature seperately (aws#5437)

* fix: Update telmetry constants to track v2 feature seperately

* fix: Update telmetry constants to track v2 feature seperately

---------

Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>
Co-authored-by: aviruthen <91846056+aviruthen@users.noreply.github.com>

Author: 3 people

src/sagemaker/estimator.py

@@ -1350,7 +1350,7 @@ def latest_job_profiler_artifacts_path(self):
             )
         return None
 
-    @_telemetry_emitter(feature=Feature.ESTIMATOR, func_name="estimator.fit")
+    @_telemetry_emitter(feature=Feature.ESTIMATOR_V2, func_name="estimator.fit")
     @runnable_by_pipeline
     def fit(
         self,

src/sagemaker/experiments/experiment.py

@@ -95,7 +95,7 @@ def load(cls, experiment_name, sagemaker_session=None):
         )
 
     @classmethod
-    @_telemetry_emitter(feature=Feature.MLOPS, func_name="experiment.create")
+    @_telemetry_emitter(feature=Feature.MLOPS_V2, func_name="experiment.create")
     def create(
         cls,
         experiment_name,

src/sagemaker/jumpstart/estimator.py

@@ -62,7 +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")
+    @_telemetry_emitter(feature=Feature.JUMPSTART_V2, func_name="jumpstart_estimator.create")
     def __init__(
         self,
         model_id: Optional[str] = None,
@@ -649,7 +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")
+    @_telemetry_emitter(feature=Feature.JUMPSTART_V2, func_name="jumpstart_estimator.fit")
     def fit(
         self,
         inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None,
@@ -837,7 +837,7 @@ def attach(
             additional_kwargs=additional_kwargs,
         )
 
-    @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_estimator.deploy")
+    @_telemetry_emitter(feature=Feature.JUMPSTART_V2, func_name="jumpstart_estimator.deploy")
     def deploy(
         self,
         initial_instance_count: Optional[int] = None,

src/sagemaker/jumpstart/model.py

@@ -86,7 +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")
+    @_telemetry_emitter(feature=Feature.JUMPSTART_V2, func_name="jumpstart_model.create")
     def __init__(
         self,
         model_id: Optional[str] = None,
@@ -643,7 +643,7 @@ def _create_sagemaker_model(
                 **kwargs,
             )
 
-    @_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_model.deploy")
+    @_telemetry_emitter(feature=Feature.JUMPSTART_V2, func_name="jumpstart_model.deploy")
     def deploy(
         self,
         initial_instance_count: Optional[int] = None,

src/sagemaker/local/local_session.py

@@ -85,7 +85,7 @@ def __init__(self, sagemaker_session=None):
         """
         self.sagemaker_session = sagemaker_session or LocalSession()
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_processing_job")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_processing_job")
     def create_processing_job(
         self,
         ProcessingJobName,
@@ -168,7 +168,7 @@ def describe_processing_job(self, ProcessingJobName):
             raise ClientError(error_response, "describe_processing_job")
         return LocalSagemakerClient._processing_jobs[ProcessingJobName].describe()
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_training_job")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_training_job")
     def create_training_job(
         self,
         TrainingJobName,
@@ -239,7 +239,7 @@ def describe_training_job(self, TrainingJobName):
             raise ClientError(error_response, "describe_training_job")
         return LocalSagemakerClient._training_jobs[TrainingJobName].describe()
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_transform_job")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_transform_job")
     def create_transform_job(
         self,
         TransformJobName,
@@ -285,7 +285,7 @@ def describe_transform_job(self, TransformJobName):
             raise ClientError(error_response, "describe_transform_job")
         return LocalSagemakerClient._transform_jobs[TransformJobName].describe()
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_model")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_model")
     def create_model(
         self, ModelName, PrimaryContainer, *args, **kwargs
     ):  # pylint: disable=unused-argument
@@ -335,7 +335,7 @@ def describe_endpoint_config(self, EndpointConfigName):
             raise ClientError(error_response, "describe_endpoint_config")
         return LocalSagemakerClient._endpoint_configs[EndpointConfigName].describe()
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_endpoint_config")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_endpoint_config")
     def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
         """Create the endpoint configuration.
 
@@ -367,7 +367,7 @@ def describe_endpoint(self, EndpointName):
             raise ClientError(error_response, "describe_endpoint")
         return LocalSagemakerClient._endpoints[EndpointName].describe()
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_endpoint")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_endpoint")
     def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
         """Create the endpoint.
 
@@ -436,7 +436,7 @@ def delete_model(self, ModelName):
         if ModelName in LocalSagemakerClient._models:
             del LocalSagemakerClient._models[ModelName]
 
-    @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_pipeline")
+    @_telemetry_emitter(Feature.LOCAL_MODE_V2, "local_session.create_pipeline")
     def create_pipeline(
         self, pipeline, pipeline_description, **kwargs  # pylint: disable=unused-argument
     ):

src/sagemaker/modules/train/model_trainer.py

@@ -807,7 +807,7 @@ def _create_training_job_args(
             args["session_chaining_config"] = self._session_chaining_config
             return args
 
-    @_telemetry_emitter(feature=Feature.MODEL_TRAINER, func_name="model_trainer.train")
+    @_telemetry_emitter(feature=Feature.MODEL_TRAINER_V2, func_name="model_trainer.train")
     @validate_call
     def train(
         self,

src/sagemaker/remote_function/client.py

@@ -59,7 +59,7 @@
 logger = logging_config.get_logger()
 
 
-@_telemetry_emitter(feature=Feature.REMOTE_FUNCTION, func_name="remote_function.remote")
+@_telemetry_emitter(feature=Feature.REMOTE_FUNCTION_V2, func_name="remote_function.remote")
 def remote(
     _func=None,
     *,

src/sagemaker/telemetry/constants.py

@@ -22,15 +22,14 @@
 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
+    SDK_DEFAULTS_V2 = 1
+    LOCAL_MODE_V2 = 2
+    REMOTE_FUNCTION_V2 = 3
+    MODEL_TRAINER_V2 = 4
+    ESTIMATOR_V2 = 5
     # Note: HyperPod CLI uses codes 6 and 7
-    JUMPSTART = 8  # Added to support JumpStart telemetry
-    MLOPS = 9  # Added to support MLOps telemetry
+    JUMPSTART_V2 = 8  # Added to support JumpStart telemetry
+    MLOPS_V2 = 9  # Added to support MLOps telemetry
 
     def __str__(self):  # pylint: disable=E0307
         """Return the feature name."""

src/sagemaker/telemetry/telemetry_logging.py

@@ -50,15 +50,14 @@
 )
 
 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.ESTIMATOR): 5,
-    str(Feature.HYPERPOD): 6,  # Added to support telemetry in sagemaker-hyperpod-cli
+    str(Feature.SDK_DEFAULTS_V2): 1,
+    str(Feature.LOCAL_MODE_V2): 2,
+    str(Feature.REMOTE_FUNCTION_V2): 3,
+    str(Feature.MODEL_TRAINER_V2): 4,
+    str(Feature.ESTIMATOR_V2): 5,
     # Note: HyperPod CLI uses codes 6 and 7
-    str(Feature.JUMPSTART): 8,
-    str(Feature.MLOPS): 9,
+    str(Feature.JUMPSTART_V2): 8,
+    str(Feature.MLOPS_V2): 9,
 }
 
 STATUS_TO_CODE = {
@@ -84,7 +83,7 @@ def wrapper(*args, **kwargs):
             if len(args) > 0 and hasattr(args[0], "sagemaker_session"):
                 # Get the sagemaker_session from the instance method args
                 sagemaker_session = args[0].sagemaker_session
-            elif feature == Feature.REMOTE_FUNCTION:
+            elif feature == Feature.REMOTE_FUNCTION_V2:
                 # Get the sagemaker_session from the function keyword arguments for remote function
                 sagemaker_session = kwargs.get(
                     "sagemaker_session", _get_default_sagemaker_session()
@@ -112,16 +111,16 @@ def wrapper(*args, **kwargs):
                 if (
                     hasattr(sagemaker_session, "sagemaker_config")
                     and sagemaker_session.sagemaker_config
-                    and feature != Feature.SDK_DEFAULTS
+                    and feature != Feature.SDK_DEFAULTS_V2
                 ):
-                    feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS)])
+                    feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS_V2)])
 
                 if (
                     hasattr(sagemaker_session, "local_mode")
                     and sagemaker_session.local_mode
-                    and feature != Feature.LOCAL_MODE
+                    and feature != Feature.LOCAL_MODE_V2
                 ):
-                    feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE)])
+                    feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE_V2)])
 
                 # Construct the extra info to track platform and environment usage metadata
                 extra = (

src/sagemaker/workflow/pipeline.py

@@ -136,7 +136,7 @@ def latest_pipeline_version_id(self):
         else:
             return summaries[0].get("PipelineVersionId")
 
-    @_telemetry_emitter(feature=Feature.MLOPS, func_name="pipeline.create")
+    @_telemetry_emitter(feature=Feature.MLOPS_V2, func_name="pipeline.create")
     def create(
         self,
         role_arn: str = None,
@@ -345,7 +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")
+    @_telemetry_emitter(feature=Feature.MLOPS_V2, func_name="pipeline.start")
     def start(
         self,
         parameters: Dict[str, Union[str, bool, int, float]] = None,