skipped _run_multi_process() and fixed run issue on wait()

Author: verayu43

src/sagemaker/feature_store/feature_group.py

@@ -240,10 +240,8 @@ def wait(self, timeout=None):
         """
         results = None
         try:
-            self._async_result.wait(timeout=timeout)
-            if self._async_result.ready() and self._async_result._success:
-                results = self._async_result._value
-        except KeyboardInterrupt as i:
+            results = self._async_result.get(timeout=timeout)
+        except KeyboardInterrupt or NotImplementedError as i:
             # terminate workers abruptly on keyboard interrupt.
             self._processing_pool.terminate()
             self._processing_pool.close()
@@ -254,20 +252,20 @@ def wait(self, timeout=None):
             self._processing_pool.close()
             self._processing_pool.clear()
 
-        if not results or results == NotImplementedError:
+        if results == None:
             return
-
-        self._failed_indices = [
-            failed_index for failed_indices in results for failed_index in failed_indices
-        ]
+        else:
+            self._failed_indices = [
+                failed_index for failed_indices in results for failed_index in failed_indices
+            ]
 
         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):
+    def _run_multi_process(self, data_frame: DataFrame, wait=True, timeout=None): # Not in use
         """Start the ingestion process with the specified number of processes.
 
         Args:
@@ -313,7 +311,6 @@ def _run_multi_threaded(self, data_frame: DataFrame, row_offset=0, timeout=None)
         """
         executor = ThreadPoolExecutor(max_workers=self.max_workers)
         batch_size = math.ceil(data_frame.shape[0] / self.max_workers)
-
         futures = {}
         for i in range(self.max_workers):
             start_index = min(i * batch_size, data_frame.shape[0])
@@ -338,7 +335,6 @@ def _run_multi_threaded(self, data_frame: DataFrame, row_offset=0, timeout=None)
             else:
                 logger.info("Successfully ingested row %d to %d", start, end)
             failed_indices += result
-
         executor.shutdown(wait=False)
 
         return failed_indices
@@ -352,7 +348,8 @@ 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)
+        #self._run_multi_process(data_frame=data_frame, wait=wait, timeout=timeout)
+        self._run_multi_threaded(data_frame=data_frame, timeout=timeout)
 
 
 class IngestionError(Exception):