feat: Support custom timeout for MatchingEngineIndex and `Matchin…

…gEngineIndexEndpoint` APIs.

PiperOrigin-RevId: 608815992

google/cloud/aiplatform/matching_engine/matching_engine_index.py

@@ -112,6 +112,7 @@ def _create(
         sync: bool = True,
         index_update_method: Optional[str] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndex":
         """Creates a MatchingEngineIndex resource.
 
@@ -177,6 +178,8 @@ def _create(
                 secured by this key.
                 The key needs to be in the same region as where the index is
                 created.
+            create_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndex - Index resource object
@@ -220,6 +223,7 @@ def _create(
             ),
             index=gapic_index,
             metadata=request_metadata,
+            timeout=create_request_timeout,
         )
 
         _LOGGER.log_create_with_lro(cls, create_lro)
@@ -243,6 +247,7 @@ def update_metadata(
         description: Optional[str] = None,
         labels: Optional[Dict[str, str]] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        update_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndex":
         """Updates the metadata for this index.
 
@@ -269,6 +274,8 @@ def update_metadata(
                 "aiplatform.googleapis.com/" and are immutable.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
+            update_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndex - The updated index resource object.
@@ -307,6 +314,7 @@ def update_metadata(
             index=gapic_index,
             update_mask=update_mask,
             metadata=request_metadata,
+            timeout=update_request_timeout,
         )
 
         _LOGGER.log_action_started_against_resource_with_lro(
@@ -324,6 +332,7 @@ def update_embeddings(
         contents_delta_uri: str,
         is_complete_overwrite: Optional[bool] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        update_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndex":
         """Updates the embeddings for this index.
 
@@ -341,6 +350,8 @@ def update_embeddings(
                 then existing content of the Index will be replaced by the data from the contentsDeltaUri.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
+            update_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndex - The updated index resource object.
@@ -373,6 +384,7 @@ def update_embeddings(
             index=gapic_index,
             update_mask=update_mask,
             metadata=request_metadata,
+            timeout=update_request_timeout,
         )
 
         _LOGGER.log_action_started_against_resource_with_lro(
@@ -420,6 +432,7 @@ def create_tree_ah_index(
         sync: bool = True,
         index_update_method: Optional[str] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndex":
         """Creates a MatchingEngineIndex resource that uses the tree-AH algorithm.
 
@@ -510,6 +523,8 @@ def create_tree_ah_index(
                 secured by this key.
                 The key needs to be in the same region as where the index is
                 created.
+            create_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndex - Index resource object
@@ -541,6 +556,7 @@ def create_tree_ah_index(
             sync=sync,
             index_update_method=index_update_method,
             encryption_spec_key_name=encryption_spec_key_name,
+            create_request_timeout=create_request_timeout,
         )
 
     @classmethod
@@ -561,6 +577,7 @@ def create_brute_force_index(
         sync: bool = True,
         index_update_method: Optional[str] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndex":
         """Creates a MatchingEngineIndex resource that uses the brute force algorithm.
 
@@ -640,6 +657,8 @@ def create_brute_force_index(
                 secured by this key.
                 The key needs to be in the same region as where the index is
                 created.
+            create_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndex - Index resource object
@@ -667,6 +686,7 @@ def create_brute_force_index(
             sync=sync,
             index_update_method=index_update_method,
             encryption_spec_key_name=encryption_spec_key_name,
+            create_request_timeout=create_request_timeout,
         )
 
     def upsert_datapoints(

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -346,6 +346,7 @@ def create(
         enable_private_service_connect: bool = False,
         project_allowlist: Optional[Sequence[str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Creates a MatchingEngineIndexEndpoint resource.
 
@@ -423,6 +424,8 @@ def create(
                 index endpoint will be secured by this key.
                 The key needs to be in the same region as where the index
                 endpoint is created.
+            create_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
@@ -469,6 +472,7 @@ def create(
             enable_private_service_connect=enable_private_service_connect,
             project_allowlist=project_allowlist,
             encryption_spec_key_name=encryption_spec_key_name,
+            create_request_timeout=create_request_timeout,
         )
 
     @classmethod
@@ -488,6 +492,7 @@ def _create(
         enable_private_service_connect: bool = False,
         project_allowlist: Optional[Sequence[str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Helper method to ensure network synchronization and to
         create a MatchingEngineIndexEndpoint resource.
@@ -552,6 +557,8 @@ def _create(
             project_allowlist (MutableSequence[str]):
                 A list of Projects from which the forwarding
                 rule will target the service attachment.
+            create_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
@@ -596,6 +603,7 @@ def _create(
             ),
             index_endpoint=gapic_index_endpoint,
             metadata=request_metadata,
+            timeout=create_request_timeout,
         )
 
         _LOGGER.log_create_with_lro(cls, create_lro)
@@ -716,6 +724,7 @@ def update(
         description: Optional[str] = None,
         labels: Optional[Dict[str, str]] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        update_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Updates an existing index endpoint resource.
 
@@ -742,6 +751,8 @@ def update(
                 "aiplatform.googleapis.com/" and are immutable.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
+            update_request_timeout (float):
+                Optional. The timeout for the request in seconds.
 
         Returns:
             MatchingEngineIndexEndpoint - The updated index endpoint resource object.
@@ -774,6 +785,7 @@ def update(
             index_endpoint=gapic_index_endpoint,
             update_mask=update_mask,
             metadata=request_metadata,
+            timeout=update_request_timeout,
         )
 
         return self
@@ -937,6 +949,7 @@ def deploy_index(
         auth_config_audiences: Optional[Sequence[str]] = None,
         auth_config_allowed_issuers: Optional[Sequence[str]] = None,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        deploy_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Deploys an existing index resource to this endpoint resource.
 
@@ -1030,6 +1043,9 @@ def deploy_index(
                 auth_config_audiences and auth_config_allowed_issuers must be passed together.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
+
+            deploy_request_timeout (float):
+                Optional. The timeout for the request in seconds.
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
         """
@@ -1060,6 +1076,7 @@ def deploy_index(
             index_endpoint=self.resource_name,
             deployed_index=deployed_index,
             metadata=request_metadata,
+            timeout=deploy_request_timeout,
         )
 
         _LOGGER.log_action_started_against_resource_with_lro(
@@ -1081,6 +1098,7 @@ def undeploy_index(
         self,
         deployed_index_id: str,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        undeploy_request_timeout: Optional[float] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Undeploy a deployed index endpoint resource.
 
@@ -1090,6 +1108,8 @@ def undeploy_index(
                 to be undeployed from the IndexEndpoint.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
+            undeploy_request_timeout (float):
+                Optional. The timeout for the request in seconds.
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
         """
@@ -1106,6 +1126,7 @@ def undeploy_index(
             index_endpoint=self.resource_name,
             deployed_index_id=deployed_index_id,
             metadata=request_metadata,
+            timeout=undeploy_request_timeout,
         )
 
         _LOGGER.log_action_started_against_resource_with_lro(
@@ -1126,6 +1147,7 @@ def mutate_deployed_index(
         min_replica_count: int = 1,
         max_replica_count: int = 1,
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        mutate_request_timeout: Optional[float] = None,
     ):
         """Updates an existing deployed index under this endpoint resource.
 
@@ -1157,6 +1179,8 @@ def mutate_deployed_index(
                 will automatically be increased to be min_replica_count.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
+            timeout (float):
+                Optional. The timeout for the request in seconds.
         """
 
         self.wait()
@@ -1178,6 +1202,7 @@ def mutate_deployed_index(
             index_endpoint=self.resource_name,
             deployed_index=deployed_index,
             metadata=request_metadata,
+            timeout=mutate_request_timeout,
         )
 
         _LOGGER.log_action_started_against_resource_with_lro(
@@ -1211,6 +1236,7 @@ def _undeploy(
         deployed_index_id: str,
         metadata: Optional[Sequence[Tuple[str, str]]] = (),
         sync=True,
+        undeploy_request_timeout: Optional[float] = None,
     ) -> None:
         """Undeploys a deployed index.
 
@@ -1221,6 +1247,8 @@ def _undeploy(
             metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as
                 metadata.
+            timeout (float):
+                Optional. The timeout for the request in seconds.
         """
         self._sync_gca_resource()
 
@@ -1230,6 +1258,7 @@ def _undeploy(
             index_endpoint=self.resource_name,
             deployed_index_id=deployed_index_id,
             metadata=metadata,
+            timeout=undeploy_request_timeout,
         )
 
         _LOGGER.log_action_started_against_resource_with_lro(