fix: Unbreak timeouts for Dataset create, FeatureStore ingest, and Ma…

…tchingEngine Index create. PiperOrigin-RevId: 501562819
googleapis · Jan 12, 2023 · 3096d1c · 3096d1c
1 parent fe75eba
commit 3096d1c
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/google/cloud/aiplatform/datasets/dataset.py b/google/cloud/aiplatform/datasets/dataset.py
@@ -389,7 +389,7 @@ def _import_and_wait(
             "Import", "data", self.__class__, import_lro
         )
 
-        import_lro.result()
+        import_lro.result(timeout=None)
 
         _LOGGER.log_action_completed_against_resource("data", "imported", self)
 

diff --git a/google/cloud/aiplatform/featurestore/_entity_type.py b/google/cloud/aiplatform/featurestore/_entity_type.py
@@ -961,7 +961,7 @@ def _import_feature_values(
             "Import", "feature values", self.__class__, import_lro
         )
 
-        import_lro.result()
+        import_lro.result(timeout=None)
 
         _LOGGER.log_action_completed_against_resource(
             "feature values", "imported", self
@@ -1478,14 +1478,13 @@ def read(
                     feature_selector=feature_selector,
                 )
             )
-            streaming_read_feature_values_responses = [
-                response
-                for response in self._featurestore_online_client.streaming_read_feature_values(
+            streaming_read_feature_values_responses = list(
+                self._featurestore_online_client.streaming_read_feature_values(
                     request=streaming_read_feature_values_request,
                     metadata=request_metadata,
                     timeout=read_request_timeout,
                 )
-            ]
+            )
             header = streaming_read_feature_values_responses[0].header
             entity_views = [
                 response.entity_view

diff --git a/google/cloud/aiplatform/matching_engine/matching_engine_index.py b/google/cloud/aiplatform/matching_engine/matching_engine_index.py
@@ -25,7 +25,9 @@
     matching_engine_index as gca_matching_engine_index,
 )
 from google.cloud.aiplatform import initializer
-from google.cloud.aiplatform.matching_engine import matching_engine_index_config
+from google.cloud.aiplatform.matching_engine import (
+    matching_engine_index_config,
+)
 from google.cloud.aiplatform import utils
 
 _LOGGER = base.Logger(__name__)
@@ -190,7 +192,7 @@ def _create(
 
         _LOGGER.log_create_with_lro(cls, create_lro)
 
-        created_index = create_lro.result()
+        created_index = create_lro.result(timeout=None)
 
         _LOGGER.log_create_complete(cls, created_index, "index")