Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rsareddy0329 wants to merge 2 commits into
base: master-v2
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2c6724b
fix: Update telmetry constants to track v2 feature seperately
Dec 23, 2025
d6fdf75
fix: Update telmetry constants to track v2 feature seperately
Dec 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/sagemaker/estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion src/sagemaker/experiments/experiment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 3 additions & 3 deletions src/sagemaker/jumpstart/estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/sagemaker/jumpstart/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
14 changes: 7 additions & 7 deletions src/sagemaker/local/local_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
):
Expand Down
2 changes: 1 addition & 1 deletion src/sagemaker/modules/train/model_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion src/sagemaker/remote_function/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
*,
Expand Down
15 changes: 7 additions & 8 deletions src/sagemaker/telemetry/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
25 changes: 12 additions & 13 deletions src/sagemaker/telemetry/telemetry_logging.py
View file
Open in desktop
Copy link
Collaborator

@aviruthen aviruthen Dec 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would entries be inconsistent for V2 with this setup? For new users this makes sense, but for users that are already using V2 and don't pip install again, they will emit telemetry for these features at their original codes (1-9) and only new users will emit telemetry in the range 11-17. Should we consider setting the V3 telemetry codes to be 11-17 and keeping V2 in the range of 1-9? I don't feel strongly either way since there are trade-offs with either approach

Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand All @@ -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()
Expand Down Expand Up @@ -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 = (
Expand Down
4 changes: 2 additions & 2 deletions src/sagemaker/workflow/pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/sagemaker/telemetry/test_telemetry_logging.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down
Loading