fix wait function to handle more excpections

Author: verayu43

src/sagemaker/feature_store/feature_group.py

@@ -238,8 +238,11 @@ def wait(self, timeout=None):
             timeout (Union[int, float]): ``concurrent.futures.TimeoutError`` will be raised
                 if timeout is reached.
         """
+        results = None
         try:
-            results = self._async_result.get(timeout=timeout)
+            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:
             # terminate workers abruptly on keyboard interrupt.
             self._processing_pool.terminate()
@@ -251,6 +254,9 @@ def wait(self, timeout=None):
             self._processing_pool.close()
             self._processing_pool.clear()
 
+        if not results or results == NotImplementedError:
+            return
+
         self._failed_indices = [
             failed_index for failed_indices in results for failed_index in failed_indices
         ]