add telemetry for feature store

Author: Romik Amipara

src/sagemaker/feature_store/feature_group.py

@@ -71,6 +71,8 @@
     TargetStoreEnum,
 )
 from sagemaker.utils import resolve_value_from_config, format_tags, Tags
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 
 logger = logging.getLogger(__name__)
 
@@ -672,6 +674,7 @@ class FeatureGroup:
         FeatureTypeEnum.STRING.value: "STRING",
     }
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.create")
     def create(
         self,
         s3_uri: Union[str, bool],
@@ -788,10 +791,12 @@ def create(
 
         return self.sagemaker_session.create_feature_group(**create_feature_store_args)
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.delete")
     def delete(self):
         """Delete a FeatureGroup."""
         self.sagemaker_session.delete_feature_group(feature_group_name=self.name)
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.describe")
     def describe(self, next_token: str = None) -> Dict[str, Any]:
         """Describe a FeatureGroup.
 
@@ -805,6 +810,7 @@ def describe(self, next_token: str = None) -> Dict[str, Any]:
             feature_group_name=self.name, next_token=next_token
         )
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.update")
     def update(
         self,
         feature_additions: Sequence[FeatureDefinition] = None,
@@ -842,7 +848,7 @@ def update(
             online_store_config=online_store_config_parameter,
             throughput_config=throughput_config_parameter,
         )
-
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.update_feature_metadata")
     def update_feature_metadata(
         self,
         feature_name: str,
@@ -871,6 +877,7 @@ def update_feature_metadata(
             parameter_removals=(parameter_removals or []),
         )
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.describe_feature_metadata")
     def describe_feature_metadata(self, feature_name: str) -> Dict[str, Any]:
         """Describe feature metadata by feature name.
 
@@ -1038,6 +1045,7 @@ def load_feature_definitions(
         self.feature_definitions = feature_definitions
         return self.feature_definitions
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.get_record")
     def get_record(
         self,
         record_identifier_value_as_string: str,
@@ -1057,6 +1065,7 @@ def get_record(
             feature_names=feature_names,
         ).get("Record")
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.put_record")
     def put_record(
         self,
         record: Sequence[FeatureValue],
@@ -1080,6 +1089,7 @@ def put_record(
             ttl_duration=ttl_duration.to_dict() if ttl_duration is not None else None,
         )
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_group.delete_record")
     def delete_record(
         self,
         record_identifier_value_as_string: str,

src/sagemaker/feature_store/feature_store.py

@@ -33,6 +33,8 @@
     SortOrderEnum,
     Identifier,
 )
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 
 
 @attr.s
@@ -47,6 +49,7 @@ class FeatureStore:
 
     sagemaker_session: Session = attr.ib(default=Session)
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_store.create_dataset")
     def create_dataset(
         self,
         base: Union[FeatureGroup, pd.DataFrame],
@@ -92,6 +95,7 @@ def create_dataset(
             kms_key_id,
         )
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_store.list_feature_groups")
     def list_feature_groups(
         self,
         name_contains: str = None,
@@ -137,6 +141,7 @@ def list_feature_groups(
             next_token=next_token,
         )
 
+    @_telemetry_emitter(feature=Feature.FEATURE_STORE_V2, func_name="feature_store.batch_get_record")
     def batch_get_record(
         self,
         identifiers: Sequence[Identifier],

src/sagemaker/telemetry/constants.py

@@ -30,6 +30,7 @@ class Feature(Enum):
     # Note: HyperPod CLI uses codes 6 and 7
     JUMPSTART_V2 = 8  # Added to support JumpStart telemetry
     MLOPS_V2 = 9  # Added to support MLOps telemetry
+    FEATURE_STORE_V2 = 10
 
     def __str__(self):  # pylint: disable=E0307
         """Return the feature name."""

src/sagemaker/telemetry/telemetry_logging.py

@@ -58,6 +58,7 @@
     # Note: HyperPod CLI uses codes 6 and 7
     str(Feature.JUMPSTART_V2): 8,
     str(Feature.MLOPS_V2): 9,
+    str(Feature.FEATURE_STORE_V2): 10,
 }
 
 STATUS_TO_CODE = {