Add feature store tests

Author: adishaa

sagemaker-mlops/src/sagemaker/mlops/feature_store/MIGRATION_GUIDE.md

@@ -3,11 +3,10 @@
 """SageMaker FeatureStore V3 - powered by sagemaker-core."""
 
 # Resources from core
-from sagemaker_core.main.resources import FeatureGroup, FeatureMetadata
-from sagemaker_core.main.resources import FeatureStore
+from sagemaker.core.resources import FeatureGroup, FeatureMetadata
 
 # Shapes from core (Pydantic - no to_dict() needed)
-from sagemaker_core.main.shapes import (
+from sagemaker.core.shapes import (
     DataCatalogConfig,
     FeatureParameter,
     FeatureValue,
@@ -52,7 +51,6 @@
 from sagemaker.mlops.feature_store.feature_utils import (
     as_hive_ddl,
     create_athena_query,
-    create_dataset,
     get_session_from_role,
     ingest_dataframe,
     load_feature_definitions_from_dataframe,
@@ -76,7 +74,6 @@
     # Resources
     "FeatureGroup",
     "FeatureMetadata",
-    "FeatureStore",
     # Shapes
     "DataCatalogConfig",
     "FeatureParameter",
@@ -113,7 +110,6 @@
     # Utility functions
     "as_hive_ddl",
     "create_athena_query",
-    "create_dataset",
     "get_session_from_role",
     "ingest_dataframe",
     "load_feature_definitions_from_dataframe",

sagemaker-mlops/src/sagemaker/mlops/feature_store/__init__.py

@@ -178,7 +178,9 @@ def construct_feature_group_to_be_merged(
         database=catalog_config.database,
         table_name=catalog_config.table_name,
         record_identifier_feature_name=record_id,
-        event_time_identifier_feature=FeatureDefinition(event_time_name, FeatureTypeEnum(event_time_type)),
+        event_time_identifier_feature=FeatureDefinition(
+            feature_name=event_time_name, feature_type=FeatureTypeEnum(event_time_type).value
+        ),
         target_feature_name_in_base=target_feature_name_in_base,
         table_type=TableType.FEATURE_GROUP,
         feature_name_in_target=feature_name_in_target,
@@ -256,6 +258,47 @@ class DatasetBuilder:
     _event_time_ending_timestamp: datetime.datetime = field(default=None, init=False)
     _feature_groups_to_be_merged: List[FeatureGroupToBeMerged] = field(default_factory=list, init=False)
 
+    @classmethod
+    def create(
+        cls,
+        base: Union[FeatureGroup, pd.DataFrame],
+        output_path: str,
+        session: Session,
+        record_identifier_feature_name: str = None,
+        event_time_identifier_feature_name: str = None,
+        included_feature_names: List[str] = None,
+        kms_key_id: str = None,
+    ) -> "DatasetBuilder":
+        """Create a DatasetBuilder for generating a Dataset.
+
+        Args:
+            base: A FeatureGroup or DataFrame to use as the base.
+            output_path: S3 URI for output.
+            session: SageMaker session.
+            record_identifier_feature_name: Required if base is DataFrame.
+            event_time_identifier_feature_name: Required if base is DataFrame.
+            included_feature_names: Features to include in output.
+            kms_key_id: KMS key for encryption.
+
+        Returns:
+            DatasetBuilder instance.
+        """
+        if isinstance(base, pd.DataFrame):
+            if not record_identifier_feature_name or not event_time_identifier_feature_name:
+                raise ValueError(
+                    "record_identifier_feature_name and event_time_identifier_feature_name "
+                    "are required when base is a DataFrame."
+                )
+        return cls(
+            _sagemaker_session=session,
+            _base=base,
+            _output_path=output_path,
+            _record_identifier_feature_name=record_identifier_feature_name,
+            _event_time_identifier_feature_name=event_time_identifier_feature_name,
+            _included_feature_names=included_feature_names,
+            _kms_key_id=kms_key_id,
+        )
+
     def with_feature_group(
         self,
         feature_group: FeatureGroup,

sagemaker-mlops/src/sagemaker/mlops/feature_store/dataset_builder.py

@@ -13,7 +13,6 @@
 from sagemaker.mlops.feature_store import FeatureGroup as CoreFeatureGroup, FeatureGroup
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.core.s3.client import S3Uploader, S3Downloader
-from sagemaker.mlops.feature_store.dataset_builder import DatasetBuilder
 from sagemaker.mlops.feature_store.feature_definition import (
     FeatureDefinition,
     FractionalFeatureDefinition,
@@ -23,7 +22,7 @@
 )
 from sagemaker.mlops.feature_store.ingestion_manager_pandas import IngestionManagerPandas
 
-from sagemaker import utils
+from sagemaker.core.utils import unique_name_from_base
 
 
 logger = logging.getLogger(__name__)
@@ -207,7 +206,7 @@ def upload_dataframe_to_s3(
         Tuple of (s3_folder, temp_table_name).
     """
 
-    temp_id = utils.unique_name_from_base("dataframe-base")
+    temp_id = unique_name_from_base("dataframe-base")
     local_file = f"{temp_id}.csv"
     s3_folder = os.path.join(output_path, temp_id)
 
@@ -460,29 +459,3 @@ def ingest_dataframe(
     manager.run(data_frame=data_frame, wait=wait, timeout=timeout)
     return manager
 
-def create_dataset(
-    base: Union[FeatureGroup, pd.DataFrame],
-    output_path: str,
-    session: Session,
-    record_identifier_feature_name: str = None,
-    event_time_identifier_feature_name: str = None,
-    included_feature_names: Sequence[str] = None,
-    kms_key_id: str = None,
-) -> DatasetBuilder:
-    """Create a DatasetBuilder for generating a Dataset."""
-    if isinstance(base, pd.DataFrame):
-        if not record_identifier_feature_name or not event_time_identifier_feature_name:
-            raise ValueError(
-                "record_identifier_feature_name and event_time_identifier_feature_name "
-                "are required when base is a DataFrame."
-            )
-    return DatasetBuilder(
-        _sagemaker_session=session,
-        _base=base,
-        _output_path=output_path,
-        _record_identifier_feature_name=record_identifier_feature_name,
-        _event_time_identifier_feature_name=event_time_identifier_feature_name,
-        _included_feature_names=included_feature_names,
-        _kms_key_id=kms_key_id,
-    )
-

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py

@@ -6,13 +6,12 @@
 import signal
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
-from multiprocessing.pool import AsyncResult
+from multiprocessing import Pool
 from typing import Any, Dict, Iterable, List, Sequence, Union
 
 import pandas as pd
 from pandas import DataFrame
 from pandas.api.types import is_list_like
-from pathos.multiprocessing import ProcessingPool
 
 from sagemaker.core.resources import FeatureGroup as CoreFeatureGroup
 from sagemaker.core.shapes import FeatureValue
@@ -54,8 +53,8 @@ class IngestionManagerPandas:
     feature_definitions: Dict[str, Dict[Any, Any]]
     max_workers: int = 1
     max_processes: int = 1
-    _async_result: AsyncResult = field(default=None, init=False)
-    _processing_pool: ProcessingPool = field(default=None, init=False)
+    _async_result: Any = field(default=None, init=False)
+    _processing_pool: Pool = field(default=None, init=False)
     _failed_indices: List[int] = field(default_factory=list, init=False)
 
     @property
@@ -100,12 +99,11 @@ def wait(self, timeout: Union[int, float] = None):
             results = self._async_result.get(timeout=timeout)
         except KeyboardInterrupt as e:
             self._processing_pool.terminate()
-            self._processing_pool.close()
-            self._processing_pool.clear()
+            self._processing_pool.join()
             raise e
         else:
             self._processing_pool.close()
-            self._processing_pool.clear()
+            self._processing_pool.join()
 
         self._failed_indices = [idx for failed in results for idx in failed]
 
@@ -170,11 +168,10 @@ def _run_multi_process(
         def init_worker():
             signal.signal(signal.SIGINT, signal.SIG_IGN)
 
-        self._processing_pool = ProcessingPool(self.max_processes, init_worker)
-        self._processing_pool.restart(force=True)
+        self._processing_pool = Pool(self.max_processes, init_worker)
 
-        self._async_result = self._processing_pool.amap(
-            lambda x: IngestionManagerPandas._run_multi_threaded(*x),
+        self._async_result = self._processing_pool.starmap_async(
+            IngestionManagerPandas._run_multi_threaded,
             args,
         )
 

sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0