fix: Allow online store only FeatureGroups (#2015)

Co-authored-by: Rui Wang Napieralski <[email protected]>

Author: icywang86rui

src/sagemaker/feature_store/feature_group.py

@@ -290,7 +290,7 @@ class FeatureGroup:
 
     def create(
         self,
-        s3_uri: str,
+        s3_uri: Union[str, bool],
         record_identifier_name: str,
         event_time_feature_name: str,
         role_arn: str,
@@ -305,7 +305,8 @@ def create(
         """Create a SageMaker FeatureStore FeatureGroup.
 
         Args:
-            s3_uri (str): S3 URI of the offline store.
+            s3_uri (Union[str, bool]): S3 URI of the offline store, set to
+                ``False`` to disable offline store.
             record_identifier_name (str): name of the record identifier feature.
             event_time_feature_name (str): name of the event time feature.
             role_arn (str): ARN of the role used to call CreateFeatureGroup.
@@ -342,15 +343,18 @@ def create(
             create_feature_store_args.update({"online_store_config": online_store_config.to_dict()})
 
         # offline store configuration
-        s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
-        if offline_store_kms_key_id:
-            s3_storage_config.kms_key_id = offline_store_kms_key_id
-        offline_store_config = OfflineStoreConfig(
-            s3_storage_config=s3_storage_config,
-            disable_glue_table_creation=disable_glue_table_creation,
-            data_catalog_config=data_catalog_config,
-        )
-        create_feature_store_args.update({"offline_store_config": offline_store_config.to_dict()})
+        if s3_uri:
+            s3_storage_config = S3StorageConfig(s3_uri=s3_uri)
+            if offline_store_kms_key_id:
+                s3_storage_config.kms_key_id = offline_store_kms_key_id
+            offline_store_config = OfflineStoreConfig(
+                s3_storage_config=s3_storage_config,
+                disable_glue_table_creation=disable_glue_table_creation,
+                data_catalog_config=data_catalog_config,
+            )
+            create_feature_store_args.update(
+                {"offline_store_config": offline_store_config.to_dict()}
+            )
 
         return self.sagemaker_session.create_feature_group(**create_feature_store_args)
 
@@ -367,7 +371,9 @@ def describe(self, next_token: str = None) -> Dict[str, Any]:
         Returns:
             Response dict from the service.
         """
-        return self.sagemaker_session.describe_feature_group(self.name, next_token)
+        return self.sagemaker_session.describe_feature_group(
+            feature_group_name=self.name, next_token=next_token
+        )
 
     def load_feature_definitions(
         self,

tests/integ/test_feature_store.py

@@ -141,6 +141,28 @@ def create_table_ddl():
     )
 
 
+def test_create_feature_store_online_only(
+    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):
+        output = feature_group.create(
+            s3_uri=False,
+            record_identifier_name="feature1",
+            event_time_feature_name="feature3",
+            role_arn=role,
+            enable_online_store=True,
+        )
+        _wait_for_feature_group_create(feature_group)
+
+    assert output["FeatureGroupArn"].endswith(f"feature-group/{feature_group_name}")
+
+
 def test_create_feature_store(
     feature_store_session,
     role,

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -79,36 +79,66 @@ def test_feature_store_create(
         role_arn=role_arn,
         enable_online_store=True,
     )
-    assert sagemaker_session_mock.create_feature_group.called_with(
+    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},
+        offline_store_config={
+            "DisableGlueTableCreation": False,
+            "S3StorageConfig": {"S3Uri": s3_uri},
+        },
+    )
+
+
+def test_feature_store_create_online_only(
+    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,
+    )
+    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},
     )
 
 
 def test_feature_store_delete(sagemaker_session_mock):
     feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
     feature_group.delete()
-    assert sagemaker_session_mock.delete_feature_group.called_with(
+    sagemaker_session_mock.delete_feature_group.assert_called_with(
         feature_group_name="MyFeatureGroup"
     )
 
 
 def test_feature_store_describe(sagemaker_session_mock):
     feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
     feature_group.describe()
-    assert sagemaker_session_mock.describe_feature_group.called_with(
-        feature_group_name="MyFeatureGroup"
+    sagemaker_session_mock.describe_feature_group.assert_called_with(
+        feature_group_name="MyFeatureGroup", next_token=None
     )
 
 
 def test_put_record(sagemaker_session_mock):
     feature_group = FeatureGroup(name="MyFeatureGroup", sagemaker_session=sagemaker_session_mock)
     feature_group.put_record(record=[])
-    assert sagemaker_session_mock.put_record.called_with(
+    sagemaker_session_mock.put_record.assert_called_with(
         feature_group_name="MyFeatureGroup", record=[]
     )
 
@@ -268,7 +298,7 @@ def query(sagemaker_session_mock):
 def test_athena_query_run(sagemaker_session_mock, query):
     sagemaker_session_mock.start_query_execution.return_value = {"QueryExecutionId": "query_id"}
     query.run(query_string="query", output_location="s3://some-bucket/some-path")
-    assert sagemaker_session_mock.start_query_execution.called_with(
+    sagemaker_session_mock.start_query_execution.assert_called_with(
         catalog="catalog",
         database="database",
         query_string="query",
@@ -283,13 +313,13 @@ def test_athena_query_run(sagemaker_session_mock, query):
 def test_athena_query_wait(sagemaker_session_mock, query):
     query._current_query_execution_id = "query_id"
     query.wait()
-    assert sagemaker_session_mock.wait_for_athena_query.called_with(query_execution_id="query_id")
+    sagemaker_session_mock.wait_for_athena_query.assert_called_with(query_execution_id="query_id")
 
 
 def test_athena_query_get_query_execution(sagemaker_session_mock, query):
     query._current_query_execution_id = "query_id"
     query.get_query_execution()
-    assert sagemaker_session_mock.wait_for_athena_query.called_with(query_execution_id="query_id")
+    sagemaker_session_mock.get_query_execution.assert_called_with(query_execution_id="query_id")
 
 
 @patch("tempfile.gettempdir", Mock(return_value="tmp"))
@@ -302,13 +332,13 @@ def test_athena_query_as_dataframe(read_csv, sagemaker_session_mock, query):
     query._result_bucket = "bucket"
     query._result_file_prefix = "prefix"
     query.as_dataframe()
-    assert sagemaker_session_mock.download_athena_query_result.called_with(
+    sagemaker_session_mock.download_athena_query_result.assert_called_with(
         bucket="bucket",
         prefix="prefix",
         query_execution_id="query_id",
         filename="tmp/query_id.csv",
     )
-    assert read_csv.called_with("tmp/query_id.csv", delimiter=",")
+    read_csv.assert_called_with("tmp/query_id.csv", delimiter=",")
 
 
 @patch("tempfile.gettempdir", Mock(return_value="tmp"))