From 3096d1c72de7c280e56d8d58192230ba36ccfce3 Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Thu, 12 Jan 2023 07:42:36 -0800 Subject: [PATCH] fix: Unbreak timeouts for Dataset create, FeatureStore ingest, and MatchingEngine Index create. PiperOrigin-RevId: 501562819 --- google/cloud/aiplatform/datasets/dataset.py | 2 +- google/cloud/aiplatform/featurestore/_entity_type.py | 9 ++++----- .../aiplatform/matching_engine/matching_engine_index.py | 6 ++++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/google/cloud/aiplatform/datasets/dataset.py b/google/cloud/aiplatform/datasets/dataset.py index a44d21d784..a6741a0cc6 100644 --- 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 index cec9598c04..0e6c22a4fb 100644 --- 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 index d382a126f1..9b594699d8 100644 --- 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")