Skip to content

Commit 3089712

Browse files
author
Mohamed Zeidan
committed
added telemetry for js and mlops
1 parent fe97782 commit 3089712

File tree

5 files changed

+18
-0
lines changed

5 files changed

+18
-0
lines changed

src/sagemaker/experiments/experiment.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
from sagemaker.experiments.trial import _Trial
2222
from sagemaker.experiments.trial_component import _TrialComponent
2323
from sagemaker.utils import format_tags
24+
from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
25+
from sagemaker.telemetry.constants import Feature
2426

2527

2628
class Experiment(_base_types.Record):
@@ -93,6 +95,7 @@ def load(cls, experiment_name, sagemaker_session=None):
9395
)
9496

9597
@classmethod
98+
@_telemetry_emitter(feature=Feature.MLOPS, func_name="experiment.create")
9699
def create(
97100
cls,
98101
experiment_name,

src/sagemaker/jumpstart/model.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,13 +76,17 @@
7676
from sagemaker.drift_check_baselines import DriftCheckBaselines
7777
from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
7878

79+
from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
80+
from sagemaker.telemetry.constants import Feature
81+
7982

8083
class JumpStartModel(Model):
8184
"""JumpStartModel class.
8285
8386
This class sets defaults based on the model ID and version.
8487
"""
8588

89+
@_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_model.create")
8690
def __init__(
8791
self,
8892
model_id: Optional[str] = None,
@@ -639,6 +643,7 @@ def _create_sagemaker_model(
639643
**kwargs,
640644
)
641645

646+
@_telemetry_emitter(feature=Feature.JUMPSTART, func_name="jumpstart_model.deploy")
642647
def deploy(
643648
self,
644649
initial_instance_count: Optional[int] = None,

src/sagemaker/telemetry/constants.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@ class Feature(Enum):
2828
MODEL_TRAINER = 4
2929
ESTIMATOR = 5
3030
HYPERPOD = 6 # Added to support telemetry in sagemaker-hyperpod-cli
31+
# Note: HyperPod CLI uses codes 6 and 7
32+
JUMPSTART = 8 # Added to support JumpStart telemetry
33+
MLOPS = 9 # Added to support MLOps telemetry
3134

3235
def __str__(self): # pylint: disable=E0307
3336
"""Return the feature name."""

src/sagemaker/telemetry/telemetry_logging.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@
5656
str(Feature.MODEL_TRAINER): 4,
5757
str(Feature.ESTIMATOR): 5,
5858
str(Feature.HYPERPOD): 6, # Added to support telemetry in sagemaker-hyperpod-cli
59+
# Note: HyperPod CLI uses codes 6 and 7
60+
str(Feature.JUMPSTART): 8,
61+
str(Feature.MLOPS): 9,
5962
}
6063

6164
STATUS_TO_CODE = {

src/sagemaker/workflow/pipeline.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@
6464
)
6565
from sagemaker.workflow.utilities import list_to_request
6666
from sagemaker.workflow._steps_compiler import StepsCompiler
67+
from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
68+
from sagemaker.telemetry.constants import Feature
6769

6870
logger = logging.getLogger(__name__)
6971

@@ -134,6 +136,7 @@ def latest_pipeline_version_id(self):
134136
else:
135137
return summaries[0].get("PipelineVersionId")
136138

139+
@_telemetry_emitter(feature=Feature.MLOPS, func_name="pipeline.create")
137140
def create(
138141
self,
139142
role_arn: str = None,
@@ -342,6 +345,7 @@ def delete(self) -> Dict[str, Any]:
342345
)
343346
return self.sagemaker_session.sagemaker_client.delete_pipeline(PipelineName=self.name)
344347

348+
@_telemetry_emitter(feature=Feature.MLOPS, func_name="pipeline.start")
345349
def start(
346350
self,
347351
parameters: Dict[str, Union[str, bool, int, float]] = None,

0 commit comments

Comments
 (0)