Feature store v3

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(