fix: use FeatureGroup's Session in nonconcurrency ingestion (#3617)

Author: ericpark3

src/sagemaker/feature_store/feature_group.py

@@ -165,6 +165,7 @@ class IngestionManagerPandas:
         feature_group_name (str): name of the Feature Group.
         sagemaker_fs_runtime_client_config (Config): instance of the Config class
             for boto calls.
+        sagemaker_session (Session): session instance to perform boto calls.
         data_frame (DataFrame): pandas DataFrame to be ingested to the given feature group.
         max_workers (int): number of threads to create.
         max_processes (int): number of processes to create. Each process spawns
@@ -174,7 +175,8 @@ class IngestionManagerPandas:
     """
 
     feature_group_name: str = attr.ib()
-    sagemaker_fs_runtime_client_config: Config = attr.ib()
+    sagemaker_fs_runtime_client_config: Config = attr.ib(default=None)
+    sagemaker_session: Session = attr.ib(default=None)
     max_workers: int = attr.ib(default=1)
     max_processes: int = attr.ib(default=1)
     profile_name: str = attr.ib(default=None)
@@ -210,29 +212,20 @@ def _ingest_single_batch(
         if "max_attempts" not in retry_config and "total_max_attempts" not in retry_config:
             client_config = copy.deepcopy(client_config)
             client_config.retries = {"max_attempts": 10, "mode": "standard"}
-        sagemaker_featurestore_runtime_client = boto3.Session(profile_name=profile_name).client(
+        sagemaker_fs_runtime_client = boto3.Session(profile_name=profile_name).client(
             service_name="sagemaker-featurestore-runtime", config=client_config
         )
 
         logger.info("Started ingesting index %d to %d", start_index, end_index)
         failed_rows = list()
         for row in data_frame[start_index:end_index].itertuples():
-            record = [
-                FeatureValue(
-                    feature_name=data_frame.columns[index - 1],
-                    value_as_string=str(row[index]),
-                )
-                for index in range(1, len(row))
-                if pd.notna(row[index])
-            ]
-            try:
-                sagemaker_featurestore_runtime_client.put_record(
-                    FeatureGroupName=feature_group_name,
-                    Record=[value.to_dict() for value in record],
-                )
-            except Exception as e:  # pylint: disable=broad-except
-                logger.error("Failed to ingest row %d: %s", row[0], e)
-                failed_rows.append(row[0])
+            IngestionManagerPandas._ingest_row(
+                data_frame=data_frame,
+                row=row,
+                feature_group_name=feature_group_name,
+                sagemaker_fs_runtime_client=sagemaker_fs_runtime_client,
+                failed_rows=failed_rows,
+            )
         return failed_rows
 
     @property
@@ -274,6 +267,69 @@ def wait(self, timeout=None):
                 f"Failed to ingest some data into FeatureGroup {self.feature_group_name}",
             )
 
+    @staticmethod
+    def _ingest_row(
+        data_frame: DataFrame,
+        row: int,
+        feature_group_name: str,
+        sagemaker_fs_runtime_client: Session,
+        failed_rows: List[int],
+    ):
+        """Ingest a single Dataframe row into FeatureStore.
+
+        Args:
+            data_frame (DataFrame): source DataFrame to be ingested.
+            row (int): current row that is being ingested
+            feature_group_name (str): name of the Feature Group.
+            sagemaker_featurestore_runtime_client (Session): session instance to perform boto calls.
+            failed_rows (List[int]): list of indices from the data frame for which ingestion failed.
+
+
+        Returns:
+            int of row indices that failed to be ingested.
+        """
+        record = [
+            FeatureValue(
+                feature_name=data_frame.columns[index - 1],
+                value_as_string=str(row[index]),
+            )
+            for index in range(1, len(row))
+            if pd.notna(row[index])
+        ]
+        try:
+            sagemaker_fs_runtime_client.put_record(
+                FeatureGroupName=feature_group_name,
+                Record=[value.to_dict() for value in record],
+            )
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error("Failed to ingest row %d: %s", row[0], e)
+            failed_rows.append(row[0])
+
+    def _run_single_process_single_thread(self, data_frame: DataFrame):
+        """Ingest a utilizing single process and single thread.
+
+        Args:
+            data_frame (DataFrame): source DataFrame to be ingested.
+        """
+        logger.info("Started ingesting index %d to %d")
+        failed_rows = list()
+        sagemaker_fs_runtime_client = self.sagemaker_session.sagemaker_featurestore_runtime_client
+        for row in data_frame.itertuples():
+            IngestionManagerPandas._ingest_row(
+                data_frame=data_frame,
+                row=row,
+                feature_group_name=self.feature_group_name,
+                sagemaker_fs_runtime_client=sagemaker_fs_runtime_client,
+                failed_rows=failed_rows,
+            )
+        self._failed_indices = failed_rows
+
+        if len(self._failed_indices) > 0:
+            raise IngestionError(
+                self._failed_indices,
+                f"Failed to ingest some data into FeatureGroup {self.feature_group_name}",
+            )
+
     def _run_multi_process(self, data_frame: DataFrame, wait=True, timeout=None):
         """Start the ingestion process with the specified number of processes.
 
