From 2c6724b7cef8fe0f5719e7c60b1e4b73fda348ec Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Mon, 22 Dec 2025 22:57:59 -0800 Subject: [PATCH 1/2] fix: Update telmetry constants to track v2 feature seperately --- src/sagemaker/estimator.py | 2 +- src/sagemaker/experiments/experiment.py | 2 +- src/sagemaker/jumpstart/estimator.py | 6 ++--- src/sagemaker/jumpstart/model.py | 4 +-- src/sagemaker/local/local_session.py | 14 +++++------ src/sagemaker/modules/train/model_trainer.py | 2 +- src/sagemaker/remote_function/client.py | 2 +- src/sagemaker/telemetry/constants.py | 15 ++++++----- src/sagemaker/telemetry/telemetry_logging.py | 25 +++++++++---------- src/sagemaker/workflow/pipeline.py | 4 +-- .../telemetry/test_telemetry_logging.py | 2 +- 11 files changed, 38 insertions(+), 40 deletions(-) diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 413f948af6..8fb3fca54e 100644 --- a/src/sagemaker/estimator.py +++ b/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, diff --git a/src/sagemaker/experiments/experiment.py b/src/sagemaker/experiments/experiment.py index 5ee31a7934..7c68ff2896 100644 --- a/src/sagemaker/experiments/experiment.py +++ b/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, diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py index 989f520f42..b8b75eaa83 100644 --- a/src/sagemaker/jumpstart/estimator.py +++ b/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, diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index 4e5d059c2c..2853df2edc 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/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, diff --git a/src/sagemaker/local/local_session.py b/src/sagemaker/local/local_session.py index 89a2df2135..5ae6bfba82 100644 --- a/src/sagemaker/local/local_session.py +++ b/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 ): diff --git a/src/sagemaker/modules/train/model_trainer.py b/src/sagemaker/modules/train/model_trainer.py index 179d03519d..b598875c00 100644 --- a/src/sagemaker/modules/train/model_trainer.py +++ b/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, diff --git a/src/sagemaker/remote_function/client.py b/src/sagemaker/remote_function/client.py index f70eb48c58..81ad47fe8b 100644 --- a/src/sagemaker/remote_function/client.py +++ b/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, *, diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index e860e5ced3..1d98fec587 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/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 = 11 + LOCAL_MODE_V2 = 12 + REMOTE_FUNCTION_V2 = 13 + MODEL_TRAINER_V2 = 14 + ESTIMATOR_V2 = 15 # Note: HyperPod CLI uses codes 6 and 7 - JUMPSTART = 8 # Added to support JumpStart telemetry - MLOPS = 9 # Added to support MLOps telemetry + JUMPSTART_V2 = 16 # Added to support JumpStart telemetry + MLOPS_V2 = 17 # 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 f8261a8c2d..27370241fc 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/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): 11, + str(Feature.LOCAL_MODE_V2): 12, + str(Feature.REMOTE_FUNCTION_V2): 13, + str(Feature.MODEL_TRAINER_V2): 14, + str(Feature.ESTIMATOR_V2): 15, # Note: HyperPod CLI uses codes 6 and 7 - str(Feature.JUMPSTART): 8, - str(Feature.MLOPS): 9, + str(Feature.JUMPSTART_V2): 16, + str(Feature.MLOPS_V2): 17, } 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 = ( diff --git a/src/sagemaker/workflow/pipeline.py b/src/sagemaker/workflow/pipeline.py index 0c60401d7e..25d236f25e 100644 --- a/src/sagemaker/workflow/pipeline.py +++ b/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, diff --git a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py index bd8db82a16..031ddfc291 100644 --- a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py +++ b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py @@ -34,7 +34,7 @@ MOCK_SESSION = Mock() MOCK_EXCEPTION = LocalModelOutOfMemoryException("mock raise ex") -MOCK_FEATURE = Feature.SDK_DEFAULTS +MOCK_FEATURE = Feature.SDK_DEFAULTS_V2 MOCK_FUNC_NAME = "Mock.local_session.create_model" MOCK_ENDPOINT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:endpoint/test" From d6fdf75c1f1121bd8d630c762ba73f5da3c31f95 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 23 Dec 2025 11:39:12 -0800 Subject: [PATCH 2/2] fix: Update telmetry constants to track v2 feature seperately --- src/sagemaker/telemetry/constants.py | 14 +++++++------- src/sagemaker/telemetry/telemetry_logging.py | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index 1d98fec587..78b49c888c 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/src/sagemaker/telemetry/constants.py @@ -22,14 +22,14 @@ class Feature(Enum): """Enumeration of feature names used in telemetry.""" - SDK_DEFAULTS_V2 = 11 - LOCAL_MODE_V2 = 12 - REMOTE_FUNCTION_V2 = 13 - MODEL_TRAINER_V2 = 14 - ESTIMATOR_V2 = 15 + 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_V2 = 16 # Added to support JumpStart telemetry - MLOPS_V2 = 17 # 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.""" diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index 27370241fc..7a33840752 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -50,14 +50,14 @@ ) FEATURE_TO_CODE = { - str(Feature.SDK_DEFAULTS_V2): 11, - str(Feature.LOCAL_MODE_V2): 12, - str(Feature.REMOTE_FUNCTION_V2): 13, - str(Feature.MODEL_TRAINER_V2): 14, - str(Feature.ESTIMATOR_V2): 15, + 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_V2): 16, - str(Feature.MLOPS_V2): 17, + str(Feature.JUMPSTART_V2): 8, + str(Feature.MLOPS_V2): 9, } STATUS_TO_CODE = {