Add feature store telemetry

Author: nargokul

sagemaker-core/src/sagemaker/core/telemetry/constants.py

@@ -28,6 +28,7 @@ class Feature(Enum):
     MODEL_TRAINER = 14
     MODEL_CUSTOMIZATION = 15
     MLOPS = 16
+    FEATURE_STORE = 17
 
     def __str__(self):  # pylint: disable=E0307
         """Return the feature name."""

sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py

@@ -56,6 +56,7 @@
     str(Feature.MODEL_TRAINER): 14,
     str(Feature.MODEL_CUSTOMIZATION): 15,
     str(Feature.MLOPS): 16,
+    str(Feature.FEATURE_STORE): 17,
 }
 
 STATUS_TO_CODE = {
@@ -81,6 +82,9 @@ 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 len(args) > 0 and hasattr(args[0], "_sagemaker_session"):
+                # Get the sagemaker_session from the instance method args (private attribute)
+                sagemaker_session = args[0]._sagemaker_session
             elif feature == Feature.REMOTE_FUNCTION:
                 # Get the sagemaker_session from the function keyword arguments for remote function
                 sagemaker_session = kwargs.get(

sagemaker-mlops/src/sagemaker/mlops/feature_store/athena_query.py

@@ -14,6 +14,7 @@
 )
 
 from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.telemetry import Feature, _telemetry_emitter
 
 @dataclass
 class AthenaQuery:
@@ -37,6 +38,7 @@ class AthenaQuery:
     _result_bucket: str = field(default=None, init=False)
     _result_file_prefix: str = field(default=None, init=False)
 
+    @_telemetry_emitter(Feature.FEATURE_STORE, "AthenaQuery.run")
     def run(
         self, query_string: str, output_location: str, kms_key: str = None, workgroup: str = None
     ) -> str:
@@ -82,6 +84,7 @@ def get_query_execution(self) -> Dict[str, Any]:
         """
         return get_query_execution(self.sagemaker_session, self._current_query_execution_id)
 
+    @_telemetry_emitter(Feature.FEATURE_STORE, "AthenaQuery.as_dataframe")
     def as_dataframe(self, **kwargs) -> DataFrame:
         """Download the result of the current query and load it into a DataFrame.
 

sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py

@@ -9,6 +9,7 @@
 import pandas as pd
 
 from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.telemetry import Feature, _telemetry_emitter
 from sagemaker.mlops.feature_store import FeatureGroup
 from sagemaker.mlops.feature_store.feature_definition import FeatureDefinition, FeatureTypeEnum
 from sagemaker.mlops.feature_store.feature_utils import (
@@ -422,6 +423,7 @@ def with_event_time_range(
         self._event_time_ending_timestamp = ending_timestamp
         return self
 
+    @_telemetry_emitter(Feature.FEATURE_STORE, "DatasetBuilder.to_csv_file")
     def to_csv_file(self) -> tuple[str, str]:
         """Get query string and result in .csv format file.
 
@@ -435,6 +437,7 @@ def to_csv_file(self) -> tuple[str, str]:
             return self._to_csv_from_feature_group()
         raise ValueError("Base must be either a FeatureGroup or a DataFrame.")
 
+    @_telemetry_emitter(Feature.FEATURE_STORE, "DatasetBuilder.to_dataframe")
     def to_dataframe(self) -> tuple[pd.DataFrame, str]:
         """Get query string and result in pandas.DataFrame.
 

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py

@@ -16,6 +16,7 @@
 from sagemaker.mlops.feature_store import FeatureGroup as CoreFeatureGroup, FeatureGroup
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.core.s3.client import S3Uploader, S3Downloader
+from sagemaker.core.telemetry import Feature, _telemetry_emitter
 from sagemaker.mlops.feature_store.feature_definition import (
     FeatureDefinition,
     FractionalFeatureDefinition,
@@ -420,6 +421,7 @@ def as_hive_ddl(
     return ddl
 
 
+@_telemetry_emitter(Feature.FEATURE_STORE, "ingest_dataframe")
 def ingest_dataframe(
     feature_group_name: str,
     data_frame: DataFrame,
@@ -471,6 +473,7 @@ def ingest_dataframe(
     manager.run(data_frame=data_frame, wait=wait, timeout=timeout)
     return manager
 
+@_telemetry_emitter(Feature.FEATURE_STORE, "get_feature_group_as_dataframe")
 def get_feature_group_as_dataframe(
     feature_group_name: str,
     athena_bucket: str,

sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py

@@ -16,6 +16,7 @@
 from sagemaker.core.resources import FeatureGroup as CoreFeatureGroup
 from sagemaker.core.shapes import FeatureValue
 from sagemaker.core.utils.utils import Unassigned
+from sagemaker.core.telemetry import Feature, _telemetry_emitter
 
 logger = logging.getLogger(__name__)
 
@@ -67,6 +68,7 @@ def failed_rows(self) -> List[int]:
         """
         return self._failed_indices
 
+    @_telemetry_emitter(Feature.FEATURE_STORE, "IngestionManagerPandas.run")
     def run(
         self,
         data_frame: DataFrame,