feat: add timeout arg across SDK

google/cloud/aiplatform/datasets/dataset.py

@@ -118,6 +118,7 @@ def create(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "_Dataset":
         """Creates a new dataset and optionally imports data into dataset when
@@ -199,6 +200,10 @@ def create(
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
 
                 Overrides encryption_spec_key_name set in aiplatform.init.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -239,6 +244,7 @@ def create(
             encryption_spec=initializer.global_config.get_encryption_spec(
                 encryption_spec_key_name=encryption_spec_key_name
             ),
+            create_request_timeout=create_request_timeout,
             sync=sync,
         )
 
@@ -257,6 +263,7 @@ def _create_and_import(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None,
+        create_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "_Dataset":
         """Creates a new dataset and optionally imports data into dataset when
@@ -309,6 +316,10 @@ def _create_and_import(
                 resource is created.
 
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -328,6 +339,7 @@ def _create_and_import(
             request_metadata=request_metadata,
             labels=labels,
             encryption_spec=encryption_spec,
+            create_request_timeout=create_request_timeout,
         )
 
         _LOGGER.log_create_with_lro(cls, create_dataset_lro)
@@ -344,17 +356,20 @@ def _create_and_import(
         )
 
         # Import if import datasource is DatasourceImportable
+        # import_request_timeout is None since user is issuing a single request with create and import
         if isinstance(datasource, _datasources.DatasourceImportable):
-            dataset_obj._import_and_wait(datasource)
+            dataset_obj._import_and_wait(datasource, import_request_timeout=None)
 
         return dataset_obj
 
-    def _import_and_wait(self, datasource):
+    def _import_and_wait(self, datasource, import_request_timeout):
         _LOGGER.log_action_start_against_resource(
             "Importing", "data", self,
         )
 
-        import_lro = self._import(datasource=datasource)
+        import_lro = self._import(
+            datasource=datasource, import_request_timeout=import_request_timeout
+        )
 
         _LOGGER.log_action_started_against_resource_with_lro(
             "Import", "data", self.__class__, import_lro
@@ -375,6 +390,7 @@ def _create(
         request_metadata: Sequence[Tuple[str, str]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None,
+        create_request_timeout: Optional[float] = None,
     ) -> operation.Operation:
         """Creates a new managed dataset by directly calling API client.
 
@@ -417,6 +433,10 @@ def _create(
                 resource is created.
 
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
         Returns:
             operation (Operation):
                 An object representing a long-running operation.
@@ -431,11 +451,16 @@ def _create(
         )
 
         return api_client.create_dataset(
-            parent=parent, dataset=gapic_dataset, metadata=request_metadata
+            parent=parent,
+            dataset=gapic_dataset,
+            metadata=request_metadata,
+            timeout=create_request_timeout,
         )
 
     def _import(
-        self, datasource: _datasources.DatasourceImportable,
+        self,
+        datasource: _datasources.DatasourceImportable,
+        import_request_timeout: Optional[float] = None,
     ) -> operation.Operation:
         """Imports data into managed dataset by directly calling API client.
 
@@ -448,7 +473,9 @@ def _import(
                 An object representing a long-running operation.
         """
         return self.api_client.import_data(
-            name=self.resource_name, import_configs=[datasource.import_data_config]
+            name=self.resource_name,
+            import_configs=[datasource.import_data_config],
+            timeout=import_request_timeout,
         )
 
     @base.optional_sync(return_input_arg="self")
@@ -457,6 +484,7 @@ def import_data(
         gcs_source: Union[str, Sequence[str]],
         import_schema_uri: str,
         data_item_labels: Optional[Dict] = None,
+        import_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "_Dataset":
         """Upload data to existing managed dataset.
@@ -491,6 +519,10 @@ def import_data(
                 labels specified inside index file referenced by
                 ``import_schema_uri``,
                 e.g. jsonl file.
+            import_request_timeout (float):
+                Optional. The timeout for initiating this import request in seconds. Note:
+                this does not set the timeout on the underlying import job, only on the time
+                to initiate the import request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -507,7 +539,9 @@ def import_data(
             data_item_labels=data_item_labels,
         )
 
-        self._import_and_wait(datasource=datasource)
+        self._import_and_wait(
+            datasource=datasource, import_request_timeout=import_request_timeout
+        )
         return self
 
     # TODO(b/174751568) add optional sync support

google/cloud/aiplatform/datasets/image_dataset.py

@@ -46,6 +46,7 @@ def create(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "ImageDataset":
         """Creates a new image dataset and optionally imports data into dataset
@@ -117,6 +118,10 @@ def create(
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
 
                 Overrides encryption_spec_key_name set in aiplatform.init.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -158,5 +163,6 @@ def create(
             encryption_spec=initializer.global_config.get_encryption_spec(
                 encryption_spec_key_name=encryption_spec_key_name
             ),
+            create_request_timeout=create_request_timeout,
             sync=sync,
         )

google/cloud/aiplatform/datasets/tabular_dataset.py

@@ -45,6 +45,7 @@ def create(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "TabularDataset":
         """Creates a new tabular dataset.
@@ -98,6 +99,10 @@ def create(
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
 
                 Overrides encryption_spec_key_name set in aiplatform.init.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -138,6 +143,7 @@ def create(
             encryption_spec=initializer.global_config.get_encryption_spec(
                 encryption_spec_key_name=encryption_spec_key_name
             ),
+            create_request_timeout=create_request_timeout,
             sync=sync,
         )
 

google/cloud/aiplatform/datasets/text_dataset.py

@@ -46,6 +46,7 @@ def create(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "TextDataset":
         """Creates a new text dataset and optionally imports data into dataset
@@ -124,6 +125,10 @@ def create(
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
 
                 Overrides encryption_spec_key_name set in aiplatform.init.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -165,5 +170,6 @@ def create(
             encryption_spec=initializer.global_config.get_encryption_spec(
                 encryption_spec_key_name=encryption_spec_key_name
             ),
+            create_request_timeout=create_request_timeout,
             sync=sync,
         )

google/cloud/aiplatform/datasets/video_dataset.py

@@ -46,6 +46,7 @@ def create(
         request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
         labels: Optional[Dict[str, str]] = None,
         encryption_spec_key_name: Optional[str] = None,
+        create_request_timeout: Optional[float] = None,
         sync: bool = True,
     ) -> "VideoDataset":
         """Creates a new video dataset and optionally imports data into dataset
@@ -117,6 +118,10 @@ def create(
                 If set, this Dataset and all sub-resources of this Dataset will be secured by this key.
 
                 Overrides encryption_spec_key_name set in aiplatform.init.
+            create_request_timeout (float):
+                Optional. The timeout for initiating this create request in seconds. Note:
+                this does not set the timeout on the underlying create job, only on the time
+                to initiate the create request.
             sync (bool):
                 Whether to execute this method synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
@@ -158,5 +163,6 @@ def create(
             encryption_spec=initializer.global_config.get_encryption_spec(
                 encryption_spec_key_name=encryption_spec_key_name
             ),
+            create_request_timeout=create_request_timeout,
             sync=sync,
         )