refactor(feature-processor): Migrate to FeatureGroup resource API

- Replace sagemaker_session.describe_feature_group() calls with FeatureGroup.get()
- Update _input_loader.py to use FeatureGroup resource attributes instead of dictionary access
- Update feature_scheduler.py to use FeatureGroup.get() and access creation_time as attribute
- Update _feature_group_lineage_entity_handler.py to return FeatureGroup resource instead of Dict
- Remove unused imports (Dict, Any, FEATURE_GROUP, CREATION_TIME constants)
- Replace dictionary key access with typed resource properties (offline_store_config, data_catalog_config, event_time_feature_name, etc.)
- Update unit tests to reflect new FeatureGroup resource API usage
- Improves type safety and reduces reliance on dictionary-based API responses

Author: bhhalim

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/_input_loader.py

@@ -34,6 +34,7 @@
     InputOffsetParser,
 )
 from sagemaker.mlops.feature_store.feature_processor._env import EnvironmentHelper
+from sagemaker.core.resources import FeatureGroup
 
 T = TypeVar("T")
 
@@ -96,26 +97,30 @@ def load_from_feature_group(
         sagemaker_session: Session = self.sagemaker_session or Session()
 
         feature_group_name = feature_group_data_source.name
-        feature_group = sagemaker_session.describe_feature_group(
-            self._parse_name_from_arn(feature_group_name)
+        feature_group = FeatureGroup.get(
+            feature_group_name=self._parse_name_from_arn(feature_group_name),
+            session=sagemaker_session.boto_session,
         )
         logger.debug(
             "Called describe_feature_group with %s and received: %s",
             feature_group_name,
             feature_group,
         )
 
-        if "OfflineStoreConfig" not in feature_group:
+        if not feature_group.offline_store_config:
             raise ValueError(
                 f"Input Feature Groups must have an enabled Offline Store."
                 f" Feature Group: {feature_group_name} does not have an Offline Store enabled."
             )
 
-        offline_store_uri = feature_group["OfflineStoreConfig"]["S3StorageConfig"][
-            "ResolvedOutputS3Uri"
-        ]
+        offline_store_config = feature_group.offline_store_config
+        offline_store_uri = offline_store_config.s3_storage_config.resolved_output_s3_uri
 
-        table_format = feature_group["OfflineStoreConfig"].get("TableFormat", None)
+        table_format = (
+            offline_store_config.table_format
+            if offline_store_config.table_format
+            else None
+        )
 
         if table_format not in self._supported_table_format:
             raise ValueError(
@@ -127,15 +132,15 @@ def load_from_feature_group(
         end_offset = feature_group_data_source.input_end_offset
 
         if table_format == "Iceberg":
-            data_catalog_config = feature_group["OfflineStoreConfig"]["DataCatalogConfig"]
+            data_catalog_config = offline_store_config.data_catalog_config
             return self.load_from_iceberg_table(
                 IcebergTableDataSource(
                     offline_store_uri,
-                    data_catalog_config["Catalog"],
-                    data_catalog_config["Database"],
-                    data_catalog_config["TableName"],
+                    data_catalog_config.catalog,
+                    data_catalog_config.database,
+                    data_catalog_config.table_name,
                 ),
-                feature_group["EventTimeFeatureName"],
+                feature_group.event_time_feature_name,
                 start_offset,
                 end_offset,
             )

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/feature_scheduler.py

@@ -73,6 +73,7 @@
 )
 
 from sagemaker.core.s3 import s3_path_join
+from sagemaker.core.resources import FeatureGroup
 
 from sagemaker.core.helper.session_helper import Session, get_execution_role
 from sagemaker.mlops.feature_store.feature_processor._event_bridge_scheduler_helper import (
@@ -770,8 +771,8 @@ def _validate_fg_lineage_resources(feature_group_name: str, sagemaker_session: S
         groups.
     """
 
-    feature_group = sagemaker_session.describe_feature_group(feature_group_name=feature_group_name)
-    feature_group_creation_time = feature_group["CreationTime"].strftime("%s")
+    feature_group = FeatureGroup.get(feature_group_name=feature_group_name, session=sagemaker_session.boto_session)
+    feature_group_creation_time = feature_group.creation_time.strftime("%s")
     feature_group_context = _get_feature_group_lineage_context_name(
         feature_group_name=feature_group_name,
         feature_group_creation_time=feature_group_creation_time,

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/lineage/_feature_group_lineage_entity_handler.py

@@ -14,7 +14,6 @@
 from __future__ import absolute_import
 
 import re
-from typing import Dict, Any
 import logging
 
 from sagemaker.core.helper.session_helper import Session
@@ -24,10 +23,9 @@
 )
 from sagemaker.mlops.feature_store.feature_processor.lineage.constants import (
     SAGEMAKER,
-    FEATURE_GROUP,
-    CREATION_TIME,
 )
 from sagemaker.core.lineage.context import Context
+from sagemaker.core.resources import FeatureGroup
 
 # pylint: disable=C0301
 from sagemaker.mlops.feature_store.feature_processor.lineage._feature_processor_lineage_name_helper import (
@@ -62,8 +60,8 @@ def retrieve_feature_group_context_arns(
             ),
             sagemaker_session=sagemaker_session,
         )
-        feature_group_name = feature_group[FEATURE_GROUP]
-        feature_group_creation_time = feature_group[CREATION_TIME].strftime("%s")
+        feature_group_name = feature_group.feature_group_name
+        feature_group_creation_time = feature_group.creation_time.strftime("%s")
         feature_group_pipeline_context = (
             FeatureGroupLineageEntityHandler._load_feature_group_pipeline_context(
                 feature_group_name=feature_group_name,
@@ -87,7 +85,7 @@ def retrieve_feature_group_context_arns(
     @staticmethod
     def _describe_feature_group(
         feature_group_name: str, sagemaker_session: Session
-    ) -> Dict[str, Any]:
+    ) -> FeatureGroup:
         """Retrieve the Feature Group.
 
         Arguments:
@@ -97,9 +95,9 @@ def _describe_feature_group(
                 function creates one using the default AWS configuration chain.
 
         Returns:
-            Dict[str, Any]: The Feature Group details.
+            FeatureGroup: The Feature Group resource.
         """
-        feature_group = sagemaker_session.describe_feature_group(feature_group_name)
+        feature_group = FeatureGroup.get(feature_group_name=feature_group_name, session=sagemaker_session.boto_session)
         logger.debug(
             "Called describe_feature_group with %s and received: %s",
             feature_group_name,

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_constants.py

@@ -46,6 +46,7 @@
 CREATION_TIME = "123123123"
 LAST_UPDATE_TIME = "234234234"
 SAGEMAKER_SESSION_MOCK = Mock(Session)
+SAGEMAKER_SESSION_MOCK.boto_session = Mock()
 CONTEXT_MOCK_01 = Mock(Context)
 CONTEXT_MOCK_02 = Mock(Context)
 

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/lineage/test_feature_group_lineage_entity_handler.py

@@ -11,12 +11,13 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
-from mock import patch, call
+from mock import patch, call, Mock
 
 from sagemaker.mlops.feature_store.feature_processor.lineage._feature_group_lineage_entity_handler import (
     FeatureGroupLineageEntityHandler,
 )
 from sagemaker.core.lineage.context import Context
+from sagemaker.core.resources import FeatureGroup
 
 from test_constants import (
     SAGEMAKER_SESSION_MOCK,
@@ -26,11 +27,15 @@
     FEATURE_GROUP_NAME,
 )
 
+FEATURE_GROUP_MOCK = Mock()
+FEATURE_GROUP_MOCK.feature_group_name = FEATURE_GROUP["FeatureGroupName"]
+FEATURE_GROUP_MOCK.creation_time = FEATURE_GROUP["CreationTime"]
+
 
 def test_retrieve_feature_group_context_arns():
     with patch.object(
-        SAGEMAKER_SESSION_MOCK, "describe_feature_group", return_value=FEATURE_GROUP
-    ) as fg_describe_method:
+        FeatureGroup, "get", return_value=FEATURE_GROUP_MOCK
+    ) as fg_get_method:
         with patch.object(
             Context, "load", side_effect=[CONTEXT_MOCK_01, CONTEXT_MOCK_02]
         ) as context_load:
@@ -44,16 +49,17 @@ def test_retrieve_feature_group_context_arns():
     assert result.name == FEATURE_GROUP_NAME
     assert result.pipeline_context_arn == "context-arn-fep"
     assert result.pipeline_version_context_arn == "context-arn-fep-ver"
-    fg_describe_method.assert_called_once_with(FEATURE_GROUP_NAME)
+    fg_get_method.assert_called_once_with(feature_group_name=FEATURE_GROUP_NAME, session=SAGEMAKER_SESSION_MOCK.boto_session)
+    creation_time_str = FEATURE_GROUP_MOCK.creation_time.strftime("%s")
     context_load.assert_has_calls(
         [
             call(
-                context_name=f'{FEATURE_GROUP_NAME}-{FEATURE_GROUP["CreationTime"].strftime("%s")}'
+                context_name=f"{FEATURE_GROUP_NAME}-{creation_time_str}"
                 f"-feature-group-pipeline",
                 sagemaker_session=SAGEMAKER_SESSION_MOCK,
             ),
             call(
-                context_name=f'{FEATURE_GROUP_NAME}-{FEATURE_GROUP["CreationTime"].strftime("%s")}'
+                context_name=f"{FEATURE_GROUP_NAME}-{creation_time_str}"
                 f"-feature-group-pipeline-version",
                 sagemaker_session=SAGEMAKER_SESSION_MOCK,
             ),

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/test_feature_scheduler.py

@@ -29,6 +29,7 @@
 )
 from sagemaker.core.lineage.context import Context
 from sagemaker.core.remote_function.spark_config import SparkConfig
+from sagemaker.core.resources import FeatureGroup
 
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.mlops.feature_store.feature_processor._enums import FeatureProcessorMode
@@ -89,10 +90,13 @@
 PIPELINE_VERSION_CONTEXT_NAME_TAG_KEY = "sm-fs-fe:feature-engineering-pipeline-version-context-name"
 NOW = datetime.now()
 SAGEMAKER_SESSION_MOCK = Mock(Session)
+SAGEMAKER_SESSION_MOCK.boto_session = Mock()
 CONTEXT_MOCK_01 = Mock(Context)
 CONTEXT_MOCK_02 = Mock(Context)
 CONTEXT_MOCK_03 = Mock(Context)
 FEATURE_GROUP = tdh.DESCRIBE_FEATURE_GROUP_RESPONSE.copy()
+FEATURE_GROUP_MOCK = Mock()
+FEATURE_GROUP_MOCK.creation_time = FEATURE_GROUP["CreationTime"]
 PIPELINE = tdh.PIPELINE.copy()
 TAGS = [dict(Key="key_1", Value="value_1"), dict(Key="key_2", Value="value_2")]
 
@@ -783,8 +787,8 @@ def test_execute(validation):
 
 def test_validate_fg_lineage_resources_happy_case():
     with patch.object(
-        SAGEMAKER_SESSION_MOCK, "describe_feature_group", return_value=FEATURE_GROUP
-    ) as fg_describe_method:
+        FeatureGroup, "get", return_value=FEATURE_GROUP_MOCK
+    ) as fg_get_method:
         with patch.object(
             Context, "load", side_effect=[CONTEXT_MOCK_01, CONTEXT_MOCK_02, CONTEXT_MOCK_03]
         ) as context_load:
@@ -795,16 +799,17 @@ def test_validate_fg_lineage_resources_happy_case():
                 feature_group_name="some_fg",
                 sagemaker_session=SAGEMAKER_SESSION_MOCK,
             )
-    fg_describe_method.assert_called_once_with(feature_group_name="some_fg")
+    fg_get_method.assert_called_once_with(feature_group_name="some_fg", session=SAGEMAKER_SESSION_MOCK.boto_session)
+    creation_time_str = FEATURE_GROUP_MOCK.creation_time.strftime("%s")
     context_load.assert_has_calls(
         [
             call(
-                context_name=f'{"some_fg"}-{FEATURE_GROUP["CreationTime"].strftime("%s")}'
+                context_name=f'{"some_fg"}-{creation_time_str}'
                 f"-feature-group-pipeline",
                 sagemaker_session=SAGEMAKER_SESSION_MOCK,
             ),
             call(
-                context_name=f'{"some_fg"}-{FEATURE_GROUP["CreationTime"].strftime("%s")}'
+                context_name=f'{"some_fg"}-{creation_time_str}'
                 f"-feature-group-pipeline-version",
                 sagemaker_session=SAGEMAKER_SESSION_MOCK,
             ),
@@ -814,7 +819,7 @@ def test_validate_fg_lineage_resources_happy_case():
 
 
 def test_validete_fg_lineage_resources_rnf():
-    with patch.object(SAGEMAKER_SESSION_MOCK, "describe_feature_group", return_value=FEATURE_GROUP):
+    with patch.object(FeatureGroup, "get", return_value=FEATURE_GROUP_MOCK):
         with patch.object(
             Context,
             "load",
@@ -824,7 +829,7 @@ def test_validete_fg_lineage_resources_rnf():
             ),
         ):
             feature_group_name = "some_fg"
-            feature_group_creation_time = FEATURE_GROUP["CreationTime"].strftime("%s")
+            feature_group_creation_time = FEATURE_GROUP_MOCK.creation_time.strftime("%s")
             context_name = f"{feature_group_name}-{feature_group_creation_time}"
             with pytest.raises(
                 ValueError,