feature: support online store ttl for records (#4030)

* feature: support online store ttl for records

* Added integ tests for online store ttl

---------

Co-authored-by: Guifan Weng <[email protected]>

Author: gwengusc

doc/amazon_sagemaker_featurestore.rst

@@ -202,6 +202,7 @@ location of your offline store.
        role_arn = role,
        s3_uri = offline_feature_store_bucket,
        enable_online_store = True,
+       ttl_duration = None,
        online_store_kms_key_id = None,
        offline_store_kms_key_id = None,
        disable_glue_table_creation = False,

src/sagemaker/feature_store/feature_group.py

@@ -61,6 +61,8 @@
     FeatureParameter,
     TableFormatEnum,
     DeletionModeEnum,
+    TtlDuration,
+    OnlineStoreConfigUpdate,
 )
 from sagemaker.utils import resolve_value_from_config
 
@@ -523,6 +525,7 @@ def create(
         role_arn: str = None,
         online_store_kms_key_id: str = None,
         enable_online_store: bool = False,
+        ttl_duration: TtlDuration = None,
         offline_store_kms_key_id: str = None,
         disable_glue_table_creation: bool = False,
         data_catalog_config: DataCatalogConfig = None,
@@ -539,6 +542,7 @@ def create(
             event_time_feature_name (str): name of the event time feature.
             role_arn (str): ARN of the role used to call CreateFeatureGroup.
             online_store_kms_key_id (str): KMS key ARN for online store (default: None).
+            ttl_duration (TtlDuration): Default time to live duration for records (default: None).
             enable_online_store (bool): whether to enable online store or not (default: False).
             offline_store_kms_key_id (str): KMS key ARN for offline store (default: None).
                 If a KMS encryption key is not specified, SageMaker encrypts all data at
@@ -592,7 +596,10 @@ def create(
 
         # online store configuration
         if enable_online_store:
-            online_store_config = OnlineStoreConfig(enable_online_store=enable_online_store)
+            online_store_config = OnlineStoreConfig(
+                enable_online_store=enable_online_store,
+                ttl_duration=ttl_duration,
+            )
             if online_store_kms_key_id is not None:
                 online_store_config.online_store_security_config = OnlineStoreSecurityConfig(
                     kms_key_id=online_store_kms_key_id
@@ -633,21 +640,37 @@ def describe(self, next_token: str = None) -> Dict[str, Any]:
             feature_group_name=self.name, next_token=next_token
         )
 
-    def update(self, feature_additions: Sequence[FeatureDefinition]) -> Dict[str, Any]:
+    def update(
+        self,
+        feature_additions: Sequence[FeatureDefinition] = None,
+        online_store_config: OnlineStoreConfigUpdate = None,
+    ) -> Dict[str, Any]:
         """Update a FeatureGroup and add new features from the given feature definitions.
 
         Args:
             feature_additions (Sequence[Dict[str, str]): list of feature definitions to be updated.
+            online_store_config (OnlineStoreConfigUpdate): online store config to be updated.
 
         Returns:
             Response dict from service.
         """
 
+        if feature_additions is None:
+            feature_additions_parameter = None
+        else:
+            feature_additions_parameter = [
+                feature_addition.to_dict() for feature_addition in feature_additions
+            ]
+
+        if online_store_config is None:
+            online_store_config_parameter = None
+        else:
+            online_store_config_parameter = online_store_config.to_dict()
+
         return self.sagemaker_session.update_feature_group(
             feature_group_name=self.name,
-            feature_additions=[
-                feature_addition.to_dict() for feature_addition in feature_additions
-            ],
+            feature_additions=feature_additions_parameter,
+            online_store_config=online_store_config_parameter,
         )
 
     def update_feature_metadata(
@@ -756,7 +779,9 @@ def load_feature_definitions(
         return self.feature_definitions
 
     def get_record(
-        self, record_identifier_value_as_string: str, feature_names: Sequence[str] = None
+        self,
+        record_identifier_value_as_string: str,
+        feature_names: Sequence[str] = None,
     ) -> Sequence[Dict[str, str]]:
         """Get a single record in a FeatureGroup
 
@@ -772,14 +797,24 @@ def get_record(
             feature_names=feature_names,
         ).get("Record")
 
-    def put_record(self, record: Sequence[FeatureValue]):
+    def put_record(self, record: Sequence[FeatureValue], ttl_duration: TtlDuration = None):
         """Put a single record in the FeatureGroup.
 
         Args:
             record (Sequence[FeatureValue]): a list contains feature values.
+            ttl_duration (TtlDuration): customer specified ttl duration.
         """
+
+        if ttl_duration is not None:
+            return self.sagemaker_session.put_record(
+                feature_group_name=self.name,
+                record=[value.to_dict() for value in record],
+                ttl_duration=ttl_duration.to_dict(),
+            )
+
         return self.sagemaker_session.put_record(
-            feature_group_name=self.name, record=[value.to_dict() for value in record]
+            feature_group_name=self.name,
+            record=[value.to_dict() for value in record],
         )
 
     def delete_record(

src/sagemaker/feature_store/feature_store.py

@@ -137,18 +137,27 @@ def list_feature_groups(
             next_token=next_token,
         )
 
-    def batch_get_record(self, identifiers: Sequence[Identifier]) -> Dict[str, Any]:
+    def batch_get_record(
+        self,
+        identifiers: Sequence[Identifier],
+        expiration_time_response: str = None,
+    ) -> Dict[str, Any]:
         """Get record in batch from FeatureStore
 
         Args:
             identifiers (Sequence[Identifier]): A list of identifiers to uniquely identify records
                 in FeatureStore.
+            expiration_time_response (str): the field of expiration time response
+                to toggle returning of expiresAt.
 
         Returns:
             Response dict from service.
         """
         batch_get_record_identifiers = [identifier.to_dict() for identifier in identifiers]
-        return self.sagemaker_session.batch_get_record(identifiers=batch_get_record_identifiers)
+        return self.sagemaker_session.batch_get_record(
+            identifiers=batch_get_record_identifiers,
+            expiration_time_response=expiration_time_response,
+        )
 
     def search(
         self,

src/sagemaker/feature_store/inputs.py

@@ -84,17 +84,43 @@ def to_dict(self) -> Dict[str, Any]:
         return Config.construct_dict(KmsKeyId=self.kms_key_id)
 
 
+@attr.s
+class TtlDuration(Config):
+    """TtlDuration for records in online FeatureStore.
+
+    Attributes:
+        unit (str): time unit.
+        value (int): time value.
+    """
+
+    unit: str = attr.ib()
+    value: int = attr.ib()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Construct a dictionary based on the attributes.
+
+        Returns:
+            dict represents the attributes.
+        """
+        return Config.construct_dict(
+            Unit=self.unit,
+            Value=self.value,
+        )
+
+
 @attr.s
 class OnlineStoreConfig(Config):
     """OnlineStoreConfig for FeatureStore.
 
     Attributes:
         enable_online_store (bool): whether to enable the online store.
         online_store_security_config (OnlineStoreSecurityConfig): configuration of security setting.
+        ttl_duration (TtlDuration): Default time to live duration for records.
     """
 
     enable_online_store: bool = attr.ib(default=True)
     online_store_security_config: OnlineStoreSecurityConfig = attr.ib(default=None)
+    ttl_duration: TtlDuration = attr.ib(default=None)
 
     def to_dict(self) -> Dict[str, Any]:
         """Construct a dictionary based on the attributes.
@@ -105,6 +131,28 @@ def to_dict(self) -> Dict[str, Any]:
         return Config.construct_dict(
             EnableOnlineStore=self.enable_online_store,
             SecurityConfig=self.online_store_security_config,
+            TtlDuration=self.ttl_duration,
+        )
+
+
+@attr.s
+class OnlineStoreConfigUpdate(Config):
+    """OnlineStoreConfigUpdate for FeatureStore.
+
+    Attributes:
+        ttl_duration (TtlDuration): Default time to live duration for records.
+    """
+
+    ttl_duration: TtlDuration = attr.ib(default=None)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Construct a dictionary based on the attributes.
+
+        Returns:
+            dict represents the attributes.
+        """
+        return Config.construct_dict(
+            TtlDuration=self.ttl_duration,
         )
 
 
@@ -379,3 +427,13 @@ class DeletionModeEnum(Enum):
 
     SOFT_DELETE = "SoftDelete"
     HARD_DELETE = "HardDelete"
+
+
+class ExpirationTimeResponseEnum(Enum):
+    """Enum of toggling the response of ExpiresAt.
+
+    The ExpirationTimeResponse for toggling the response of ExpiresAt can be Disabled or Enabled.
+    """
+
+    DISABLED = "Disabled"
+    ENABLED = "Enabled"

src/sagemaker/session.py

@@ -2639,7 +2639,7 @@ def tune(  # noqa: C901
             warm_start_config (dict): Configuration defining the type of warm start and
                 other required configurations.
             max_runtime_in_seconds (int or PipelineVariable): The maximum time in seconds
-                that a hyperparameter tuning job can run.
+                that a training job launched by a hyperparameter tuning job can run.
             completion_criteria_config (sagemaker.tuner.TuningJobCompletionCriteriaConfig): A
                 configuration for the completion criteria.
             early_stopping_type (str): Specifies whether early stopping is enabled for the job.
@@ -2894,7 +2894,7 @@ def _map_tuning_config(
                 tuning job.
             max_parallel_jobs (int): Maximum number of parallel training jobs to start.
             max_runtime_in_seconds (int or PipelineVariable): The maximum time in seconds
-                that a hyperparameter tuning job can run.
+                that a training job launched by a hyperparameter tuning job can run.
             early_stopping_type (str): Specifies whether early stopping is enabled for the job.
                 Can be either 'Auto' or 'Off'. If set to 'Off', early stopping will not be
                 attempted. If set to 'Auto', early stopping of some training jobs may happen,
@@ -5097,9 +5097,15 @@ def describe_feature_group(
         return self.sagemaker_client.describe_feature_group(**kwargs)
 
     def update_feature_group(
-        self, feature_group_name: str, feature_additions: Sequence[Dict[str, str]]
+        self,
+        feature_group_name: str,
+        feature_additions: Sequence[Dict[str, str]] = None,
+        online_store_config: Dict[str, any] = None,
     ) -> Dict[str, Any]:
-        """Update a FeatureGroup and add new features from the given feature definitions.
+        """Update a FeatureGroup
+
+            either adding new features from the given feature definitions
+            or updating online store config
 
         Args:
             feature_group_name (str): name of the FeatureGroup to update.
@@ -5108,6 +5114,12 @@ def update_feature_group(
             Response dict from service.
         """
 
+        if feature_additions is None:
+            return self.sagemaker_client.update_feature_group(
+                FeatureGroupName=feature_group_name,
+                OnlineStoreConfig=online_store_config,
+            )
+
         return self.sagemaker_client.update_feature_group(
             FeatureGroupName=feature_group_name, FeatureAdditions=feature_additions
         )
@@ -5258,6 +5270,7 @@ def put_record(
         self,
         feature_group_name: str,
         record: Sequence[Dict[str, str]],
+        ttl_duration: Dict[str, str] = None,
     ):
         """Puts a single record in the FeatureGroup.
 
@@ -5266,6 +5279,14 @@ def put_record(
             record (Sequence[Dict[str, str]]): list of FeatureValue dicts to be ingested
                 into FeatureStore.
         """
+
+        if ttl_duration:
+            return self.sagemaker_featurestore_runtime_client.put_record(
+                FeatureGroupName=feature_group_name,
+                Record=record,
+                TtlDuration=ttl_duration,
+            )
+
         return self.sagemaker_featurestore_runtime_client.put_record(
             FeatureGroupName=feature_group_name,
             Record=record,
@@ -5298,36 +5319,51 @@ def get_record(
         record_identifier_value_as_string: str,
         feature_group_name: str,
         feature_names: Sequence[str],
+        expiration_time_response: str = None,
     ) -> Dict[str, Sequence[Dict[str, str]]]:
         """Gets a single record in the FeatureGroup.
 
         Args:
             record_identifier_value_as_string (str): name of the record identifier.
             feature_group_name (str): name of the FeatureGroup.
             feature_names (Sequence[str]): list of feature names.
+            expiration_time_response (str): the field of expiration time response
+                to toggle returning of expiresAt.
         """
         get_record_args = {
             "FeatureGroupName": feature_group_name,
             "RecordIdentifierValueAsString": record_identifier_value_as_string,
         }
 
+        if expiration_time_response:
+            get_record_args["ExpirationTimeResponse"] = expiration_time_response
+
         if feature_names:
             get_record_args["FeatureNames"] = feature_names
 
         return self.sagemaker_featurestore_runtime_client.get_record(**get_record_args)
 
-    def batch_get_record(self, identifiers: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
+    def batch_get_record(
+        self,
+        identifiers: Sequence[Dict[str, Any]],
+        expiration_time_response: str = None,
+    ) -> Dict[str, Any]:
         """Gets a batch of record from FeatureStore.
 
         Args:
             identifiers (Sequence[Dict[str, Any]]): list of identifiers to uniquely identify records
                 in FeatureStore.
+            expiration_time_response (str): the field of expiration time response
+                to toggle returning of expiresAt.
 
         Returns:
             Response dict from service.
         """
         batch_get_record_args = {"Identifiers": identifiers}
 
+        if expiration_time_response:
+            batch_get_record_args["ExpirationTimeResponse"] = expiration_time_response
+
         return self.sagemaker_featurestore_runtime_client.batch_get_record(**batch_get_record_args)
 
     def start_query_execution(