Skip to content

Commit 3096d1c

Browse files
sasha-gitgcopybara-github
authored andcommitted
fix: Unbreak timeouts for Dataset create, FeatureStore ingest, and MatchingEngine Index create.
PiperOrigin-RevId: 501562819
1 parent fe75eba commit 3096d1c

File tree

3 files changed

+9
-8
lines changed

3 files changed

+9
-8
lines changed

google/cloud/aiplatform/datasets/dataset.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,7 @@ def _import_and_wait(
389389
"Import", "data", self.__class__, import_lro
390390
)
391391

392-
import_lro.result()
392+
import_lro.result(timeout=None)
393393

394394
_LOGGER.log_action_completed_against_resource("data", "imported", self)
395395

google/cloud/aiplatform/featurestore/_entity_type.py

+4-5
Original file line numberDiff line numberDiff line change
@@ -961,7 +961,7 @@ def _import_feature_values(
961961
"Import", "feature values", self.__class__, import_lro
962962
)
963963

964-
import_lro.result()
964+
import_lro.result(timeout=None)
965965

966966
_LOGGER.log_action_completed_against_resource(
967967
"feature values", "imported", self
@@ -1478,14 +1478,13 @@ def read(
14781478
feature_selector=feature_selector,
14791479
)
14801480
)
1481-
streaming_read_feature_values_responses = [
1482-
response
1483-
for response in self._featurestore_online_client.streaming_read_feature_values(
1481+
streaming_read_feature_values_responses = list(
1482+
self._featurestore_online_client.streaming_read_feature_values(
14841483
request=streaming_read_feature_values_request,
14851484
metadata=request_metadata,
14861485
timeout=read_request_timeout,
14871486
)
1488-
]
1487+
)
14891488
header = streaming_read_feature_values_responses[0].header
14901489
entity_views = [
14911490
response.entity_view

google/cloud/aiplatform/matching_engine/matching_engine_index.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@
2525
matching_engine_index as gca_matching_engine_index,
2626
)
2727
from google.cloud.aiplatform import initializer
28-
from google.cloud.aiplatform.matching_engine import matching_engine_index_config
28+
from google.cloud.aiplatform.matching_engine import (
29+
matching_engine_index_config,
30+
)
2931
from google.cloud.aiplatform import utils
3032

3133
_LOGGER = base.Logger(__name__)
@@ -190,7 +192,7 @@ def _create(
190192

191193
_LOGGER.log_create_with_lro(cls, create_lro)
192194

193-
created_index = create_lro.result()
195+
created_index = create_lro.result(timeout=None)
194196

195197
_LOGGER.log_create_complete(cls, created_index, "index")
196198

0 commit comments

Comments
 (0)