@@ -385,7 +441,10 @@ def run(self, data_frame: DataFrame, wait=True, timeout=None):
             timeout (Union[int, float]): ``concurrent.futures.TimeoutError`` will be raised
                 if timeout is reached.
         """
-        self._run_multi_process(data_frame=data_frame, wait=wait, timeout=timeout)
+        if self.max_workers == 1 and self.max_processes == 1 and self.profile_name is None:
+            self._run_single_process_single_thread(data_frame=data_frame)
+        else:
+            self._run_multi_process(data_frame=data_frame, wait=wait, timeout=timeout)
 
 
 class IngestionError(Exception):
@@ -749,6 +808,7 @@ def ingest(
 
         manager = IngestionManagerPandas(
             feature_group_name=self.name,
+            sagemaker_session=self.sagemaker_session,
             sagemaker_fs_runtime_client_config=self.sagemaker_session.sagemaker_featurestore_runtime_client.meta.config,
             max_workers=max_workers,
             max_processes=max_processes,

tests/unit/sagemaker/feature_store/test_feature_group.py

@@ -307,6 +307,7 @@ def test_ingest(ingestion_manager_init, sagemaker_session_mock, fs_runtime_clien
 
     ingestion_manager_init.assert_called_once_with(
         feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
         sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
         max_workers=10,
         max_processes=1,
@@ -317,6 +318,32 @@ def test_ingest(ingestion_manager_init, sagemaker_session_mock, fs_runtime_clien
     )
 
 
+@patch("sagemaker.feature_store.feature_group.IngestionManagerPandas")
+def test_ingest_default(ingestion_manager_init, sagemaker_session_mock):
+    sagemaker_session_mock.sagemaker_featurestore_runtime_client.meta.config = (
+        fs_runtime_client_config_mock
+    )
+
+    feature_group = FeatureGroup(name="MyGroup", sagemaker_session=sagemaker_session_mock)
+    df = pd.DataFrame(dict((f"float{i}", pd.Series([2.0], dtype="float64")) for i in range(300)))
+
+    mock_ingestion_manager_instance = Mock()
+    ingestion_manager_init.return_value = mock_ingestion_manager_instance
+    feature_group.ingest(data_frame=df)
+
+    ingestion_manager_init.assert_called_once_with(
+        feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
+        sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
+        max_workers=1,
+        max_processes=1,
+        profile_name=None,
+    )
+    mock_ingestion_manager_instance.run.assert_called_once_with(
+        data_frame=df, wait=True, timeout=None
+    )
+
+
 @patch("sagemaker.feature_store.feature_group.IngestionManagerPandas")
 def test_ingest_with_profile_name(
     ingestion_manager_init, sagemaker_session_mock, fs_runtime_client_config_mock
@@ -334,6 +361,7 @@ def test_ingest_with_profile_name(
 
     ingestion_manager_init.assert_called_once_with(
         feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
         sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
         max_workers=10,
         max_processes=1,
@@ -403,6 +431,7 @@ def test_ingestion_manager_run_success():
     df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
     manager = IngestionManagerPandas(
         feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
         sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
         max_workers=10,
     )
@@ -421,6 +450,7 @@ def test_ingestion_manager_run_multi_process_with_multi_thread_success(
     df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
     manager = IngestionManagerPandas(
         feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
         sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
         max_workers=2,
         max_processes=2,
@@ -436,16 +466,17 @@ def test_ingestion_manager_run_failure():
     df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
     manager = IngestionManagerPandas(
         feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
         sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
-        max_workers=1,
+        max_workers=2,
     )
 
     with pytest.raises(IngestionError) as error:
         manager.run(df)
 
     assert "Failed to ingest some data into FeatureGroup MyGroup" in str(error)
-    assert error.value.failed_rows == [1]
-    assert manager.failed_rows == [1]
+    assert error.value.failed_rows == [1, 1]
+    assert manager.failed_rows == [1, 1]
 
 
 @patch(
@@ -456,6 +487,7 @@ def test_ingestion_manager_with_profile_name_run_failure():
     df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
     manager = IngestionManagerPandas(
         feature_group_name="MyGroup",
+        sagemaker_session=sagemaker_session_mock,
         sagemaker_fs_runtime_client_config=fs_runtime_client_config_mock,
         max_workers=1,
         profile_name="non_exist",
@@ -475,6 +507,7 @@ def test_ingestion_manager_run_multi_process_failure():
     df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
     manager = IngestionManagerPandas(
         feature_group_name="MyGroup",
+        sagemaker_session=None,
         sagemaker_fs_runtime_client_config=None,
         max_workers=2,
         max_processes=2,