feature: Add support for online store storage type in create feature … (aws#1098)

Author: AWS-pratab

doc/api/prep_data/feature_store.rst

@@ -77,6 +77,10 @@ Inputs
     :members:
     :show-inheritance:
 
+.. autoclass:: sagemaker.feature_store.inputs.OnlineStoreStorageTypeEnum
+    :members:
+    :show-inheritance:
+
 .. autoclass:: sagemaker.feature_store.inputs.ResourceEnum
     :members:
     :show-inheritance:

src/sagemaker/feature_store/feature_group.py

@@ -63,6 +63,7 @@
     DeletionModeEnum,
     TtlDuration,
     OnlineStoreConfigUpdate,
+    OnlineStoreStorageTypeEnum,
 )
 from sagemaker.utils import resolve_value_from_config
 
@@ -539,6 +540,7 @@ def create(
         description: str = None,
         tags: List[Dict[str, str]] = None,
         table_format: TableFormatEnum = None,
+        online_store_storage_type: OnlineStoreStorageTypeEnum = None,
     ) -> Dict[str, Any]:
         """Create a SageMaker FeatureStore FeatureGroup.
 
@@ -566,6 +568,8 @@ def create(
             description (str): description of the FeatureGroup (default: None).
             tags (List[Dict[str, str]]): list of tags for labeling a FeatureGroup (default: None).
             table_format (TableFormatEnum): format of the offline store table (default: None).
+            online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
+                online store (default: None).
 
         Returns:
             Response dict from service.
@@ -606,6 +610,7 @@ def create(
             online_store_config = OnlineStoreConfig(
                 enable_online_store=enable_online_store,
                 ttl_duration=ttl_duration,
+                storage_type=online_store_storage_type,
             )
             if online_store_kms_key_id is not None:
                 online_store_config.online_store_security_config = OnlineStoreSecurityConfig(

src/sagemaker/feature_store/inputs.py

@@ -108,6 +108,16 @@ def to_dict(self) -> Dict[str, Any]:
         )
 
 
+class OnlineStoreStorageTypeEnum(Enum):
+    """Enum of storage types for online store.
+
+    The online store storage types can be Standard or InMemory.
+    """
+
+    STANDARD = "Standard"
+    IN_MEMORY = "InMemory"
+
+
 @attr.s
 class OnlineStoreConfig(Config):
     """OnlineStoreConfig for FeatureStore.
@@ -121,6 +131,7 @@ class OnlineStoreConfig(Config):
     enable_online_store: bool = attr.ib(default=True)
     online_store_security_config: OnlineStoreSecurityConfig = attr.ib(default=None)
     ttl_duration: TtlDuration = attr.ib(default=None)
+    storage_type: OnlineStoreStorageTypeEnum = attr.ib(default=None)
 
     def to_dict(self) -> Dict[str, Any]:
         """Construct a dictionary based on the attributes.
@@ -132,6 +143,7 @@ def to_dict(self) -> Dict[str, Any]:
             EnableOnlineStore=self.enable_online_store,
             SecurityConfig=self.online_store_security_config,
             TtlDuration=self.ttl_duration,
+            StorageType=self.storage_type.value if self.storage_type else None,
         )
 
 

tests/integ/test_feature_store.py

@@ -38,6 +38,7 @@
     DeletionModeEnum,
     TtlDuration,
     OnlineStoreConfigUpdate,
+    OnlineStoreStorageTypeEnum,
 )
 from sagemaker.feature_store.dataset_builder import (
     JoinTypeEnum,
@@ -347,6 +348,54 @@ def test_create_feature_group_glue_table_format(
         assert table_format == "Glue"
 
 
+def test_create_feature_group_in_memory_storage_type(
+    feature_store_session,
+    role,
+    feature_group_name,
+    pandas_data_frame,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=False,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+            online_store_storage_type=OnlineStoreStorageTypeEnum.IN_MEMORY,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        storage_type = feature_group.describe().get("OnlineStoreConfig").get("StorageType")
+        assert storage_type == "InMemory"
+
+
+def test_create_feature_group_standard_storage_type(
+    feature_store_session,
+    role,
+    feature_group_name,
+    pandas_data_frame,
+):
+    feature_group = FeatureGroup(name=feature_group_name, sagemaker_session=feature_store_session)
+    feature_group.load_feature_definitions(data_frame=pandas_data_frame)
+
+    with cleanup_feature_group(feature_group):
+        feature_group.create(
+            s3_uri=False,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+            online_store_storage_type=OnlineStoreStorageTypeEnum.STANDARD,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+        storage_type = feature_group.describe().get("OnlineStoreConfig").get("StorageType")
+        assert storage_type == "Standard"
+
+
 def test_ttl_duration(
     feature_store_session,
     role,

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -36,6 +36,7 @@
     DeletionModeEnum,
     TtlDuration,
     OnlineStoreConfigUpdate,
+    OnlineStoreStorageTypeEnum,
 )
 
 from tests.unit import SAGEMAKER_CONFIG_FEATURE_GROUP
@@ -236,6 +237,31 @@ def test_feature_store_create_online_only(
     )
 
 
+def test_feature_store_create_online_only_with_in_memory(
+    sagemaker_session_mock, role_arn, feature_group_dummy_definitions
+):
+    feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
+    feature_group.feature_definitions = feature_group_dummy_definitions
+    feature_group.create(
+        s3_uri=False,
+        record_identifier_name="feature1",
+        event_time_feature_name="feature2",
+        role_arn=role_arn,
+        enable_online_store=True,
+        online_store_storage_type=OnlineStoreStorageTypeEnum.IN_MEMORY,
+    )
+    sagemaker_session_mock.create_feature_group.assert_called_with(
+        feature_group_name="MyFeatureGroup",
+        record_identifier_name="feature1",
+        event_time_feature_name="feature2",
+        feature_definitions=[fd.to_dict() for fd in feature_group_dummy_definitions],
+        role_arn=role_arn,
+        description=None,
+        tags=None,
+        online_store_config={"EnableOnlineStore": True, "StorageType": "InMemory"},
+    )
+
+
 def test_feature_store_delete(sagemaker_session_mock):
     feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
     feature_group.delete()

tests/unit/sagemaker/feature_store/test_inputs.py

@@ -21,6 +21,7 @@
     OfflineStoreConfig,
     FeatureParameter,
     TableFormatEnum,
+    OnlineStoreStorageTypeEnum,
     Filter,
     FilterOperatorEnum,
     Identifier,
@@ -59,6 +60,32 @@ def test_online_store_config():
         }
     )
 
+    config_with_inmemory = OnlineStoreConfig(
+        enable_online_store=True,
+        storage_type=OnlineStoreStorageTypeEnum.IN_MEMORY,
+    )
+    assert ordered(config_with_inmemory.to_dict()) == ordered(
+        {
+            "EnableOnlineStore": True,
+            "StorageType": "InMemory",
+        }
+    )
+
+    config_with_kms_standard = OnlineStoreConfig(
+        enable_online_store=True,
+        online_store_security_config=OnlineStoreSecurityConfig(kms_key_id="kms"),
+        storage_type=OnlineStoreStorageTypeEnum.STANDARD,
+    )
+    assert ordered(config_with_kms_standard.to_dict()) == ordered(
+        {
+            "EnableOnlineStore": True,
+            "SecurityConfig": {
+                "KmsKeyId": "kms",
+            },
+            "StorageType": "Standard",
+        }
+    )
+
 
 def test_s3_store_config():
     config = S3StorageConfig(s3_uri="uri", kms_key_id="kms")