diff --git a/google/cloud/aiplatform/datasets/dataset.py b/google/cloud/aiplatform/datasets/dataset.py index 48398bf005..757af1922d 100644 --- a/google/cloud/aiplatform/datasets/dataset.py +++ b/google/cloud/aiplatform/datasets/dataset.py @@ -119,6 +119,7 @@ def create( labels: Optional[Dict[str, str]] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "_Dataset": """Creates a new dataset and optionally imports data into dataset when source and import_schema_uri are passed. @@ -203,6 +204,8 @@ def create( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: dataset (Dataset): @@ -240,6 +243,7 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) @classmethod @@ -258,6 +262,8 @@ def _create_and_import( labels: Optional[Dict[str, str]] = None, encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, + import_request_timeout: Optional[float] = None, ) -> "_Dataset": """Creates a new dataset and optionally imports data into dataset when source and import_schema_uri are passed. @@ -313,6 +319,10 @@ def _create_and_import( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. + import_request_timeout (float): + Optional. The timeout for the import request in seconds. Returns: dataset (Dataset): @@ -328,6 +338,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) @@ -345,18 +356,26 @@ def _create_and_import( # Import if import datasource is DatasourceImportable if isinstance(datasource, _datasources.DatasourceImportable): - dataset_obj._import_and_wait(datasource) + dataset_obj._import_and_wait( + datasource, import_request_timeout=import_request_timeout + ) return dataset_obj - def _import_and_wait(self, datasource): + def _import_and_wait( + self, + datasource, + import_request_timeout: Optional[float] = None, + ): _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 @@ -377,6 +396,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. @@ -419,6 +439,8 @@ 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 the create request in seconds. Returns: operation (Operation): An object representing a long-running operation. @@ -433,25 +455,33 @@ 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, + import_request_timeout: Optional[float] = None, ) -> operation.Operation: """Imports data into managed dataset by directly calling API client. Args: datasource (_datasources.DatasourceImportable): Required. Datasource for importing data to an existing dataset for Vertex AI. + import_request_timeout (float): + Optional. The timeout for the import request in seconds. Returns: operation (Operation): 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") @@ -461,6 +491,7 @@ def import_data( import_schema_uri: str, data_item_labels: Optional[Dict] = None, sync: bool = True, + import_request_timeout: Optional[float] = None, ) -> "_Dataset": """Upload data to existing managed dataset. @@ -498,6 +529,8 @@ def import_data( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + import_request_timeout (float): + Optional. The timeout for the import request in seconds. Returns: dataset (Dataset): @@ -510,7 +543,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 diff --git a/google/cloud/aiplatform/datasets/image_dataset.py b/google/cloud/aiplatform/datasets/image_dataset.py index 14b26cc6de..98571f1d69 100644 --- a/google/cloud/aiplatform/datasets/image_dataset.py +++ b/google/cloud/aiplatform/datasets/image_dataset.py @@ -47,6 +47,7 @@ def create( labels: Optional[Dict[str, str]] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "ImageDataset": """Creates a new image dataset and optionally imports data into dataset when source and import_schema_uri are passed. @@ -121,6 +122,8 @@ def create( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: image_dataset (ImageDataset): @@ -159,4 +162,5 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) diff --git a/google/cloud/aiplatform/datasets/tabular_dataset.py b/google/cloud/aiplatform/datasets/tabular_dataset.py index 57ad827b31..62ddf43d9d 100644 --- a/google/cloud/aiplatform/datasets/tabular_dataset.py +++ b/google/cloud/aiplatform/datasets/tabular_dataset.py @@ -46,6 +46,7 @@ def create( labels: Optional[Dict[str, str]] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "TabularDataset": """Creates a new tabular dataset. @@ -102,6 +103,8 @@ def create( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: tabular_dataset (TabularDataset): @@ -139,6 +142,7 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) def import_data(self): diff --git a/google/cloud/aiplatform/datasets/text_dataset.py b/google/cloud/aiplatform/datasets/text_dataset.py index 8fa28c3f31..7126c4e53a 100644 --- a/google/cloud/aiplatform/datasets/text_dataset.py +++ b/google/cloud/aiplatform/datasets/text_dataset.py @@ -47,6 +47,7 @@ def create( labels: Optional[Dict[str, str]] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "TextDataset": """Creates a new text dataset and optionally imports data into dataset when source and import_schema_uri are passed. @@ -124,6 +125,8 @@ 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 the create request in seconds. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -166,4 +169,5 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) diff --git a/google/cloud/aiplatform/datasets/video_dataset.py b/google/cloud/aiplatform/datasets/video_dataset.py index a758339b71..f98a70b596 100644 --- a/google/cloud/aiplatform/datasets/video_dataset.py +++ b/google/cloud/aiplatform/datasets/video_dataset.py @@ -47,6 +47,7 @@ def create( labels: Optional[Dict[str, str]] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "VideoDataset": """Creates a new video dataset and optionally imports data into dataset when source and import_schema_uri are passed. @@ -117,6 +118,8 @@ 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 the create request in seconds. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -159,4 +162,5 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) diff --git a/google/cloud/aiplatform/featurestore/entity_type.py b/google/cloud/aiplatform/featurestore/entity_type.py index af9152baec..edd0c7433b 100644 --- a/google/cloud/aiplatform/featurestore/entity_type.py +++ b/google/cloud/aiplatform/featurestore/entity_type.py @@ -184,6 +184,7 @@ def update( description: Optional[str] = None, labels: Optional[Dict[str, str]] = None, request_metadata: Sequence[Tuple[str, str]] = (), + update_request_timeout: Optional[float] = None, ) -> "EntityType": """Updates an existing managed entityType resource. @@ -216,6 +217,8 @@ def update( "aiplatform.googleapis.com/" and are immutable. request_metadata (Sequence[Tuple[str, str]]): Required. Strings which should be sent along with the request as metadata. + update_request_timeout (float): + Optional. The timeout for the update request in seconds. Returns: EntityType - The updated entityType resource object. """ @@ -247,6 +250,7 @@ def update( entity_type=gapic_entity_type, update_mask=update_mask, metadata=request_metadata, + timeout=update_request_timeout, ) _LOGGER.log_action_started_against_resource_with_lro( @@ -480,6 +484,7 @@ def create( credentials: Optional[auth_credentials.Credentials] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "EntityType": """Creates an EntityType resource in a Featurestore. @@ -541,7 +546,8 @@ def create( Optional. Whether to execute this creation synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. - + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: EntityType - entity_type resource object @@ -580,6 +586,7 @@ def create( entity_type=gapic_entity_type, entity_type_id=entity_type_id, metadata=request_metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_with_lro(cls, created_entity_type_lro) @@ -605,6 +612,7 @@ def create_feature( labels: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "featurestore.Feature": """Creates a Feature resource in this EntityType. @@ -649,6 +657,8 @@ def create_feature( "aiplatform.googleapis.com/" and are immutable. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. sync (bool): Optional. Whether to execute this creation synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -667,6 +677,7 @@ def create_feature( labels=labels, request_metadata=request_metadata, sync=sync, + create_request_timeout=create_request_timeout, ) def _validate_and_get_create_feature_requests( @@ -918,6 +929,7 @@ def _import_feature_values( self, import_feature_values_request: gca_featurestore_service.ImportFeatureValuesRequest, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ingest_request_timeout: Optional[float] = None, ) -> "EntityType": """Imports Feature values into the Featurestore from a source storage. @@ -926,7 +938,8 @@ def _import_feature_values( Required. Request message for importing feature values. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. - + ingest_request_timeout (float): + Optional. The timeout for the ingest request in seconds. Returns: EntityType - The entityType resource object with imported feature values. """ @@ -939,6 +952,7 @@ def _import_feature_values( import_lro = self.api_client.import_feature_values( request=import_feature_values_request, metadata=request_metadata, + timeout=ingest_request_timeout, ) _LOGGER.log_action_started_against_resource_with_lro( @@ -965,6 +979,7 @@ def ingest_from_bq( worker_count: Optional[int] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + ingest_request_timeout: Optional[float] = None, ) -> "EntityType": """Ingest feature values from BigQuery. @@ -1026,6 +1041,8 @@ def ingest_from_bq( Optional. Whether to execute this import synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + ingest_request_timeout (float): + Optional. The timeout for the ingest request in seconds. Returns: EntityType - The entityType resource object with feature values imported. @@ -1050,6 +1067,7 @@ def ingest_from_bq( return self._import_feature_values( import_feature_values_request=import_feature_values_request, request_metadata=request_metadata, + ingest_request_timeout=ingest_request_timeout, ) @base.optional_sync(return_input_arg="self") @@ -1065,6 +1083,7 @@ def ingest_from_gcs( worker_count: Optional[int] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + ingest_request_timeout: Optional[float] = None, ) -> "EntityType": """Ingest feature values from GCS. @@ -1134,6 +1153,8 @@ def ingest_from_gcs( Optional. Whether to execute this import synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + ingest_request_timeout (float): + Optional. The timeout for the ingest request in seconds. Returns: EntityType - The entityType resource object with feature values imported. @@ -1175,6 +1196,7 @@ def ingest_from_gcs( return self._import_feature_values( import_feature_values_request=import_feature_values_request, request_metadata=request_metadata, + ingest_request_timeout=ingest_request_timeout, ) def ingest_from_df( @@ -1185,6 +1207,7 @@ def ingest_from_df( feature_source_fields: Optional[Dict[str, str]] = None, entity_id_field: Optional[str] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + ingest_request_timeout: Optional[float] = None, ) -> "EntityType": """Ingest feature values from DataFrame. @@ -1245,6 +1268,8 @@ def ingest_from_df( IDs are extracted from the column named ``entity_id``. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + ingest_request_timeout (float): + Optional. The timeout for the ingest request in seconds. Returns: EntityType - The entityType resource object with feature values imported. @@ -1321,6 +1346,7 @@ def ingest_from_df( feature_source_fields=feature_source_fields, entity_id_field=entity_id_field, request_metadata=request_metadata, + ingest_request_timeout=ingest_request_timeout, ) finally: @@ -1386,6 +1412,7 @@ def read( entity_ids: Union[str, List[str]], feature_ids: Union[str, List[str]] = "*", request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + read_request_timeout: Optional[float] = None, ) -> "pd.DataFrame": # noqa: F821 - skip check for undefined name 'pd' """Reads feature values for given feature IDs of given entity IDs in this EntityType. @@ -1398,6 +1425,8 @@ def read( for reading feature values. Default to "*", where value of all features will be read. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + read_request_timeout (float): + Optional. The timeout for the read request in seconds. Returns: pd.DataFrame: entities' feature values in DataFrame @@ -1420,7 +1449,9 @@ def read( ) read_feature_values_response = ( self._featurestore_online_client.read_feature_values( - request=read_feature_values_request, metadata=request_metadata + request=read_feature_values_request, + metadata=request_metadata, + timeout=read_request_timeout, ) ) header = read_feature_values_response.header @@ -1438,6 +1469,7 @@ def read( for response in 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 diff --git a/google/cloud/aiplatform/featurestore/feature.py b/google/cloud/aiplatform/featurestore/feature.py index 7c7f55778e..7a7fc0f29a 100644 --- a/google/cloud/aiplatform/featurestore/feature.py +++ b/google/cloud/aiplatform/featurestore/feature.py @@ -174,6 +174,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, ) -> "Feature": """Updates an existing managed feature resource. @@ -207,6 +208,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 update request in seconds. Returns: Feature - The updated feature resource object. @@ -239,6 +242,7 @@ def update( feature=gapic_feature, update_mask=update_mask, metadata=request_metadata, + timeout=update_request_timeout, ) _LOGGER.log_action_started_against_resource_with_lro( @@ -507,6 +511,7 @@ def create( credentials: Optional[auth_credentials.Credentials] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "Feature": """Creates a Feature resource in an EntityType. @@ -577,6 +582,8 @@ def create( Optional. Whether to execute this creation synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: Feature - feature resource object @@ -618,6 +625,7 @@ def create( created_feature_lro = api_client.create_feature( request=create_feature_request, metadata=request_metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_with_lro(cls, created_feature_lro) diff --git a/google/cloud/aiplatform/featurestore/featurestore.py b/google/cloud/aiplatform/featurestore/featurestore.py index 6c20dd4870..44e1150712 100644 --- a/google/cloud/aiplatform/featurestore/featurestore.py +++ b/google/cloud/aiplatform/featurestore/featurestore.py @@ -138,6 +138,7 @@ def update( self, labels: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + update_request_timeout: Optional[float] = None, ) -> "Featurestore": """Updates an existing managed featurestore resource. @@ -167,18 +168,25 @@ 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 update request in seconds. Returns: Featurestore - The updated featurestore resource object. """ - return self._update(labels=labels, request_metadata=request_metadata) + return self._update( + labels=labels, + request_metadata=request_metadata, + update_request_timeout=update_request_timeout, + ) # TODO(b/206818784): Add enable_online_store and disable_online_store methods def update_online_store( self, fixed_node_count: int, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + update_request_timeout: Optional[float] = None, ) -> "Featurestore": """Updates the online store of an existing managed featurestore resource. @@ -196,12 +204,16 @@ def update_online_store( Required. Config for online serving resources, can only update the node count to >= 1. 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 update request in seconds. Returns: Featurestore - The updated featurestore resource object. """ return self._update( - fixed_node_count=fixed_node_count, request_metadata=request_metadata + fixed_node_count=fixed_node_count, + request_metadata=request_metadata, + update_request_timeout=update_request_timeout, ) def _update( @@ -209,6 +221,7 @@ def _update( labels: Optional[Dict[str, str]] = None, fixed_node_count: Optional[int] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + update_request_timeout: Optional[float] = None, ) -> "Featurestore": """Updates an existing managed featurestore resource. @@ -231,6 +244,8 @@ def _update( Optional. Config for online serving resources, can only update the node count to >= 1. 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 update request in seconds. Returns: Featurestore - The updated featurestore resource object. @@ -265,6 +280,7 @@ def _update( featurestore=gapic_featurestore, update_mask=update_mask, metadata=request_metadata, + timeout=update_request_timeout, ) _LOGGER.log_action_started_against_resource_with_lro( @@ -409,6 +425,7 @@ def create( request_metadata: Optional[Sequence[Tuple[str, str]]] = (), encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "Featurestore": """Creates a Featurestore resource. @@ -470,6 +487,8 @@ def create( Optional. Whether to execute this creation synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: Featurestore - Featurestore resource object @@ -501,6 +520,7 @@ def create( featurestore=gapic_featurestore, featurestore_id=featurestore_id, metadata=request_metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_with_lro(cls, created_featurestore_lro) @@ -525,6 +545,7 @@ def create_entity_type( labels: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "featurestore.EntityType": """Creates an EntityType resource in this Featurestore. @@ -565,6 +586,8 @@ def create_entity_type( "aiplatform.googleapis.com/" and are immutable. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. sync (bool): Optional. Whether to execute this creation synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -582,12 +605,14 @@ def create_entity_type( labels=labels, request_metadata=request_metadata, sync=sync, + create_request_timeout=create_request_timeout, ) def _batch_read_feature_values( self, batch_read_feature_values_request: gca_featurestore_service.BatchReadFeatureValuesRequest, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + serve_request_timeout: Optional[float] = None, ) -> "Featurestore": """Batch read Feature values from the Featurestore to a destination storage. @@ -596,6 +621,8 @@ def _batch_read_feature_values( Required. Request of batch read feature values. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + serve_request_timeout (float): + Optional. The timeout for the serve request in seconds. Returns: Featurestore: The featurestore resource object batch read feature values from. @@ -610,6 +637,7 @@ def _batch_read_feature_values( batch_read_lro = self.api_client.batch_read_feature_values( request=batch_read_feature_values_request, metadata=request_metadata, + timeout=serve_request_timeout, ) _LOGGER.log_action_started_against_resource_with_lro( @@ -802,6 +830,7 @@ def batch_serve_to_bq( pass_through_fields: Optional[List[str]] = None, feature_destination_fields: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + serve_request_timeout: Optional[float] = None, sync: bool = True, ) -> "Featurestore": """Batch serves feature values to BigQuery destination @@ -874,7 +903,8 @@ def batch_serve_to_bq( 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', } - + serve_request_timeout (float): + Optional. The timeout for the serve request in seconds. Returns: Featurestore: The featurestore resource object batch read feature values from. @@ -900,6 +930,7 @@ def batch_serve_to_bq( return self._batch_read_feature_values( batch_read_feature_values_request=batch_read_feature_values_request, request_metadata=request_metadata, + serve_request_timeout=serve_request_timeout, ) @base.optional_sync(return_input_arg="self") @@ -913,6 +944,7 @@ def batch_serve_to_gcs( feature_destination_fields: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), sync: bool = True, + serve_request_timeout: Optional[float] = None, ) -> "Featurestore": """Batch serves feature values to GCS destination @@ -1005,6 +1037,8 @@ def batch_serve_to_gcs( 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', } + serve_request_timeout (float): + Optional. The timeout for the serve request in seconds. Returns: Featurestore: The featurestore resource object batch read feature values from. @@ -1047,6 +1081,7 @@ def batch_serve_to_gcs( return self._batch_read_feature_values( batch_read_feature_values_request=batch_read_feature_values_request, request_metadata=request_metadata, + serve_request_timeout=serve_request_timeout, ) def batch_serve_to_df( @@ -1056,6 +1091,7 @@ def batch_serve_to_df( pass_through_fields: Optional[List[str]] = None, feature_destination_fields: Optional[Dict[str, str]] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + serve_request_timeout: Optional[float] = None, ) -> "pd.DataFrame": # noqa: F821 - skip check for undefined name 'pd' """Batch serves feature values to pandas DataFrame @@ -1137,6 +1173,8 @@ def batch_serve_to_df( 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id1/features/f_id11': 'foo', 'projects/123/locations/us-central1/featurestores/fs_id/entityTypes/et_id2/features/f_id22': 'bar', } + serve_request_timeout (float): + Optional. The timeout for the serve request in seconds. Returns: pd.DataFrame: The pandas DataFrame containing feature values from batch serving. @@ -1210,6 +1248,7 @@ def batch_serve_to_df( pass_through_fields=pass_through_fields, feature_destination_fields=feature_destination_fields, request_metadata=request_metadata, + serve_request_timeout=serve_request_timeout, ) bigquery_storage_read_client = bigquery_storage.BigQueryReadClient( diff --git a/google/cloud/aiplatform/jobs.py b/google/cloud/aiplatform/jobs.py index 7ea5c36960..0a12d1b4fc 100644 --- a/google/cloud/aiplatform/jobs.py +++ b/google/cloud/aiplatform/jobs.py @@ -374,6 +374,7 @@ def create( credentials: Optional[auth_credentials.Credentials] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "BatchPredictionJob": """Create a batch prediction job. @@ -530,7 +531,8 @@ def create( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. - + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: (jobs.BatchPredictionJob): Instantiated representation of the created batch prediction job. @@ -670,6 +672,7 @@ def create( gca_batch_prediction_job=gapic_batch_prediction_job, generate_explanation=generate_explanation, sync=sync, + create_request_timeout=create_request_timeout, ) @classmethod @@ -681,6 +684,7 @@ def _create( gca_batch_prediction_job: gca_bp_job_compat.BatchPredictionJob, generate_explanation: bool, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "BatchPredictionJob": """Create a batch prediction job. @@ -695,6 +699,8 @@ def _create( generate_explanation (bool): Required. Generate explanation along with the batch prediction results. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: (jobs.BatchPredictionJob): Instantiated representation of the created batch prediction job. @@ -725,7 +731,9 @@ def _create( _LOGGER.log_create_with_lro(cls) gca_batch_prediction_job = api_client.create_batch_prediction_job( - parent=parent, batch_prediction_job=gca_batch_prediction_job + parent=parent, + batch_prediction_job=gca_batch_prediction_job, + timeout=create_request_timeout, ) empty_batch_prediction_job._gca_resource = gca_batch_prediction_job @@ -1393,6 +1401,7 @@ def run( enable_web_access: bool = False, tensorboard: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> None: """Run this configured CustomJob. @@ -1434,6 +1443,8 @@ def run( sync (bool): Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. """ if service_account: @@ -1458,7 +1469,9 @@ def run( _LOGGER.log_create_with_lro(self.__class__) self._gca_resource = self.api_client.create_custom_job( - parent=self._parent, custom_job=self._gca_resource + parent=self._parent, + custom_job=self._gca_resource, + timeout=create_request_timeout, ) _LOGGER.log_create_complete_with_getter( @@ -1776,6 +1789,7 @@ def run( enable_web_access: bool = False, tensorboard: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> None: """Run this configured CustomJob. @@ -1817,6 +1831,8 @@ def run( sync (bool): Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. """ if service_account: @@ -1843,7 +1859,9 @@ def run( _LOGGER.log_create_with_lro(self.__class__) self._gca_resource = self.api_client.create_hyperparameter_tuning_job( - parent=self._parent, hyperparameter_tuning_job=self._gca_resource + parent=self._parent, + hyperparameter_tuning_job=self._gca_resource, + timeout=create_request_timeout, ) _LOGGER.log_create_complete_with_getter( diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index bc03ea918d..ae019029fa 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -206,6 +206,7 @@ def create( credentials: Optional[auth_credentials.Credentials] = None, encryption_spec_key_name: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> "Endpoint": """Creates a new endpoint. @@ -253,6 +254,8 @@ def create( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: endpoint (endpoint.Endpoint): Created endpoint. @@ -280,6 +283,7 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) @classmethod @@ -296,6 +300,7 @@ def _create( credentials: Optional[auth_credentials.Credentials] = None, encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> "Endpoint": """Creates a new endpoint by calling the API client. @@ -339,6 +344,8 @@ def _create( If set, this Dataset and all sub-resources of this Dataset will be secured by this key. sync (bool): Whether to create this endpoint synchronously. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: endpoint (endpoint.Endpoint): Created endpoint. @@ -356,7 +363,10 @@ def _create( ) operation_future = api_client.create_endpoint( - parent=parent, endpoint=gapic_endpoint, metadata=metadata + parent=parent, + endpoint=gapic_endpoint, + metadata=metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_with_lro(cls, operation_future) @@ -603,6 +613,7 @@ def deploy( explanation_parameters: Optional[explain.ExplanationParameters] = None, metadata: Optional[Sequence[Tuple[str, str]]] = (), sync=True, + deploy_request_timeout: Optional[float] = None, ) -> None: """Deploys a Model to the Endpoint. @@ -674,6 +685,8 @@ def deploy( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + deploy_request_timeout (float): + Optional. The timeout for the deploy request in seconds. """ self._sync_gca_resource_if_skipped() @@ -703,6 +716,7 @@ def deploy( explanation_parameters=explanation_parameters, metadata=metadata, sync=sync, + deploy_request_timeout=deploy_request_timeout, ) @base.optional_sync() @@ -722,6 +736,7 @@ def _deploy( explanation_parameters: Optional[explain.ExplanationParameters] = None, metadata: Optional[Sequence[Tuple[str, str]]] = (), sync=True, + deploy_request_timeout: Optional[float] = None, ) -> None: """Deploys a Model to the Endpoint. @@ -793,6 +808,8 @@ def _deploy( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + deploy_request_timeout (float): + Optional. The timeout for the deploy request in seconds. Raises: ValueError: If there is not current traffic split and traffic percentage is not 0 or 100. @@ -818,6 +835,7 @@ def _deploy( explanation_metadata=explanation_metadata, explanation_parameters=explanation_parameters, metadata=metadata, + deploy_request_timeout=deploy_request_timeout, ) _LOGGER.log_action_completed_against_resource("model", "deployed", self) @@ -843,6 +861,7 @@ def _deploy_call( explanation_metadata: Optional[explain.ExplanationMetadata] = None, explanation_parameters: Optional[explain.ExplanationParameters] = None, metadata: Optional[Sequence[Tuple[str, str]]] = (), + deploy_request_timeout: Optional[float] = None, ): """Helper method to deploy model to endpoint. @@ -914,6 +933,8 @@ def _deploy_call( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + deploy_request_timeout (float): + Optional. The timeout for the deploy request in seconds. Raises: ValueError: If there is not current traffic split and traffic percentage is not 0 or 100. @@ -1023,6 +1044,7 @@ def _deploy_call( deployed_model=deployed_model, traffic_split=traffic_split, metadata=metadata, + timeout=deploy_request_timeout, ) _LOGGER.log_action_started_against_resource_with_lro( @@ -1655,6 +1677,7 @@ def upload( encryption_spec_key_name: Optional[str] = None, staging_bucket: Optional[str] = None, sync=True, + upload_request_timeout: Optional[float] = None, ) -> "Model": """Uploads a model and returns a Model representing the uploaded Model resource. @@ -1800,6 +1823,8 @@ def upload( staging_bucket (str): Optional. Bucket to stage local model artifacts. Overrides staging_bucket set in aiplatform.init. + upload_request_timeout (float): + Optional. The timeout for the upload request in seconds. Returns: model: Instantiated representation of the uploaded model resource. Raises: @@ -1906,6 +1931,7 @@ def upload( lro = api_client.upload_model( parent=initializer.global_config.common_location_path(project, location), model=managed_model, + timeout=upload_request_timeout, ) _LOGGER.log_create_with_lro(cls, lro) @@ -1936,6 +1962,7 @@ def deploy( metadata: Optional[Sequence[Tuple[str, str]]] = (), encryption_spec_key_name: Optional[str] = None, sync=True, + deploy_request_timeout: Optional[float] = None, ) -> Endpoint: """Deploys model to endpoint. Endpoint will be created if unspecified. @@ -2018,6 +2045,8 @@ def deploy( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + deploy_request_timeout (float): + Optional. The timeout for the deploy request in seconds. Returns: endpoint ("Endpoint"): Endpoint with the deployed model. @@ -2051,6 +2080,7 @@ def deploy( encryption_spec_key_name=encryption_spec_key_name or initializer.global_config.encryption_spec_key_name, sync=sync, + deploy_request_timeout=deploy_request_timeout, ) @base.optional_sync(return_input_arg="endpoint", bind_future_to_self=False) @@ -2071,6 +2101,7 @@ def _deploy( metadata: Optional[Sequence[Tuple[str, str]]] = (), encryption_spec_key_name: Optional[str] = None, sync: bool = True, + deploy_request_timeout: Optional[float] = None, ) -> Endpoint: """Deploys model to endpoint. Endpoint will be created if unspecified. @@ -2153,6 +2184,8 @@ def _deploy( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + deploy_request_timeout (float): + Optional. The timeout for the deploy request in seconds. Returns: endpoint ("Endpoint"): Endpoint with the deployed model. @@ -2187,6 +2220,7 @@ def _deploy( explanation_metadata=explanation_metadata, explanation_parameters=explanation_parameters, metadata=metadata, + deploy_request_timeout=deploy_request_timeout, ) _LOGGER.log_action_completed_against_resource("model", "deployed", endpoint) @@ -2217,6 +2251,7 @@ def batch_predict( credentials: Optional[auth_credentials.Credentials] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> jobs.BatchPredictionJob: """Creates a batch prediction job using this Model and outputs prediction results to the provided destination prefix in the specified @@ -2373,6 +2408,8 @@ def batch_predict( If set, this Model and all sub-resources of this Model will be secured by this key. Overrides encryption_spec_key_name set in aiplatform.init. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: (jobs.BatchPredictionJob): Instantiated representation of the created batch prediction job. @@ -2402,6 +2439,7 @@ def batch_predict( credentials=credentials or self.credentials, encryption_spec_key_name=encryption_spec_key_name, sync=sync, + create_request_timeout=create_request_timeout, ) @classmethod @@ -2612,6 +2650,7 @@ def upload_xgboost_model_file( encryption_spec_key_name: Optional[str] = None, staging_bucket: Optional[str] = None, sync=True, + upload_request_timeout: Optional[float] = None, ) -> "Model": """Uploads a model and returns a Model representing the uploaded Model resource. @@ -2721,6 +2760,8 @@ def upload_xgboost_model_file( staging_bucket (str): Optional. Bucket to stage local model artifacts. Overrides staging_bucket set in aiplatform.init. + upload_request_timeout (float): + Optional. The timeout for the upload request in seconds. Returns: model: Instantiated representation of the uploaded model resource. Raises: @@ -2785,6 +2826,7 @@ def upload_xgboost_model_file( encryption_spec_key_name=encryption_spec_key_name, staging_bucket=staging_bucket, sync=True, + upload_request_timeout=upload_request_timeout, ) @classmethod @@ -2807,6 +2849,7 @@ def upload_scikit_learn_model_file( encryption_spec_key_name: Optional[str] = None, staging_bucket: Optional[str] = None, sync=True, + upload_request_timeout: Optional[float] = None, ) -> "Model": """Uploads a model and returns a Model representing the uploaded Model resource. @@ -2917,6 +2960,8 @@ def upload_scikit_learn_model_file( staging_bucket (str): Optional. Bucket to stage local model artifacts. Overrides staging_bucket set in aiplatform.init. + upload_request_timeout (float): + Optional. The timeout for the upload request in seconds. Returns: model: Instantiated representation of the uploaded model resource. Raises: @@ -2980,6 +3025,7 @@ def upload_scikit_learn_model_file( encryption_spec_key_name=encryption_spec_key_name, staging_bucket=staging_bucket, sync=True, + upload_request_timeout=upload_request_timeout, ) @classmethod @@ -3002,6 +3048,7 @@ def upload_tensorflow_saved_model( encryption_spec_key_name: Optional[str] = None, staging_bucket: Optional[str] = None, sync=True, + upload_request_timeout: Optional[str] = None, ) -> "Model": """Uploads a model and returns a Model representing the uploaded Model resource. @@ -3114,6 +3161,8 @@ def upload_tensorflow_saved_model( staging_bucket (str): Optional. Bucket to stage local model artifacts. Overrides staging_bucket set in aiplatform.init. + upload_request_timeout (float): + Optional. The timeout for the upload request in seconds. Returns: model: Instantiated representation of the uploaded model resource. Raises: @@ -3145,4 +3194,5 @@ def upload_tensorflow_saved_model( encryption_spec_key_name=encryption_spec_key_name, staging_bucket=staging_bucket, sync=sync, + upload_request_timeout=upload_request_timeout, ) diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index b1724e47d6..bc2ddee0e8 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -235,6 +235,7 @@ def run( service_account: Optional[str] = None, network: Optional[str] = None, sync: Optional[bool] = True, + create_request_timeout: Optional[float] = None, ) -> None: """Run this configured PipelineJob and monitor the job until completion. @@ -250,8 +251,14 @@ def run( If left unspecified, the job is not peered with any network. sync (bool): Optional. Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. """ - self.submit(service_account=service_account, network=network) + self.submit( + service_account=service_account, + network=network, + create_request_timeout=create_request_timeout, + ) self._block_until_complete() @@ -259,6 +266,7 @@ def submit( self, service_account: Optional[str] = None, network: Optional[str] = None, + create_request_timeout: Optional[float] = None, ) -> None: """Run this configured PipelineJob. @@ -272,6 +280,8 @@ def submit( Private services access must already be configured for the network. If left unspecified, the job is not peered with any network. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. """ if service_account: self._gca_resource.service_account = service_account @@ -289,6 +299,7 @@ def submit( parent=self._parent, pipeline_job=self._gca_resource, pipeline_job_id=self.job_id, + timeout=create_request_timeout, ) _LOGGER.log_create_complete_with_getter( diff --git a/google/cloud/aiplatform/tensorboard/tensorboard_resource.py b/google/cloud/aiplatform/tensorboard/tensorboard_resource.py index 9821df6636..2bfab2c41b 100644 --- a/google/cloud/aiplatform/tensorboard/tensorboard_resource.py +++ b/google/cloud/aiplatform/tensorboard/tensorboard_resource.py @@ -90,6 +90,7 @@ def create( credentials: Optional[auth_credentials.Credentials] = None, request_metadata: Optional[Sequence[Tuple[str, str]]] = (), encryption_spec_key_name: Optional[str] = None, + create_request_timeout: Optional[float] = None, ) -> "Tensorboard": """Creates a new tensorboard. @@ -143,6 +144,8 @@ def create( If set, this Tensorboard and all sub-resources of this Tensorboard will be secured by this key. Overrides encryption_spec_key_name set in aiplatform.init. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: tensorboard (Tensorboard): @@ -171,7 +174,10 @@ def create( ) create_tensorboard_lro = api_client.create_tensorboard( - parent=parent, tensorboard=gapic_tensorboard, metadata=request_metadata + parent=parent, + tensorboard=gapic_tensorboard, + metadata=request_metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_with_lro(cls, create_tensorboard_lro) @@ -364,6 +370,7 @@ def create( location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, request_metadata: Sequence[Tuple[str, str]] = (), + create_request_timeout: Optional[float] = None, ) -> "TensorboardExperiment": """Creates a new TensorboardExperiment. @@ -423,6 +430,8 @@ def create( credentials set in aiplatform.init. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: TensorboardExperiment: The TensorboardExperiment resource. """ @@ -457,6 +466,7 @@ def create( tensorboard_experiment=gapic_tensorboard_experiment, tensorboard_experiment_id=tensorboard_experiment_id, metadata=request_metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_complete(cls, tensorboard_experiment, "tb experiment") @@ -618,6 +628,7 @@ def create( location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, request_metadata: Sequence[Tuple[str, str]] = (), + create_request_timeout: Optional[float] = None, ) -> "TensorboardRun": """Creates a new tensorboard. @@ -680,6 +691,8 @@ def create( credentials set in aiplatform.init. request_metadata (Sequence[Tuple[str, str]]): Optional. Strings which should be sent along with the request as metadata. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: TensorboardExperiment: The TensorboardExperiment resource. """ @@ -717,6 +730,7 @@ def create( tensorboard_run=gapic_tensorboard_run, tensorboard_run_id=tensorboard_run_id, metadata=request_metadata, + timeout=create_request_timeout, ) _LOGGER.log_create_complete(cls, tensorboard_run, "tb_run") diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py index 86b7ef68d4..dd64b7f2af 100644 --- a/google/cloud/aiplatform/training_jobs.py +++ b/google/cloud/aiplatform/training_jobs.py @@ -568,6 +568,7 @@ def _run_job( model: Optional[gca_model.Model] = None, gcs_destination_uri_prefix: Optional[str] = None, bigquery_destination: Optional[str] = None, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Runs the training job. @@ -713,6 +714,8 @@ def _run_job( - AIP_TRAINING_DATA_URI ="bigquery_destination.dataset_*.training" - AIP_VALIDATION_DATA_URI = "bigquery_destination.dataset_*.validation" - AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test" + create_request_timeout (float): + Optional. The timeout for the create request in seconds. """ input_data_config = self._create_input_data_config( @@ -746,6 +749,7 @@ def _run_job( self.project, self.location ), training_pipeline=training_pipeline, + timeout=create_request_timeout, ) self._gca_resource = training_pipeline @@ -1826,6 +1830,7 @@ def run( enable_web_access: bool = False, tensorboard: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Runs the custom training job. @@ -2069,6 +2074,8 @@ def run( `service_account` is required with provided `tensorboard`. For more information on configuring your service account please visit: https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training + create_request_timeout (float): + Optional. The timeout for the create request in seconds. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -2124,6 +2131,7 @@ def run( if reduction_server_replica_count > 0 else None, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync(construct_object_on_arg="managed_model") @@ -2161,6 +2169,7 @@ def _run( tensorboard: Optional[str] = None, reduction_server_container_uri: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Packages local script and launches training_job. @@ -2309,6 +2318,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float) + Optional. The timeout for the create request in seconds Returns: model: The trained Vertex AI Model resource or None if training did not @@ -2378,6 +2389,7 @@ def _run( model=managed_model, gcs_destination_uri_prefix=base_output_dir, bigquery_destination=bigquery_destination, + create_request_timeout=create_request_timeout, ) return model @@ -2652,6 +2664,7 @@ def run( enable_web_access: bool = False, tensorboard: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Runs the custom training job. @@ -2892,6 +2905,8 @@ def run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -2942,6 +2957,7 @@ def run( if reduction_server_replica_count > 0 else None, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync(construct_object_on_arg="managed_model") @@ -2978,6 +2994,7 @@ def _run( tensorboard: Optional[str] = None, reduction_server_container_uri: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Packages local script and launches training_job. Args: @@ -3122,6 +3139,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -3185,6 +3204,7 @@ def _run( model=managed_model, gcs_destination_uri_prefix=base_output_dir, bigquery_destination=bigquery_destination, + create_request_timeout=create_request_timeout, ) return model @@ -3392,6 +3412,7 @@ def run( export_evaluated_data_items_override_destination: bool = False, additional_experiments: Optional[List[str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -3523,6 +3544,8 @@ def run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not produce a Vertex AI Model. @@ -3561,6 +3584,7 @@ def run( export_evaluated_data_items_bigquery_destination_uri=export_evaluated_data_items_bigquery_destination_uri, export_evaluated_data_items_override_destination=export_evaluated_data_items_override_destination, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync() @@ -3582,6 +3606,7 @@ def _run( export_evaluated_data_items_bigquery_destination_uri: Optional[str] = None, export_evaluated_data_items_override_destination: bool = False, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -3711,6 +3736,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -3784,6 +3811,7 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=timestamp_split_column_name, model=model, + create_request_timeout=create_request_timeout, ) @property @@ -3991,6 +4019,7 @@ def run( model_labels: Optional[Dict[str, str]] = None, additional_experiments: Optional[List[str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -4140,6 +4169,8 @@ def run( and examples of labels. additional_experiments (List[str]): Optional. Additional experiment flags for the time series forcasting training. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will @@ -4194,6 +4225,7 @@ def run( model_display_name=model_display_name, model_labels=model_labels, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync() @@ -4224,6 +4256,7 @@ def _run( model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -4383,6 +4416,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not produce a Vertex AI Model. @@ -4462,6 +4497,7 @@ def _run( predefined_split_column_name=predefined_split_column_name, timestamp_split_column_name=None, # Not supported by AutoMLForecasting model=model, + create_request_timeout=create_request_timeout, ) if export_evaluated_data_items: @@ -4682,6 +4718,7 @@ def run( model_labels: Optional[Dict[str, str]] = None, disable_early_stopping: bool = False, sync: bool = True, + create_request_timeout: Optional[float] = False, ) -> models.Model: """Runs the AutoML Image training job and returns a model. @@ -4787,6 +4824,8 @@ def run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not produce a Vertex AI Model. @@ -4820,6 +4859,7 @@ def run( model_labels=model_labels, disable_early_stopping=disable_early_stopping, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync() @@ -4838,6 +4878,7 @@ def _run( model_labels: Optional[Dict[str, str]] = None, disable_early_stopping: bool = False, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -4943,6 +4984,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -4991,6 +5034,7 @@ def _run( validation_filter_split=validation_filter_split, test_filter_split=test_filter_split, model=model_tbt, + create_request_timeout=create_request_timeout, ) @property @@ -5281,6 +5325,7 @@ def run( enable_web_access: bool = False, tensorboard: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Runs the custom training job. @@ -5521,6 +5566,8 @@ def run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -5566,6 +5613,7 @@ def run( if reduction_server_replica_count > 0 else None, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync(construct_object_on_arg="managed_model") @@ -5602,6 +5650,7 @@ def _run( tensorboard: Optional[str] = None, reduction_server_container_uri: Optional[str] = None, sync=True, + create_request_timeout: Optional[float] = None, ) -> Optional[models.Model]: """Packages local script and launches training_job. @@ -5733,6 +5782,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -5796,6 +5847,7 @@ def _run( model=managed_model, gcs_destination_uri_prefix=base_output_dir, bigquery_destination=bigquery_destination, + create_request_timeout=create_request_timeout, ) return model @@ -5944,6 +5996,7 @@ def run( model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the AutoML Video training job and returns a model. @@ -6011,6 +6064,8 @@ def run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not produce a Vertex AI Model. @@ -6039,6 +6094,7 @@ def run( model_display_name=model_display_name, model_labels=model_labels, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync() @@ -6052,6 +6108,7 @@ def _run( model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -6121,6 +6178,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -6158,6 +6217,7 @@ def _run( validation_filter_split=validation_filter_split, test_filter_split=test_filter_split, model=model_tbt, + create_request_timeout=create_request_timeout, ) @property @@ -6322,6 +6382,7 @@ def run( model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -6396,11 +6457,13 @@ def run( underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information - and examples of labels. + and examples of labels.. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds Returns: model: The trained Vertex AI Model resource. @@ -6430,6 +6493,7 @@ def run( model_display_name=model_display_name, model_labels=model_labels, sync=sync, + create_request_timeout=create_request_timeout, ) @base.optional_sync() @@ -6445,6 +6509,7 @@ def _run( model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> models.Model: """Runs the training job and returns a model. @@ -6526,6 +6591,8 @@ def _run( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: model: The trained Vertex AI Model resource or None if training did not @@ -6549,6 +6616,7 @@ def _run( validation_filter_split=validation_filter_split, test_filter_split=test_filter_split, model=model, + create_request_timeout=create_request_timeout, ) @property diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py index 0167cb8f20..d8d8bd53e3 100644 --- a/tests/system/aiplatform/test_dataset.py +++ b/tests/system/aiplatform/test_dataset.py @@ -196,6 +196,7 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client, shared_state): my_dataset.import_data( gcs_source=_TEST_TEXT_ENTITY_EXTRACTION_GCS_SOURCE, import_schema_uri=_TEST_TEXT_ENTITY_IMPORT_SCHEMA, + import_request_timeout=None, ) data_items_post_import = dataset_gapic_client.list_data_items( @@ -216,6 +217,7 @@ def test_create_and_import_image_dataset(self, dataset_gapic_client, shared_stat display_name=f"temp_sdk_integration_create_and_import_dataset_{uuid.uuid4()}", gcs_source=_TEST_IMAGE_OBJECT_DETECTION_GCS_SOURCE, import_schema_uri=_TEST_IMAGE_OBJ_DET_IMPORT_SCHEMA, + create_request_timeout=None, ) shared_state["dataset_name"] = img_dataset.resource_name @@ -236,6 +238,7 @@ def test_create_tabular_dataset(self, dataset_gapic_client, shared_state): tabular_dataset = aiplatform.TabularDataset.create( display_name=f"temp_sdk_integration_create_and_import_dataset_{uuid.uuid4()}", gcs_source=[_TEST_TABULAR_CLASSIFICATION_GCS_SOURCE], + create_request_timeout=None, ) shared_state["dataset_name"] = tabular_dataset.resource_name diff --git a/tests/system/aiplatform/test_e2e_tabular.py b/tests/system/aiplatform/test_e2e_tabular.py index ee0692014a..31b9bb9769 100644 --- a/tests/system/aiplatform/test_e2e_tabular.py +++ b/tests/system/aiplatform/test_e2e_tabular.py @@ -83,6 +83,7 @@ def test_end_to_end_tabular(self, shared_state): display_name=self._make_display_name("dataset"), gcs_source=[dataset_gcs_source], sync=False, + create_request_timeout=180.0, ) shared_state["resources"].extend([ds]) @@ -113,6 +114,7 @@ def test_end_to_end_tabular(self, shared_state): restart_job_on_worker_restart=True, enable_web_access=True, sync=False, + create_request_timeout=None, ) automl_model = automl_job.run( diff --git a/tests/system/aiplatform/test_tensorboard.py b/tests/system/aiplatform/test_tensorboard.py index 3b7d892214..501205122f 100644 --- a/tests/system/aiplatform/test_tensorboard.py +++ b/tests/system/aiplatform/test_tensorboard.py @@ -32,7 +32,10 @@ def test_create_and_get_tensorboard(self, shared_state): display_name = self._make_display_name("tensorboard") - tb = aiplatform.Tensorboard.create(display_name=display_name) + tb = aiplatform.Tensorboard.create( + display_name=display_name, + create_request_timeout=None, + ) shared_state["resources"] = [tb] @@ -50,6 +53,7 @@ def test_create_and_get_tensorboard(self, shared_state): display_name=self._make_display_name("tensorboard_experiment"), description="Vertex SDK Integration test.", labels={"test": "labels"}, + create_request_timeout=None, ) shared_state["resources"].append(tb_experiment) @@ -71,6 +75,7 @@ def test_create_and_get_tensorboard(self, shared_state): tensorboard_experiment_name=tb_experiment.resource_name, description="Vertex SDK Integration test run", labels={"test": "labels"}, + create_request_timeout=None, ) shared_state["resources"].append(tb_run) diff --git a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py index bf678e91a9..953da7af91 100644 --- a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py @@ -280,6 +280,7 @@ def test_run_call_pipeline_service_create( validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, additional_experiments=_TEST_ADDITIONAL_EXPERIMENTS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -308,6 +309,7 @@ def test_run_call_pipeline_service_create( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource is mock_pipeline_service_get.return_value @@ -324,6 +326,80 @@ def test_run_call_pipeline_service_create( assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_dataset_time_series, + mock_model_service_get, + sync, + ): + aiplatform.init(project=_TEST_PROJECT, staging_bucket=_TEST_BUCKET_NAME) + + job = AutoMLForecastingTrainingJob( + display_name=_TEST_DISPLAY_NAME, + optimization_objective=_TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME, + column_transformations=_TEST_TRAINING_COLUMN_TRANSFORMATIONS, + labels=_TEST_LABELS, + ) + + model_from_job = job.run( + dataset=mock_dataset_time_series, + target_column=_TEST_TRAINING_TARGET_COLUMN, + time_column=_TEST_TRAINING_TIME_COLUMN, + time_series_identifier_column=_TEST_TRAINING_TIME_SERIES_IDENTIFIER_COLUMN, + unavailable_at_forecast_columns=_TEST_TRAINING_UNAVAILABLE_AT_FORECAST_COLUMNS, + available_at_forecast_columns=_TEST_TRAINING_AVAILABLE_AT_FORECAST_COLUMNS, + forecast_horizon=_TEST_TRAINING_FORECAST_HORIZON, + data_granularity_unit=_TEST_TRAINING_DATA_GRANULARITY_UNIT, + data_granularity_count=_TEST_TRAINING_DATA_GRANULARITY_COUNT, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + predefined_split_column_name=_TEST_PREDEFINED_SPLIT_COLUMN_NAME, + weight_column=_TEST_TRAINING_WEIGHT_COLUMN, + time_series_attribute_columns=_TEST_TRAINING_TIME_SERIES_ATTRIBUTE_COLUMNS, + context_window=_TEST_TRAINING_CONTEXT_WINDOW, + budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, + export_evaluated_data_items=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS, + export_evaluated_data_items_bigquery_destination_uri=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_BIGQUERY_DESTINATION_URI, + export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, + quantiles=_TEST_TRAINING_QUANTILES, + validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + additional_experiments=_TEST_ADDITIONAL_EXPERIMENTS, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, labels=_TEST_MODEL_LABELS + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + predefined_split=gca_training_pipeline.PredefinedSplit( + key=_TEST_PREDEFINED_SPLIT_COLUMN_NAME + ), + dataset_id=mock_dataset_time_series.name, + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.automl_forecasting, + training_task_inputs=_TEST_TRAINING_TASK_INPUTS_WITH_ADDITIONAL_EXPERIMENTS, + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.usefixtures("mock_pipeline_service_get") @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( @@ -362,6 +438,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -389,6 +466,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.usefixtures("mock_pipeline_service_get") @@ -430,6 +508,7 @@ def test_run_call_pipeline_if_set_additional_experiments( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -453,6 +532,7 @@ def test_run_call_pipeline_if_set_additional_experiments( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.usefixtures( @@ -634,6 +714,7 @@ def test_splits_fraction( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -667,6 +748,7 @@ def test_splits_fraction( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -716,6 +798,7 @@ def test_splits_predefined( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -747,6 +830,7 @@ def test_splits_predefined( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -795,6 +879,7 @@ def test_splits_default( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -821,4 +906,5 @@ def test_splits_default( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) diff --git a/tests/unit/aiplatform/test_automl_image_training_jobs.py b/tests/unit/aiplatform/test_automl_image_training_jobs.py index 41d95251e8..80187df7ae 100644 --- a/tests/unit/aiplatform/test_automl_image_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_image_training_jobs.py @@ -283,6 +283,7 @@ def test_run_call_pipeline_service_create( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -319,6 +320,7 @@ def test_run_call_pipeline_service_create( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) mock_model_service_get.assert_called_once_with( @@ -330,6 +332,79 @@ def test_run_call_pipeline_service_create( assert not job.has_failed assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_dataset_image, + mock_model_service_get, + mock_model, + sync, + ): + """Create and run an AutoML ICN training job, verify calls and return value""" + + aiplatform.init( + project=_TEST_PROJECT, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + job = training_jobs.AutoMLImageTrainingJob( + display_name=_TEST_DISPLAY_NAME, + base_model=mock_model, + labels=_TEST_LABELS, + ) + + model_from_job = job.run( + dataset=mock_dataset_image, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + training_filter_split=_TEST_FILTER_SPLIT_TRAINING, + validation_filter_split=_TEST_FILTER_SPLIT_VALIDATION, + test_filter_split=_TEST_FILTER_SPLIT_TEST, + budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, + disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_filter_split = gca_training_pipeline.FilterSplit( + training_filter=_TEST_FILTER_SPLIT_TRAINING, + validation_filter=_TEST_FILTER_SPLIT_VALIDATION, + test_filter=_TEST_FILTER_SPLIT_TEST, + ) + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=mock_model._gca_resource.labels, + description=mock_model._gca_resource.description, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + filter_split=true_filter_split, + dataset_id=mock_dataset_image.name, + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.automl_image_classification, + training_task_inputs=_TEST_TRAINING_TASK_INPUTS_WITH_BASE_MODEL, + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.usefixtures("mock_pipeline_service_get") @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( @@ -352,6 +427,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( dataset=mock_dataset_image, budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, + create_request_timeout=None, ) if not sync: @@ -381,6 +457,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.usefixtures( @@ -520,6 +597,7 @@ def test_splits_fraction( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -554,6 +632,7 @@ def test_splits_fraction( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -589,6 +668,7 @@ def test_splits_filter( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -623,6 +703,7 @@ def test_splits_filter( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -655,6 +736,7 @@ def test_splits_default( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -682,6 +764,7 @@ def test_splits_default( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) def test_splits_filter_incomplete( diff --git a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py index cf19048249..b15dbd6a0e 100644 --- a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py @@ -341,6 +341,7 @@ def test_run_call_pipeline_service_create( disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, additional_experiments=_TEST_ADDITIONAL_EXPERIMENTS, sync=sync, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -373,6 +374,7 @@ def test_run_call_pipeline_service_create( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource is mock_pipeline_service_get.return_value @@ -389,6 +391,75 @@ def test_run_call_pipeline_service_create( assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_dataset_tabular, + mock_model_service_get, + sync, + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_BUCKET_NAME, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + job = training_jobs.AutoMLTabularTrainingJob( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + optimization_objective=_TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME, + optimization_prediction_type=_TEST_TRAINING_OPTIMIZATION_PREDICTION_TYPE, + column_transformations=_TEST_TRAINING_COLUMN_TRANSFORMATIONS, + optimization_objective_recall_value=None, + optimization_objective_precision_value=None, + ) + + model_from_job = job.run( + dataset=mock_dataset_tabular, + target_column=_TEST_TRAINING_TARGET_COLUMN, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + weight_column=_TEST_TRAINING_WEIGHT_COLUMN, + budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, + disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, + additional_experiments=_TEST_ADDITIONAL_EXPERIMENTS, + sync=sync, + create_request_timeout=180.0, + ) + + job.wait_for_resource_creation() + + if not sync: + model_from_job.wait() + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + dataset_id=mock_dataset_tabular.name, + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.automl_tabular, + training_task_inputs=_TEST_TRAINING_TASK_INPUTS_WITH_ADDITIONAL_EXPERIMENTS, + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_service_create_with_export_eval_data_items( self, @@ -424,6 +495,7 @@ def test_run_call_pipeline_service_create_with_export_eval_data_items( export_evaluated_data_items_bigquery_destination_uri=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_BIGQUERY_DESTINATION_URI, export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, sync=sync, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -454,6 +526,7 @@ def test_run_call_pipeline_service_create_with_export_eval_data_items( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource is mock_pipeline_service_get.return_value @@ -499,6 +572,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( weight_column=_TEST_TRAINING_WEIGHT_COLUMN, budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -532,6 +606,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -567,6 +642,7 @@ def test_run_call_pipeline_service_create_if_no_column_transformations( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -597,6 +673,7 @@ def test_run_call_pipeline_service_create_if_no_column_transformations( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -634,6 +711,7 @@ def test_run_call_pipeline_service_create_if_set_additional_experiments( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -664,6 +742,7 @@ def test_run_call_pipeline_service_create_if_set_additional_experiments( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -701,6 +780,7 @@ def test_run_call_pipeline_service_create_with_column_specs( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -723,6 +803,7 @@ def test_run_call_pipeline_service_create_with_column_specs( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -803,6 +884,7 @@ def test_run_call_pipeline_service_create_with_column_specs_not_auto( budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -825,6 +907,7 @@ def test_run_call_pipeline_service_create_with_column_specs_not_auto( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.usefixtures( @@ -1111,6 +1194,7 @@ def test_splits_fraction( test_fraction_split=_TEST_FRACTION_SPLIT_TEST, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1144,6 +1228,7 @@ def test_splits_fraction( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1185,6 +1270,7 @@ def test_splits_timestamp( timestamp_split_column_name=_TEST_SPLIT_TIMESTAMP_COLUMN_NAME, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1219,6 +1305,7 @@ def test_splits_timestamp( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1257,6 +1344,7 @@ def test_splits_predefined( predefined_split_column_name=_TEST_SPLIT_PREDEFINED_COLUMN_NAME, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1288,6 +1376,7 @@ def test_splits_predefined( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1325,6 +1414,7 @@ def test_splits_default( model_display_name=_TEST_MODEL_DISPLAY_NAME, disable_early_stopping=_TEST_TRAINING_DISABLE_EARLY_STOPPING, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1351,4 +1441,5 @@ def test_splits_default( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) diff --git a/tests/unit/aiplatform/test_automl_text_training_jobs.py b/tests/unit/aiplatform/test_automl_text_training_jobs.py index ec1fc26d76..4980034140 100644 --- a/tests/unit/aiplatform/test_automl_text_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_text_training_jobs.py @@ -287,6 +287,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( dataset=mock_dataset_text, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -313,6 +314,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -345,6 +347,7 @@ def test_run_call_pipeline_service_create_classification( validation_filter_split=_TEST_FILTER_SPLIT_VALIDATION, test_filter_split=_TEST_FILTER_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -380,6 +383,7 @@ def test_run_call_pipeline_service_create_classification( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) mock_model_service_get.assert_called_once_with( @@ -391,6 +395,75 @@ def test_run_call_pipeline_service_create_classification( assert not job.has_failed assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_classification_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_dataset_text, + mock_model_service_get, + sync, + ): + """Create and run an AutoML Text Classification training job, verify calls and return value""" + + aiplatform.init(project=_TEST_PROJECT) + + job = training_jobs.AutoMLTextTrainingJob( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + prediction_type=_TEST_PREDICTION_TYPE_CLASSIFICATION, + multi_label=_TEST_CLASSIFICATION_MULTILABEL, + training_encryption_spec_key_name=_TEST_PIPELINE_ENCRYPTION_KEY_NAME, + model_encryption_spec_key_name=_TEST_MODEL_ENCRYPTION_KEY_NAME, + ) + + model_from_job = job.run( + dataset=mock_dataset_text, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + training_filter_split=_TEST_FILTER_SPLIT_TRAINING, + validation_filter_split=_TEST_FILTER_SPLIT_VALIDATION, + test_filter_split=_TEST_FILTER_SPLIT_TEST, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_filter_split = gca_training_pipeline.FilterSplit( + training_filter=_TEST_FILTER_SPLIT_TRAINING, + validation_filter=_TEST_FILTER_SPLIT_VALIDATION, + test_filter=_TEST_FILTER_SPLIT_TEST, + ) + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + filter_split=true_filter_split, + dataset_id=mock_dataset_text.name, + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.automl_text_classification, + training_task_inputs=_TEST_TRAINING_TASK_INPUTS_CLASSIFICATION, + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + encryption_spec=_TEST_PIPELINE_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_service_create_extraction( self, @@ -418,6 +491,7 @@ def test_run_call_pipeline_service_create_extraction( validation_fraction_split=_TEST_FRACTION_SPLIT_VALIDATION, test_fraction_split=_TEST_FRACTION_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -451,6 +525,7 @@ def test_run_call_pipeline_service_create_extraction( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) mock_model_service_get.assert_called_once_with( @@ -490,6 +565,7 @@ def test_run_call_pipeline_service_create_sentiment( validation_filter_split=_TEST_FILTER_SPLIT_VALIDATION, test_filter_split=_TEST_FILTER_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -522,6 +598,7 @@ def test_run_call_pipeline_service_create_sentiment( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) mock_model_service_get.assert_called_once_with( @@ -556,6 +633,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( dataset=mock_dataset_text, model_display_name=None, # Omit model_display_name sync=sync, + create_request_timeout=None, ) if not sync: @@ -583,6 +661,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.usefixtures( @@ -712,6 +791,7 @@ def test_splits_fraction( validation_fraction_split=_TEST_FRACTION_SPLIT_VALIDATION, test_fraction_split=_TEST_FRACTION_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -746,6 +826,7 @@ def test_splits_fraction( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -781,6 +862,7 @@ def test_splits_filter( validation_filter_split=_TEST_FILTER_SPLIT_VALIDATION, test_filter_split=_TEST_FILTER_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -815,6 +897,7 @@ def test_splits_filter( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -847,6 +930,7 @@ def test_splits_default( dataset=mock_dataset_text, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -874,4 +958,5 @@ def test_splits_default( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) diff --git a/tests/unit/aiplatform/test_automl_video_training_jobs.py b/tests/unit/aiplatform/test_automl_video_training_jobs.py index 7333e4a580..e0e2a524e2 100644 --- a/tests/unit/aiplatform/test_automl_video_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_video_training_jobs.py @@ -253,6 +253,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( dataset=mock_dataset_video, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -280,6 +281,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) mock_model_service_get.assert_called_once_with( @@ -323,6 +325,7 @@ def test_splits_fraction( training_fraction_split=_TEST_FRACTION_SPLIT_TRAINING, test_fraction_split=_TEST_FRACTION_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -357,6 +360,7 @@ def test_splits_fraction( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -391,6 +395,7 @@ def test_splits_filter( training_filter_split=_TEST_FILTER_SPLIT_TRAINING, test_filter_split=_TEST_FILTER_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -425,6 +430,7 @@ def test_splits_filter( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -457,6 +463,7 @@ def test_splits_default( dataset=mock_dataset_video, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -484,6 +491,7 @@ def test_splits_default( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -516,6 +524,7 @@ def test_run_call_pipeline_service_create( training_filter_split=_TEST_FILTER_SPLIT_TRAINING, test_filter_split=_TEST_FILTER_SPLIT_TEST, sync=sync, + create_request_timeout=None, ) if not sync: @@ -552,6 +561,7 @@ def test_run_call_pipeline_service_create( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) mock_model_service_get.assert_called_once_with( @@ -563,6 +573,76 @@ def test_run_call_pipeline_service_create( assert not job.has_failed assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_dataset_video, + mock_model_service_get, + mock_model, + sync, + ): + """Create and run an AutoML ICN training job, verify calls and return value""" + + aiplatform.init(project=_TEST_PROJECT) + + job = training_jobs.AutoMLVideoTrainingJob( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + prediction_type=_TEST_PREDICTION_TYPE_VCN, + model_type=_TEST_MODEL_TYPE_CLOUD, + training_encryption_spec_key_name=_TEST_PIPELINE_ENCRYPTION_KEY_NAME, + model_encryption_spec_key_name=_TEST_MODEL_ENCRYPTION_KEY_NAME, + ) + + model_from_job = job.run( + dataset=mock_dataset_video, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + training_filter_split=_TEST_FILTER_SPLIT_TRAINING, + test_filter_split=_TEST_FILTER_SPLIT_TEST, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_filter_split = gca_training_pipeline.FilterSplit( + training_filter=_TEST_FILTER_SPLIT_TRAINING, + validation_filter=_TEST_FILTER_SPLIT_VALIDATION, + test_filter=_TEST_FILTER_SPLIT_TEST, + ) + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + description=mock_model._gca_resource.description, + encryption_spec=_TEST_MODEL_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + filter_split=true_filter_split, + dataset_id=mock_dataset_video.name, + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.automl_video_classification, + training_task_inputs=_TEST_TRAINING_TASK_INPUTS, + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + encryption_spec=_TEST_PIPELINE_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.usefixtures("mock_pipeline_service_get") @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( @@ -585,6 +665,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( dataset=mock_dataset_video, training_fraction_split=_TEST_ALTERNATE_FRACTION_SPLIT_TRAINING, test_fraction_split=_TEST_ALTERNATE_FRACTION_SPLIT_TEST, + create_request_timeout=None, ) if not sync: @@ -619,6 +700,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) @pytest.mark.usefixtures( diff --git a/tests/unit/aiplatform/test_custom_job.py b/tests/unit/aiplatform/test_custom_job.py index 0d13c3388c..f16f3a1cc8 100644 --- a/tests/unit/aiplatform/test_custom_job.py +++ b/tests/unit/aiplatform/test_custom_job.py @@ -296,6 +296,7 @@ def test_create_custom_job(self, create_custom_job_mock, get_custom_job_mock, sy timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -307,7 +308,9 @@ def test_create_custom_job(self, create_custom_job_mock, get_custom_job_mock, sy expected_custom_job = _get_custom_job_proto() create_custom_job_mock.assert_called_once_with( - parent=_TEST_PARENT, custom_job=expected_custom_job + parent=_TEST_PARENT, + custom_job=expected_custom_job, + timeout=None, ) assert job.job_spec == expected_custom_job.job_spec @@ -316,6 +319,48 @@ def test_create_custom_job(self, create_custom_job_mock, get_custom_job_mock, sy ) assert job.network == _TEST_NETWORK + @pytest.mark.parametrize("sync", [True, False]) + def test_create_custom_job_with_timeout( + self, create_custom_job_mock, get_custom_job_mock, sync + ): + + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + staging_bucket=_TEST_STAGING_BUCKET, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + job = aiplatform.CustomJob( + display_name=_TEST_DISPLAY_NAME, + worker_pool_specs=_TEST_WORKER_POOL_SPEC, + base_output_dir=_TEST_BASE_OUTPUT_DIR, + labels=_TEST_LABELS, + ) + + job.run( + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + timeout=_TEST_TIMEOUT, + restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, + sync=sync, + create_request_timeout=180.0, + ) + + job.wait_for_resource_creation() + + assert job.resource_name == _TEST_CUSTOM_JOB_NAME + + job.wait() + + expected_custom_job = _get_custom_job_proto() + + create_custom_job_mock.assert_called_once_with( + parent=_TEST_PARENT, + custom_job=expected_custom_job, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_run_custom_job_with_fail_raises( self, create_custom_job_mock, get_custom_job_mock_with_fail, sync @@ -345,6 +390,7 @@ def test_run_custom_job_with_fail_raises( timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) job.wait() @@ -356,7 +402,9 @@ def test_run_custom_job_with_fail_raises( expected_custom_job = _get_custom_job_proto() create_custom_job_mock.assert_called_once_with( - parent=_TEST_PARENT, custom_job=expected_custom_job + parent=_TEST_PARENT, + custom_job=expected_custom_job, + timeout=None, ) assert job.job_spec == expected_custom_job.job_spec @@ -520,6 +568,7 @@ def test_create_custom_job_with_enable_web_access( timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) job.wait_for_resource_creation() @@ -533,7 +582,9 @@ def test_create_custom_job_with_enable_web_access( expected_custom_job = _get_custom_job_proto_with_enable_web_access() create_custom_job_mock_with_enable_web_access.assert_called_once_with( - parent=_TEST_PARENT, custom_job=expected_custom_job + parent=_TEST_PARENT, + custom_job=expected_custom_job, + timeout=None, ) assert job.job_spec == expected_custom_job.job_spec @@ -588,6 +639,7 @@ def test_create_custom_job_with_tensorboard( timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) job.wait() @@ -596,7 +648,9 @@ def test_create_custom_job_with_tensorboard( expected_custom_job.job_spec.tensorboard = _TEST_TENSORBOARD_NAME create_custom_job_mock_with_tensorboard.assert_called_once_with( - parent=_TEST_PARENT, custom_job=expected_custom_job + parent=_TEST_PARENT, + custom_job=expected_custom_job, + timeout=None, ) expected_custom_job = _get_custom_job_proto() diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index 9400f33cef..bab93280b5 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -578,6 +578,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_dataset( display_name=_TEST_DISPLAY_NAME, metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_NONTABULAR, sync=sync, + create_request_timeout=None, ) if not sync: @@ -594,6 +595,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_dataset( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_dataset_mock") @@ -606,6 +608,7 @@ def test_create_dataset_nontabular(self, create_dataset_mock, sync): metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_NONTABULAR, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -622,6 +625,37 @@ def test_create_dataset_nontabular(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.usefixtures("get_dataset_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_create_dataset_nontabular_with_timeout(self, create_dataset_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_dataset = datasets._Dataset.create( + display_name=_TEST_DISPLAY_NAME, + metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_NONTABULAR, + encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + my_dataset.wait() + + expected_dataset = gca_dataset.Dataset( + display_name=_TEST_DISPLAY_NAME, + metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_NONTABULAR, + metadata=_TEST_NONTABULAR_DATASET_METADATA, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + + create_dataset_mock.assert_called_once_with( + parent=_TEST_PARENT, + dataset=expected_dataset, + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.usefixtures("get_dataset_mock") @@ -633,6 +667,7 @@ def test_create_dataset_tabular(self, create_dataset_mock): metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_TABULAR, bq_source=_TEST_SOURCE_URI_BQ, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + create_request_timeout=None, ) expected_dataset = gca_dataset.Dataset( @@ -646,6 +681,7 @@ def test_create_dataset_tabular(self, create_dataset_mock): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_dataset_mock") @@ -663,6 +699,7 @@ def test_create_and_import_dataset( data_item_labels=_TEST_DATA_LABEL_ITEMS, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -685,10 +722,13 @@ def test_create_and_import_dataset( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -706,6 +746,7 @@ def test_import_data(self, import_data_mock, sync): import_schema_uri=_TEST_IMPORT_SCHEMA_URI, data_item_labels=_TEST_DATA_LABEL_ITEMS, sync=sync, + import_request_timeout=None, ) if not sync: @@ -718,7 +759,39 @@ def test_import_data(self, import_data_mock, sync): ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, + ) + + @pytest.mark.usefixtures("get_dataset_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_import_data_with_timeout(self, import_data_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_dataset = datasets._Dataset(dataset_name=_TEST_NAME) + + my_dataset.import_data( + gcs_source=_TEST_SOURCE_URI_GCS, + import_schema_uri=_TEST_IMPORT_SCHEMA_URI, + data_item_labels=_TEST_DATA_LABEL_ITEMS, + sync=sync, + import_request_timeout=180.0, + ) + + if not sync: + my_dataset.wait() + + expected_import_config = gca_dataset.ImportDataConfig( + gcs_source=gca_io.GcsSource(uris=[_TEST_SOURCE_URI_GCS]), + import_schema_uri=_TEST_IMPORT_SCHEMA_URI, + data_item_labels=_TEST_DATA_LABEL_ITEMS, + ) + + import_data_mock.assert_called_once_with( + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=180.0, ) @pytest.mark.usefixtures("get_dataset_mock") @@ -749,6 +822,7 @@ def test_create_then_import( metadata_schema_uri=_TEST_METADATA_SCHEMA_URI_NONTABULAR, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) my_dataset.import_data( @@ -756,6 +830,7 @@ def test_create_then_import( import_schema_uri=_TEST_IMPORT_SCHEMA_URI, data_item_labels=_TEST_DATA_LABEL_ITEMS, sync=sync, + import_request_timeout=None, ) if not sync: @@ -778,6 +853,7 @@ def test_create_then_import( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) get_dataset_mock.assert_called_once_with( @@ -785,7 +861,9 @@ def test_create_then_import( ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -837,6 +915,7 @@ def test_create_dataset(self, create_dataset_mock, sync): my_dataset = datasets.ImageDataset.create( display_name=_TEST_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -853,6 +932,7 @@ def test_create_dataset(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_dataset_image_mock") @@ -868,6 +948,7 @@ def test_create_and_import_dataset( import_schema_uri=_TEST_IMPORT_SCHEMA_URI_IMAGE, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -884,6 +965,7 @@ def test_create_and_import_dataset( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) expected_import_config = gca_dataset.ImportDataConfig( @@ -891,7 +973,9 @@ def test_create_and_import_dataset( import_schema_uri=_TEST_IMPORT_SCHEMA_URI_IMAGE, ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -908,6 +992,7 @@ def test_import_data(self, import_data_mock, sync): gcs_source=[_TEST_SOURCE_URI_GCS], import_schema_uri=_TEST_IMPORT_SCHEMA_URI_IMAGE, sync=sync, + import_request_timeout=None, ) if not sync: @@ -919,7 +1004,9 @@ def test_import_data(self, import_data_mock, sync): ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -933,12 +1020,14 @@ def test_create_then_import( display_name=_TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) my_dataset.import_data( gcs_source=[_TEST_SOURCE_URI_GCS], import_schema_uri=_TEST_IMPORT_SCHEMA_URI_IMAGE, sync=sync, + import_request_timeout=None, ) if not sync: @@ -954,6 +1043,7 @@ def test_create_then_import( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) get_dataset_image_mock.assert_called_once_with( @@ -966,7 +1056,9 @@ def test_create_then_import( ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -984,6 +1076,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1001,6 +1094,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @@ -1040,6 +1134,7 @@ def test_create_dataset_with_default_encryption_key( display_name=_TEST_DISPLAY_NAME, bq_source=_TEST_SOURCE_URI_BQ, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1058,6 +1153,7 @@ def test_create_dataset_with_default_encryption_key( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("create_dataset_mock_fail") @@ -1098,6 +1194,7 @@ def test_create_dataset(self, create_dataset_mock, sync): bq_source=_TEST_SOURCE_URI_BQ, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1114,6 +1211,7 @@ def test_create_dataset(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_dataset_tabular_bq_mock") @@ -1242,6 +1340,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): labels=_TEST_LABELS, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1259,6 +1358,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @@ -1294,6 +1394,7 @@ def test_create_dataset(self, create_dataset_mock, sync): my_dataset = datasets.TextDataset.create( display_name=_TEST_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1310,6 +1411,7 @@ def test_create_dataset(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_dataset_text_mock") @@ -1325,6 +1427,7 @@ def test_create_and_import_dataset( import_schema_uri=_TEST_IMPORT_SCHEMA_URI_TEXT, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1341,6 +1444,7 @@ def test_create_and_import_dataset( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) expected_import_config = gca_dataset.ImportDataConfig( @@ -1348,7 +1452,9 @@ def test_create_and_import_dataset( import_schema_uri=_TEST_IMPORT_SCHEMA_URI_TEXT, ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -1402,6 +1508,7 @@ def test_import_data(self, import_data_mock, sync): gcs_source=[_TEST_SOURCE_URI_GCS], import_schema_uri=_TEST_IMPORT_SCHEMA_URI_TEXT, sync=sync, + import_request_timeout=None, ) if not sync: @@ -1413,7 +1520,9 @@ def test_import_data(self, import_data_mock, sync): ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1427,12 +1536,14 @@ def test_create_then_import( display_name=_TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) my_dataset.import_data( gcs_source=[_TEST_SOURCE_URI_GCS], import_schema_uri=_TEST_IMPORT_SCHEMA_URI_TEXT, sync=sync, + import_request_timeout=None, ) if not sync: @@ -1448,6 +1559,7 @@ def test_create_then_import( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) get_dataset_text_mock.assert_called_once_with( @@ -1460,7 +1572,9 @@ def test_create_then_import( ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -1478,6 +1592,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1495,6 +1610,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @@ -1529,6 +1645,7 @@ def test_create_dataset(self, create_dataset_mock, sync): my_dataset = datasets.VideoDataset.create( display_name=_TEST_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1545,6 +1662,7 @@ def test_create_dataset(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_dataset_video_mock") @@ -1560,6 +1678,7 @@ def test_create_and_import_dataset( import_schema_uri=_TEST_IMPORT_SCHEMA_URI_VIDEO, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1576,6 +1695,7 @@ def test_create_and_import_dataset( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) expected_import_config = gca_dataset.ImportDataConfig( @@ -1583,7 +1703,9 @@ def test_create_and_import_dataset( import_schema_uri=_TEST_IMPORT_SCHEMA_URI_VIDEO, ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -1600,6 +1722,7 @@ def test_import_data(self, import_data_mock, sync): gcs_source=[_TEST_SOURCE_URI_GCS], import_schema_uri=_TEST_IMPORT_SCHEMA_URI_VIDEO, sync=sync, + import_request_timeout=None, ) if not sync: @@ -1611,7 +1734,9 @@ def test_import_data(self, import_data_mock, sync): ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1625,12 +1750,14 @@ def test_create_then_import( display_name=_TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) my_dataset.import_data( gcs_source=[_TEST_SOURCE_URI_GCS], import_schema_uri=_TEST_IMPORT_SCHEMA_URI_VIDEO, sync=sync, + import_request_timeout=None, ) if not sync: @@ -1646,6 +1773,7 @@ def test_create_then_import( parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) get_dataset_video_mock.assert_called_once_with( @@ -1658,7 +1786,9 @@ def test_create_then_import( ) import_data_mock.assert_called_once_with( - name=_TEST_NAME, import_configs=[expected_import_config] + name=_TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = _TEST_NAME @@ -1675,6 +1805,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1692,4 +1823,5 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): parent=_TEST_PARENT, dataset=expected_dataset, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) diff --git a/tests/unit/aiplatform/test_end_to_end.py b/tests/unit/aiplatform/test_end_to_end.py index e637677fa3..c31b17ab1c 100644 --- a/tests/unit/aiplatform/test_end_to_end.py +++ b/tests/unit/aiplatform/test_end_to_end.py @@ -103,6 +103,7 @@ def test_dataset_create_to_model_predict( display_name=test_datasets._TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) my_dataset.import_data( @@ -110,6 +111,7 @@ def test_dataset_create_to_model_predict( import_schema_uri=test_datasets._TEST_IMPORT_SCHEMA_URI, data_item_labels=test_datasets._TEST_DATA_LABEL_ITEMS, sync=sync, + import_request_timeout=None, ) job = aiplatform.CustomTrainingJob( @@ -134,16 +136,20 @@ def test_dataset_create_to_model_predict( validation_fraction_split=test_training_jobs._TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=test_training_jobs._TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) created_endpoint = models.Endpoint.create( display_name=test_endpoints._TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) my_endpoint = model_from_job.deploy( - encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync + encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + sync=sync, + deploy_request_timeout=None, ) endpoint_deploy_return = created_endpoint.deploy(model_from_job, sync=sync) @@ -194,10 +200,13 @@ def test_dataset_create_to_model_predict( parent=test_datasets._TEST_PARENT, dataset=expected_dataset, metadata=test_datasets._TEST_REQUEST_METADATA, + timeout=None, ) import_data_mock.assert_called_once_with( - name=test_datasets._TEST_NAME, import_configs=[expected_import_config] + name=test_datasets._TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = test_datasets._TEST_NAME @@ -274,6 +283,7 @@ def test_dataset_create_to_model_predict( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -319,6 +329,7 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( my_dataset = aiplatform.ImageDataset.create( display_name=test_datasets._TEST_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) my_dataset.import_data( @@ -326,6 +337,7 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( import_schema_uri=test_datasets._TEST_IMPORT_SCHEMA_URI, data_item_labels=test_datasets._TEST_DATA_LABEL_ITEMS, sync=sync, + import_request_timeout=None, ) job = aiplatform.CustomTrainingJob( @@ -340,6 +352,7 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( created_endpoint = models.Endpoint.create( display_name=test_endpoints._TEST_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) model_from_job = job.run( @@ -355,6 +368,7 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( validation_fraction_split=test_training_jobs._TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=test_training_jobs._TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) with pytest.raises(RuntimeError): @@ -383,10 +397,13 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( parent=test_datasets._TEST_PARENT, dataset=expected_dataset, metadata=test_datasets._TEST_REQUEST_METADATA, + timeout=None, ) import_data_mock.assert_called_once_with( - name=test_datasets._TEST_NAME, import_configs=[expected_import_config] + name=test_datasets._TEST_NAME, + import_configs=[expected_import_config], + timeout=None, ) expected_dataset.name = test_datasets._TEST_NAME @@ -465,6 +482,7 @@ def test_dataset_create_to_model_predict_with_pipeline_fail( mock_pipeline_service_create_and_get_with_fail[0].assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert ( diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index b35fab20a5..a0844fa9b3 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/tests/unit/aiplatform/test_endpoints.py @@ -532,7 +532,11 @@ def test_init_aiplatform_with_encryption_key_name_and_create_endpoint( location=_TEST_LOCATION, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, ) - my_endpoint = models.Endpoint.create(display_name=_TEST_DISPLAY_NAME, sync=sync) + my_endpoint = models.Endpoint.create( + display_name=_TEST_DISPLAY_NAME, + sync=sync, + create_request_timeout=None, + ) if not sync: my_endpoint.wait() @@ -544,6 +548,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_endpoint( parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(), + timeout=None, ) expected_endpoint.name = _TEST_ENDPOINT_NAME @@ -556,6 +561,7 @@ def test_create(self, create_endpoint_mock, sync): display_name=_TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -568,11 +574,35 @@ def test_create(self, create_endpoint_mock, sync): parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(), + timeout=None, ) expected_endpoint.name = _TEST_ENDPOINT_NAME assert my_endpoint._gca_resource == expected_endpoint + @pytest.mark.usefixtures("get_endpoint_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_create_with_timeout(self, create_endpoint_mock, sync): + my_endpoint = models.Endpoint.create( + display_name=_TEST_DISPLAY_NAME, + encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + my_endpoint.wait() + + expected_endpoint = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC + ) + create_endpoint_mock.assert_called_once_with( + parent=_TEST_PARENT, + endpoint=expected_endpoint, + metadata=(), + timeout=180.0, + ) + @pytest.mark.usefixtures("get_empty_endpoint_mock") def test_accessing_properties_with_no_resource_raises( self, @@ -597,7 +627,10 @@ def test_accessing_properties_with_no_resource_raises( @pytest.mark.parametrize("sync", [True, False]) def test_create_with_description(self, create_endpoint_mock, sync): my_endpoint = models.Endpoint.create( - display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, sync=sync + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + sync=sync, + create_request_timeout=None, ) if not sync: my_endpoint.wait() @@ -610,13 +643,17 @@ def test_create_with_description(self, create_endpoint_mock, sync): parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock") @pytest.mark.parametrize("sync", [True, False]) def test_create_with_labels(self, create_endpoint_mock, sync): my_endpoint = models.Endpoint.create( - display_name=_TEST_DISPLAY_NAME, labels=_TEST_LABELS, sync=sync + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + sync=sync, + create_request_timeout=None, ) if not sync: my_endpoint.wait() @@ -629,6 +666,7 @@ def test_create_with_labels(self, create_endpoint_mock, sync): parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -639,7 +677,45 @@ def test_deploy(self, deploy_model_mock, sync): test_model._gca_resource.supported_deployment_resources_types.append( aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES ) - test_endpoint.deploy(test_model, sync=sync) + test_endpoint.deploy( + test_model, + sync=sync, + deploy_request_timeout=None, + ) + + if not sync: + test_endpoint.wait() + + automatic_resources = gca_machine_resources.AutomaticResources( + min_replica_count=1, + max_replica_count=1, + ) + deployed_model = gca_endpoint.DeployedModel( + automatic_resources=automatic_resources, + model=test_model.resource_name, + display_name=None, + ) + deploy_model_mock.assert_called_once_with( + endpoint=test_endpoint.resource_name, + deployed_model=deployed_model, + traffic_split={"0": 100}, + metadata=(), + timeout=None, + ) + + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_deploy_with_timeout(self, deploy_model_mock, sync): + test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + test_model = models.Model(_TEST_ID) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES + ) + test_endpoint.deploy( + test_model, + sync=sync, + deploy_request_timeout=180.0, + ) if not sync: test_endpoint.wait() @@ -658,6 +734,7 @@ def test_deploy(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=180.0, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -669,7 +746,10 @@ def test_deploy_with_display_name(self, deploy_model_mock, sync): aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES ) test_endpoint.deploy( - model=test_model, deployed_model_display_name=_TEST_DISPLAY_NAME, sync=sync + model=test_model, + deployed_model_display_name=_TEST_DISPLAY_NAME, + sync=sync, + deploy_request_timeout=None, ) if not sync: @@ -689,6 +769,7 @@ def test_deploy_with_display_name(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -777,7 +858,12 @@ def test_deploy_with_traffic_percent(self, deploy_model_mock, sync): test_model._gca_resource.supported_deployment_resources_types.append( aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES ) - test_endpoint.deploy(model=test_model, traffic_percentage=70, sync=sync) + test_endpoint.deploy( + model=test_model, + traffic_percentage=70, + sync=sync, + deploy_request_timeout=None, + ) if not sync: test_endpoint.wait() automatic_resources = gca_machine_resources.AutomaticResources( @@ -794,6 +880,7 @@ def test_deploy_with_traffic_percent(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"model1": 30, "0": 70}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_model_mock") @@ -814,7 +901,10 @@ def test_deploy_with_traffic_split(self, deploy_model_mock, sync): aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES ) test_endpoint.deploy( - model=test_model, traffic_split={"model1": 30, "0": 70}, sync=sync + model=test_model, + traffic_split={"model1": 30, "0": 70}, + sync=sync, + deploy_request_timeout=None, ) if not sync: @@ -833,6 +923,7 @@ def test_deploy_with_traffic_split(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"model1": 30, "0": 70}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -850,6 +941,7 @@ def test_deploy_with_dedicated_resources(self, deploy_model_mock, sync): accelerator_count=_TEST_ACCELERATOR_COUNT, service_account=_TEST_SERVICE_ACCOUNT, sync=sync, + deploy_request_timeout=None, ) if not sync: @@ -876,6 +968,7 @@ def test_deploy_with_dedicated_resources(self, deploy_model_mock, sync): deployed_model=expected_deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -894,6 +987,7 @@ def test_deploy_with_explanations(self, deploy_model_with_explanations_mock, syn explanation_metadata=_TEST_EXPLANATION_METADATA, explanation_parameters=_TEST_EXPLANATION_PARAMETERS, sync=sync, + deploy_request_timeout=None, ) if not sync: @@ -923,6 +1017,7 @@ def test_deploy_with_explanations(self, deploy_model_with_explanations_mock, syn deployed_model=expected_deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -933,7 +1028,12 @@ def test_deploy_with_min_replica_count(self, deploy_model_mock, sync): test_model._gca_resource.supported_deployment_resources_types.append( aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES ) - test_endpoint.deploy(model=test_model, min_replica_count=2, sync=sync) + test_endpoint.deploy( + model=test_model, + min_replica_count=2, + sync=sync, + deploy_request_timeout=None, + ) if not sync: test_endpoint.wait() @@ -951,6 +1051,7 @@ def test_deploy_with_min_replica_count(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @@ -961,7 +1062,12 @@ def test_deploy_with_max_replica_count(self, deploy_model_mock, sync): test_model._gca_resource.supported_deployment_resources_types.append( aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES ) - test_endpoint.deploy(model=test_model, max_replica_count=2, sync=sync) + test_endpoint.deploy( + model=test_model, + max_replica_count=2, + sync=sync, + deploy_request_timeout=None, + ) if not sync: test_endpoint.wait() automatic_resources = gca_machine_resources.AutomaticResources( @@ -978,6 +1084,7 @@ def test_deploy_with_max_replica_count(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.parametrize( diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py index 74ab1cf0c5..6d96e5ce45 100644 --- a/tests/unit/aiplatform/test_featurestores.py +++ b/tests/unit/aiplatform/test_featurestores.py @@ -880,7 +880,10 @@ def test_update_featurestore(self, update_featurestore_mock): my_featurestore = aiplatform.Featurestore( featurestore_name=_TEST_FEATURESTORE_ID ) - my_featurestore.update(labels=_TEST_LABELS_UPDATE) + my_featurestore.update( + labels=_TEST_LABELS_UPDATE, + update_request_timeout=None, + ) expected_featurestore = gca_featurestore.Featurestore( name=_TEST_FEATURESTORE_NAME, @@ -891,6 +894,31 @@ def test_update_featurestore(self, update_featurestore_mock): featurestore=expected_featurestore, update_mask=field_mask_pb2.FieldMask(paths=["labels"]), metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.usefixtures("get_featurestore_mock") + def test_update_featurestore_with_timeout(self, update_featurestore_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_ID + ) + my_featurestore.update( + labels=_TEST_LABELS_UPDATE, + update_request_timeout=180.0, + ) + + expected_featurestore = gca_featurestore.Featurestore( + name=_TEST_FEATURESTORE_NAME, + labels=_TEST_LABELS_UPDATE, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig(), + ) + update_featurestore_mock.assert_called_once_with( + featurestore=expected_featurestore, + update_mask=field_mask_pb2.FieldMask(paths=["labels"]), + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.usefixtures("get_featurestore_mock") @@ -901,7 +929,8 @@ def test_update_featurestore_online(self, update_featurestore_mock): featurestore_name=_TEST_FEATURESTORE_ID ) my_featurestore.update_online_store( - fixed_node_count=_TEST_ONLINE_SERVING_CONFIG_UPDATE + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG_UPDATE, + update_request_timeout=None, ) expected_featurestore = gca_featurestore.Featurestore( @@ -916,6 +945,7 @@ def test_update_featurestore_online(self, update_featurestore_mock): paths=["online_serving_config.fixed_node_count"] ), metadata=_TEST_REQUEST_METADATA, + timeout=None, ) def test_list_featurestores(self, list_featurestores_mock): @@ -1022,6 +1052,38 @@ def test_create_entity_type(self, create_entity_type_mock, sync): description=_TEST_DESCRIPTION, labels=_TEST_LABELS, sync=sync, + create_request_timeout=None, + ) + + if not sync: + my_entity_type.wait() + + expected_entity_type = gca_entity_type.EntityType( + labels=_TEST_LABELS, + description=_TEST_DESCRIPTION, + ) + create_entity_type_mock.assert_called_once_with( + parent=_TEST_FEATURESTORE_NAME, + entity_type=expected_entity_type, + entity_type_id=_TEST_ENTITY_TYPE_ID, + metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.usefixtures("get_featurestore_mock", "get_entity_type_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_create_entity_type_with_timeout(self, create_entity_type_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_NAME + ) + my_entity_type = my_featurestore.create_entity_type( + entity_type_id=_TEST_ENTITY_TYPE_ID, + description=_TEST_DESCRIPTION, + labels=_TEST_LABELS, + sync=sync, + create_request_timeout=180.0, ) if not sync: @@ -1036,6 +1098,7 @@ def test_create_entity_type(self, create_entity_type_mock, sync): entity_type=expected_entity_type, entity_type_id=_TEST_ENTITY_TYPE_ID, metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.usefixtures("get_featurestore_mock") @@ -1048,6 +1111,7 @@ def test_create_featurestore(self, create_featurestore_mock, sync): online_store_fixed_node_count=_TEST_ONLINE_SERVING_CONFIG, labels=_TEST_LABELS, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + create_request_timeout=None, ) if not sync: @@ -1065,6 +1129,38 @@ def test_create_featurestore(self, create_featurestore_mock, sync): featurestore=expected_featurestore, featurestore_id=_TEST_FEATURESTORE_ID, metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.usefixtures("get_featurestore_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_create_featurestore_with_timeout(self, create_featurestore_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_featurestore = aiplatform.Featurestore.create( + featurestore_id=_TEST_FEATURESTORE_ID, + online_store_fixed_node_count=_TEST_ONLINE_SERVING_CONFIG, + labels=_TEST_LABELS, + encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + create_request_timeout=180.0, + ) + + if not sync: + my_featurestore.wait() + + expected_featurestore = gca_featurestore.Featurestore( + labels=_TEST_LABELS, + online_serving_config=gca_featurestore.Featurestore.OnlineServingConfig( + fixed_node_count=_TEST_ONLINE_SERVING_CONFIG + ), + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + create_featurestore_mock.assert_called_once_with( + parent=_TEST_PARENT, + featurestore=expected_featurestore, + featurestore_id=_TEST_FEATURESTORE_ID, + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.usefixtures("get_featurestore_mock") @@ -1248,6 +1344,54 @@ def test_batch_serve_to_bq(self, batch_read_feature_values_mock, sync): serving_feature_ids=_TEST_SERVING_FEATURE_IDS, read_instances_uri=_TEST_BQ_SOURCE_URI, sync=sync, + serve_request_timeout=None, + ) + + if not sync: + my_featurestore.wait() + + batch_read_feature_values_mock.assert_called_once_with( + request=expected_batch_read_feature_values_request, + metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_featurestore_mock") + def test_batch_serve_to_bq_with_timeout(self, batch_read_feature_values_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + my_featurestore = aiplatform.Featurestore( + featurestore_name=_TEST_FEATURESTORE_NAME + ) + + expected_entity_type_specs = [ + _get_entity_type_spec_proto_with_feature_ids( + entity_type_id="my_entity_type_id_1", + feature_ids=["my_feature_id_1_1", "my_feature_id_1_2"], + ), + _get_entity_type_spec_proto_with_feature_ids( + entity_type_id="my_entity_type_id_2", + feature_ids=["my_feature_id_2_1", "my_feature_id_2_2"], + ), + ] + + expected_batch_read_feature_values_request = ( + gca_featurestore_service.BatchReadFeatureValuesRequest( + featurestore=my_featurestore.resource_name, + destination=gca_featurestore_service.FeatureValueDestination( + bigquery_destination=_TEST_BQ_DESTINATION, + ), + entity_type_specs=expected_entity_type_specs, + bigquery_read_instances=_TEST_BQ_SOURCE, + ) + ) + + my_featurestore.batch_serve_to_bq( + bq_destination_output_uri=_TEST_BQ_DESTINATION_URI, + serving_feature_ids=_TEST_SERVING_FEATURE_IDS, + read_instances_uri=_TEST_BQ_SOURCE_URI, + sync=sync, + serve_request_timeout=180.0, ) if not sync: @@ -1256,6 +1400,7 @@ def test_batch_serve_to_bq(self, batch_read_feature_values_mock, sync): batch_read_feature_values_mock.assert_called_once_with( request=expected_batch_read_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1294,6 +1439,7 @@ def test_batch_serve_to_gcs(self, batch_read_feature_values_mock, sync): serving_feature_ids=_TEST_SERVING_FEATURE_IDS, read_instances_uri=_TEST_GCS_CSV_SOURCE_URI, sync=sync, + serve_request_timeout=None, ) if not sync: @@ -1302,6 +1448,7 @@ def test_batch_serve_to_gcs(self, batch_read_feature_values_mock, sync): batch_read_feature_values_mock.assert_called_once_with( request=expected_batch_read_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_featurestore_mock") @@ -1384,11 +1531,13 @@ def test_batch_serve_to_df(self, batch_read_feature_values_mock): my_featurestore.batch_serve_to_df( serving_feature_ids=_TEST_SERVING_FEATURE_IDS, read_instances_df=read_instances_df, + serve_request_timeout=None, ) batch_read_feature_values_mock.assert_called_once_with( request=expected_batch_read_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @@ -1449,7 +1598,10 @@ def test_update_entity_type(self, update_entity_type_mock): aiplatform.init(project=_TEST_PROJECT) my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) - my_entity_type.update(labels=_TEST_LABELS_UPDATE) + my_entity_type.update( + labels=_TEST_LABELS_UPDATE, + update_request_timeout=None, + ) expected_entity_type = gca_entity_type.EntityType( name=_TEST_ENTITY_TYPE_NAME, @@ -1459,6 +1611,7 @@ def test_update_entity_type(self, update_entity_type_mock): entity_type=expected_entity_type, update_mask=field_mask_pb2.FieldMask(paths=["labels"]), metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.parametrize( @@ -1524,6 +1677,7 @@ def test_create_feature(self, create_feature_mock, sync): value_type=_TEST_FEATURE_VALUE_TYPE_STR, description=_TEST_DESCRIPTION, labels=_TEST_LABELS, + create_request_timeout=None, ) if not sync: @@ -1543,6 +1697,7 @@ def test_create_feature(self, create_feature_mock, sync): create_feature_mock.assert_called_once_with( request=expected_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_entity_type_mock") @@ -1555,6 +1710,7 @@ def test_create_entity_type(self, create_entity_type_mock, sync): featurestore_name=_TEST_FEATURESTORE_NAME, description=_TEST_DESCRIPTION, labels=_TEST_LABELS, + create_request_timeout=None, ) if not sync: @@ -1569,6 +1725,7 @@ def test_create_entity_type(self, create_entity_type_mock, sync): entity_type=expected_entity_type, entity_type_id=_TEST_ENTITY_TYPE_ID, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_entity_type_mock") @@ -1683,6 +1840,44 @@ def test_ingest_from_bq(self, import_feature_values_mock, sync): bq_source_uri=_TEST_BQ_SOURCE_URI, feature_source_fields=_TEST_IMPORTING_FEATURE_SOURCE_FIELDS, sync=sync, + ingest_request_timeout=None, + ) + + if not sync: + my_entity_type.wait() + + true_import_feature_values_request = ( + gca_featurestore_service.ImportFeatureValuesRequest( + entity_type=_TEST_ENTITY_TYPE_NAME, + feature_specs=[ + gca_featurestore_service.ImportFeatureValuesRequest.FeatureSpec( + id="my_feature_id_1", + source_field="my_feature_id_1_source_field", + ), + ], + bigquery_source=_TEST_BQ_SOURCE, + feature_time_field=_TEST_FEATURE_TIME_FIELD, + ) + ) + import_feature_values_mock.assert_called_once_with( + request=true_import_feature_values_request, + metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.usefixtures("get_entity_type_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_ingest_from_bq_with_timeout(self, import_feature_values_mock, sync): + aiplatform.init(project=_TEST_PROJECT) + + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + my_entity_type.ingest_from_bq( + feature_ids=_TEST_IMPORTING_FEATURE_IDS, + feature_time=_TEST_FEATURE_TIME_FIELD, + bq_source_uri=_TEST_BQ_SOURCE_URI, + feature_source_fields=_TEST_IMPORTING_FEATURE_SOURCE_FIELDS, + sync=sync, + ingest_request_timeout=180.0, ) if not sync: @@ -1704,6 +1899,7 @@ def test_ingest_from_bq(self, import_feature_values_mock, sync): import_feature_values_mock.assert_called_once_with( request=true_import_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.usefixtures("get_entity_type_mock") @@ -1718,6 +1914,7 @@ def test_ingest_from_gcs(self, import_feature_values_mock, sync): gcs_source_uris=_TEST_GCS_AVRO_SOURCE_URIS, gcs_source_type=_TEST_GCS_SOURCE_TYPE_AVRO, sync=sync, + ingest_request_timeout=None, ) if not sync: @@ -1738,6 +1935,7 @@ def test_ingest_from_gcs(self, import_feature_values_mock, sync): import_feature_values_mock.assert_called_once_with( request=true_import_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_entity_type_mock") @@ -1779,6 +1977,7 @@ def test_ingest_from_df_using_column( feature_time=_TEST_FEATURE_TIME_FIELD, df_source=df_source, feature_source_fields=_TEST_IMPORTING_FEATURE_SOURCE_FIELDS, + ingest_request_timeout=None, ) expected_temp_bq_dataset_name = ( f"temp_{_TEST_FEATURESTORE_ID}_{uuid.uuid4()}".replace("-", "_") @@ -1815,6 +2014,7 @@ def test_ingest_from_df_using_column( import_feature_values_mock.assert_called_once_with( request=expected_import_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures( @@ -1843,6 +2043,7 @@ def test_ingest_from_df_using_datetime( feature_time=_TEST_FEATURE_TIME_DATETIME, df_source=df_source, feature_source_fields=_TEST_IMPORTING_FEATURE_SOURCE_FIELDS, + ingest_request_timeout=None, ) expected_temp_bq_dataset_name = ( @@ -1883,6 +2084,7 @@ def test_ingest_from_df_using_datetime( import_feature_values_mock.assert_called_once_with( request=expected_import_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.parametrize( @@ -1924,16 +2126,43 @@ def test_read_single_entity(self, read_feature_values_mock): ), ) ) - result = my_entity_type.read(entity_ids=_TEST_READ_ENTITY_ID) + result = my_entity_type.read( + entity_ids=_TEST_READ_ENTITY_ID, + read_request_timeout=None, + ) read_feature_values_mock.assert_called_once_with( request=expected_read_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) assert type(result) == pd.DataFrame assert len(result) == 1 assert result.entity_id[0] == _TEST_READ_ENTITY_ID assert result.get(_TEST_FEATURE_ID)[0] == _TEST_FEATURE_VALUE + @pytest.mark.usefixtures("get_entity_type_mock", "get_feature_mock") + def test_read_single_entity_with_timeout(self, read_feature_values_mock): + aiplatform.init(project=_TEST_PROJECT) + my_entity_type = aiplatform.EntityType(entity_type_name=_TEST_ENTITY_TYPE_NAME) + expected_read_feature_values_request = ( + gca_featurestore_online_service.ReadFeatureValuesRequest( + entity_type=my_entity_type.resource_name, + entity_id=_TEST_READ_ENTITY_ID, + feature_selector=gca_feature_selector.FeatureSelector( + id_matcher=gca_feature_selector.IdMatcher(ids=["*"]) + ), + ) + ) + my_entity_type.read( + entity_ids=_TEST_READ_ENTITY_ID, + read_request_timeout=180.0, + ) + read_feature_values_mock.assert_called_once_with( + request=expected_read_feature_values_request, + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, + ) + @pytest.mark.usefixtures("get_entity_type_mock", "get_feature_mock") def test_read_multiple_entities(self, streaming_read_feature_values_mock): aiplatform.init(project=_TEST_PROJECT) @@ -1948,11 +2177,14 @@ def test_read_multiple_entities(self, streaming_read_feature_values_mock): ) ) result = my_entity_type.read( - entity_ids=_TEST_READ_ENTITY_IDS, feature_ids=_TEST_FEATURE_ID + entity_ids=_TEST_READ_ENTITY_IDS, + feature_ids=_TEST_FEATURE_ID, + read_request_timeout=None, ) streaming_read_feature_values_mock.assert_called_once_with( request=expected_streaming_read_feature_values_request, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) assert type(result) == pd.DataFrame assert len(result) == 1 @@ -2443,7 +2675,10 @@ def test_update_feature(self, update_feature_mock): aiplatform.init(project=_TEST_PROJECT) my_feature = aiplatform.Feature(feature_name=_TEST_FEATURE_NAME) - my_feature.update(labels=_TEST_LABELS_UPDATE) + my_feature.update( + labels=_TEST_LABELS_UPDATE, + update_request_timeout=None, + ) expected_feature = gca_feature.Feature( name=_TEST_FEATURE_NAME, @@ -2453,6 +2688,7 @@ def test_update_feature(self, update_feature_mock): feature=expected_feature, update_mask=field_mask_pb2.FieldMask(paths=["labels"]), metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.parametrize( @@ -2501,6 +2737,7 @@ def test_create_feature(self, create_feature_mock, sync): featurestore_id=_TEST_FEATURESTORE_ID, description=_TEST_DESCRIPTION, labels=_TEST_LABELS, + create_request_timeout=None, ) if not sync: @@ -2518,6 +2755,7 @@ def test_create_feature(self, create_feature_mock, sync): feature_id=_TEST_FEATURE_ID, ), metadata=_TEST_REQUEST_METADATA, + timeout=None, ) diff --git a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py index 4a41dde3b9..727f106fb5 100644 --- a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py +++ b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py @@ -404,6 +404,7 @@ def test_create_hyperparameter_tuning_job( timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) job.wait() @@ -413,12 +414,75 @@ def test_create_hyperparameter_tuning_job( create_hyperparameter_tuning_job_mock.assert_called_once_with( parent=_TEST_PARENT, hyperparameter_tuning_job=expected_hyperparameter_tuning_job, + timeout=None, ) assert job.state == gca_job_state_compat.JobState.JOB_STATE_SUCCEEDED assert job.network == _TEST_NETWORK assert job.trials == [] + @pytest.mark.parametrize("sync", [True, False]) + def test_create_hyperparameter_tuning_job_with_timeout( + self, + create_hyperparameter_tuning_job_mock, + get_hyperparameter_tuning_job_mock, + sync, + ): + + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + staging_bucket=_TEST_STAGING_BUCKET, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + custom_job = aiplatform.CustomJob( + display_name=test_custom_job._TEST_DISPLAY_NAME, + worker_pool_specs=test_custom_job._TEST_WORKER_POOL_SPEC, + base_output_dir=test_custom_job._TEST_BASE_OUTPUT_DIR, + ) + + job = aiplatform.HyperparameterTuningJob( + display_name=_TEST_DISPLAY_NAME, + custom_job=custom_job, + metric_spec={_TEST_METRIC_SPEC_KEY: _TEST_METRIC_SPEC_VALUE}, + parameter_spec={ + "lr": hpt.DoubleParameterSpec(min=0.001, max=0.1, scale="log"), + "units": hpt.IntegerParameterSpec(min=4, max=1028, scale="linear"), + "activation": hpt.CategoricalParameterSpec( + values=["relu", "sigmoid", "elu", "selu", "tanh"] + ), + "batch_size": hpt.DiscreteParameterSpec( + values=[16, 32], scale="linear" + ), + }, + parallel_trial_count=_TEST_PARALLEL_TRIAL_COUNT, + max_trial_count=_TEST_MAX_TRIAL_COUNT, + max_failed_trial_count=_TEST_MAX_FAILED_TRIAL_COUNT, + search_algorithm=_TEST_SEARCH_ALGORITHM, + measurement_selection=_TEST_MEASUREMENT_SELECTION, + labels=_TEST_LABELS, + ) + + job.run( + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + timeout=_TEST_TIMEOUT, + restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, + sync=sync, + create_request_timeout=180.0, + ) + + job.wait() + + expected_hyperparameter_tuning_job = _get_hyperparameter_tuning_job_proto() + + create_hyperparameter_tuning_job_mock.assert_called_once_with( + parent=_TEST_PARENT, + hyperparameter_tuning_job=expected_hyperparameter_tuning_job, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_run_hyperparameter_tuning_job_with_fail_raises( self, @@ -468,6 +532,7 @@ def test_run_hyperparameter_tuning_job_with_fail_raises( timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) job.wait() @@ -477,6 +542,7 @@ def test_run_hyperparameter_tuning_job_with_fail_raises( create_hyperparameter_tuning_job_mock.assert_called_once_with( parent=_TEST_PARENT, hyperparameter_tuning_job=expected_hyperparameter_tuning_job, + timeout=None, ) assert job._gca_resource.state == gca_job_state_compat.JobState.JOB_STATE_FAILED @@ -648,6 +714,7 @@ def test_create_hyperparameter_tuning_job_with_tensorboard( restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, tensorboard=test_custom_job._TEST_TENSORBOARD_NAME, sync=sync, + create_request_timeout=None, ) job.wait() @@ -660,6 +727,7 @@ def test_create_hyperparameter_tuning_job_with_tensorboard( create_hyperparameter_tuning_job_mock_with_tensorboard.assert_called_once_with( parent=_TEST_PARENT, hyperparameter_tuning_job=expected_hyperparameter_tuning_job, + timeout=None, ) assert ( @@ -718,6 +786,7 @@ def test_create_hyperparameter_tuning_job_with_enable_web_access( restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, enable_web_access=test_custom_job._TEST_ENABLE_WEB_ACCESS, sync=sync, + create_request_timeout=None, ) job.wait() @@ -731,6 +800,7 @@ def test_create_hyperparameter_tuning_job_with_enable_web_access( create_hyperparameter_tuning_job_mock_with_enable_web_access.assert_called_once_with( parent=_TEST_PARENT, hyperparameter_tuning_job=expected_hyperparameter_tuning_job, + timeout=None, ) assert job.state == gca_job_state_compat.JobState.JOB_STATE_SUCCEEDED diff --git a/tests/unit/aiplatform/test_jobs.py b/tests/unit/aiplatform/test_jobs.py index 862e2c7efc..6b52fb7902 100644 --- a/tests/unit/aiplatform/test_jobs.py +++ b/tests/unit/aiplatform/test_jobs.py @@ -508,6 +508,7 @@ def test_batch_predict_gcs_source_and_dest( gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, sync=sync, + create_request_timeout=None, ) batch_prediction_job.wait_for_resource_creation() @@ -535,6 +536,52 @@ def test_batch_predict_gcs_source_and_dest( create_batch_prediction_job_mock.assert_called_once_with( parent=_TEST_PARENT, batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, + ) + + @pytest.mark.parametrize("sync", [True, False]) + @pytest.mark.usefixtures("get_batch_prediction_job_mock") + def test_batch_predict_gcs_source_and_dest_with_timeout( + self, create_batch_prediction_job_mock, sync + ): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + # Make SDK batch_predict method call + batch_prediction_job = jobs.BatchPredictionJob.create( + model_name=_TEST_MODEL_NAME, + job_display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, + gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, + sync=sync, + create_request_timeout=180.0, + ) + + batch_prediction_job.wait_for_resource_creation() + + batch_prediction_job.wait() + + # Construct expected request + expected_gapic_batch_prediction_job = gca_batch_prediction_job_compat.BatchPredictionJob( + display_name=_TEST_BATCH_PREDICTION_JOB_DISPLAY_NAME, + model=_TEST_MODEL_NAME, + input_config=gca_batch_prediction_job_compat.BatchPredictionJob.InputConfig( + instances_format="jsonl", + gcs_source=gca_io_compat.GcsSource( + uris=[_TEST_BATCH_PREDICTION_GCS_SOURCE] + ), + ), + output_config=gca_batch_prediction_job_compat.BatchPredictionJob.OutputConfig( + gcs_destination=gca_io_compat.GcsDestination( + output_uri_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX + ), + predictions_format="jsonl", + ), + ) + + create_batch_prediction_job_mock.assert_called_once_with( + parent=_TEST_PARENT, + batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=180.0, ) @pytest.mark.usefixtures("get_batch_prediction_job_mock") @@ -571,6 +618,7 @@ def test_batch_predict_gcs_source_bq_dest( gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, bigquery_destination_prefix=_TEST_BATCH_PREDICTION_BQ_PREFIX, sync=sync, + create_request_timeout=None, ) batch_prediction_job.wait_for_resource_creation() @@ -603,6 +651,7 @@ def test_batch_predict_gcs_source_bq_dest( create_batch_prediction_job_mock.assert_called_once_with( parent=_TEST_PARENT, batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -631,6 +680,7 @@ def test_batch_predict_with_all_args( labels=_TEST_LABEL, credentials=creds, sync=sync, + create_request_timeout=None, ) batch_prediction_job.wait_for_resource_creation() @@ -673,6 +723,7 @@ def test_batch_predict_with_all_args( create_batch_prediction_job_with_explanations_mock.assert_called_once_with( parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, ) @pytest.mark.usefixtures("create_batch_prediction_job_mock_fail") diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index 4a82b7448f..88d9be865f 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/tests/unit/aiplatform/test_models.py @@ -551,6 +551,7 @@ def test_upload_uploads_and_gets_model( serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, sync=sync, + upload_request_timeout=None, ) if not sync: @@ -570,12 +571,40 @@ def test_upload_uploads_and_gets_model( upload_model_mock.assert_called_once_with( parent=initializer.global_config.common_location_path(), model=managed_model, + timeout=None, ) get_model_mock.assert_called_once_with( name=_TEST_MODEL_RESOURCE_NAME, retry=base._DEFAULT_RETRY ) + @pytest.mark.parametrize("sync", [True, False]) + def test_upload_with_timeout(self, upload_model_mock, get_model_mock, sync): + my_model = models.Model.upload( + display_name=_TEST_MODEL_NAME, + serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE, + sync=sync, + upload_request_timeout=180.0, + ) + + if not sync: + my_model.wait() + + container_spec = gca_model.ModelContainerSpec( + image_uri=_TEST_SERVING_CONTAINER_IMAGE, + ) + + managed_model = gca_model.Model( + display_name=_TEST_MODEL_NAME, + container_spec=container_spec, + ) + + upload_model_mock.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + model=managed_model, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_upload_uploads_and_gets_model_with_labels( self, upload_model_mock, get_model_mock, sync @@ -587,6 +616,7 @@ def test_upload_uploads_and_gets_model_with_labels( serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, labels=_TEST_LABEL, + upload_request_timeout=None, sync=sync, ) @@ -608,6 +638,7 @@ def test_upload_uploads_and_gets_model_with_labels( upload_model_mock.assert_called_once_with( parent=initializer.global_config.common_location_path(), model=managed_model, + timeout=None, ) get_model_mock.assert_called_once_with( @@ -650,6 +681,7 @@ def test_upload_uploads_and_gets_model_with_all_args( explanation_parameters=_TEST_EXPLANATION_PARAMETERS, labels=_TEST_LABEL, sync=sync, + upload_request_timeout=None, ) if not sync: @@ -695,6 +727,7 @@ def test_upload_uploads_and_gets_model_with_all_args( upload_model_mock.assert_called_once_with( parent=initializer.global_config.common_location_path(), model=managed_model, + timeout=None, ) get_model_mock.assert_called_once_with( name=_TEST_MODEL_RESOURCE_NAME, retry=base._DEFAULT_RETRY @@ -721,6 +754,7 @@ def test_upload_uploads_and_gets_model_with_custom_project( serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, project=_TEST_PROJECT_2, sync=sync, + upload_request_timeout=None, ) if not sync: @@ -741,6 +775,7 @@ def test_upload_uploads_and_gets_model_with_custom_project( upload_model_with_custom_project_mock.assert_called_once_with( parent=f"projects/{_TEST_PROJECT_2}/locations/{_TEST_LOCATION}", model=managed_model, + timeout=None, ) get_model_with_custom_project_mock.assert_called_once_with( @@ -810,6 +845,7 @@ def test_upload_uploads_and_gets_model_with_custom_location( serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, location=_TEST_LOCATION_2, sync=sync, + upload_request_timeout=None, ) if not sync: @@ -830,6 +866,7 @@ def test_upload_uploads_and_gets_model_with_custom_location( upload_model_with_custom_location_mock.assert_called_once_with( parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION_2}", model=managed_model, + timeout=None, ) get_model_with_custom_location_mock.assert_called_once_with( @@ -872,6 +909,42 @@ def test_deploy(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, + ) + + @pytest.mark.usefixtures( + "get_endpoint_mock", "get_model_mock", "create_endpoint_mock" + ) + @pytest.mark.parametrize("sync", [True, False]) + def test_deploy_with_timeout(self, deploy_model_mock, sync): + + test_model = models.Model(_TEST_ID) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES + ) + + test_endpoint = models.Endpoint(_TEST_ID) + + test_model.deploy(test_endpoint, sync=sync, deploy_request_timeout=180.0) + + if not sync: + test_endpoint.wait() + + automatic_resources = gca_machine_resources.AutomaticResources( + min_replica_count=1, + max_replica_count=1, + ) + deployed_model = gca_endpoint.DeployedModel( + automatic_resources=automatic_resources, + model=test_model.resource_name, + display_name=None, + ) + deploy_model_mock.assert_called_once_with( + endpoint=test_endpoint.resource_name, + deployed_model=deployed_model, + traffic_split={"0": 100}, + metadata=(), + timeout=180.0, ) @pytest.mark.usefixtures( @@ -903,6 +976,7 @@ def test_deploy_no_endpoint(self, deploy_model_mock, sync): deployed_model=deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures( @@ -921,6 +995,7 @@ def test_deploy_no_endpoint_dedicated_resources(self, deploy_model_mock, sync): accelerator_count=_TEST_ACCELERATOR_COUNT, service_account=_TEST_SERVICE_ACCOUNT, sync=sync, + deploy_request_timeout=None, ) if not sync: @@ -945,6 +1020,7 @@ def test_deploy_no_endpoint_dedicated_resources(self, deploy_model_mock, sync): deployed_model=expected_deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures( @@ -963,6 +1039,7 @@ def test_deploy_no_endpoint_with_explanations(self, deploy_model_mock, sync): explanation_metadata=_TEST_EXPLANATION_METADATA, explanation_parameters=_TEST_EXPLANATION_PARAMETERS, sync=sync, + deploy_request_timeout=None, ) if not sync: @@ -990,6 +1067,7 @@ def test_deploy_no_endpoint_with_explanations(self, deploy_model_mock, sync): deployed_model=expected_deployed_model, traffic_split={"0": 100}, metadata=(), + timeout=None, ) @pytest.mark.usefixtures( @@ -1031,6 +1109,7 @@ def test_init_aiplatform_with_encryption_key_name_and_batch_predict_gcs_source_a gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1062,6 +1141,7 @@ def test_init_aiplatform_with_encryption_key_name_and_batch_predict_gcs_source_a create_batch_prediction_job_mock.assert_called_once_with( parent=_TEST_PARENT, batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1078,6 +1158,7 @@ def test_batch_predict_gcs_source_and_dest( gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, gcs_destination_prefix=_TEST_BATCH_PREDICTION_GCS_DEST_PREFIX, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1108,6 +1189,7 @@ def test_batch_predict_gcs_source_and_dest( create_batch_prediction_job_mock.assert_called_once_with( parent=_TEST_PARENT, batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1124,6 +1206,7 @@ def test_batch_predict_gcs_source_bq_dest( gcs_source=_TEST_BATCH_PREDICTION_GCS_SOURCE, bigquery_destination_prefix=_TEST_BATCH_PREDICTION_BQ_PREFIX, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1154,6 +1237,7 @@ def test_batch_predict_gcs_source_bq_dest( create_batch_prediction_job_mock.assert_called_once_with( parent=_TEST_PARENT, batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1181,6 +1265,7 @@ def test_batch_predict_with_all_args(self, create_batch_prediction_job_mock, syn credentials=creds, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1227,6 +1312,7 @@ def test_batch_predict_with_all_args(self, create_batch_prediction_job_mock, syn create_batch_prediction_job_mock.assert_called_once_with( parent=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}", batch_prediction_job=expected_gapic_batch_prediction_job, + timeout=None, ) @pytest.mark.usefixtures("get_model_mock", "get_batch_prediction_job_mock") @@ -1555,6 +1641,7 @@ def test_upload_xgboost_model_file_uploads_and_gets_model( project=_TEST_PROJECT, location=_TEST_LOCATION, sync=sync, + upload_request_timeout=None, ) if not sync: @@ -1665,6 +1752,7 @@ def test_upload_scikit_learn_model_file_uploads_and_gets_model( project=_TEST_PROJECT, location=_TEST_LOCATION, sync=sync, + upload_request_timeout=None, ) if not sync: @@ -1714,6 +1802,7 @@ def test_upload_tensorflow_saved_model_uploads_and_gets_model( project=_TEST_PROJECT, location=_TEST_LOCATION, sync=sync, + upload_request_timeout=None, ) if not sync: diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 561126db12..3188360276 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -282,6 +282,7 @@ def test_run_call_pipeline_service_create( service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, sync=sync, + create_request_timeout=None, ) if not sync: @@ -314,6 +315,7 @@ def test_run_call_pipeline_service_create( parent=_TEST_PARENT, pipeline_job=expected_gapic_pipeline_job, pipeline_job_id=_TEST_PIPELINE_JOB_ID, + timeout=None, ) mock_pipeline_service_get.assert_called_with( @@ -324,6 +326,82 @@ def test_run_call_pipeline_service_create( gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED ) + @pytest.mark.parametrize( + "job_spec_json", + [_TEST_PIPELINE_SPEC, _TEST_PIPELINE_JOB], + ) + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + job_spec_json, + mock_load_json, + sync, + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_GCS_BUCKET_NAME, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + + job = pipeline_jobs.PipelineJob( + display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, + template_path=_TEST_TEMPLATE_PATH, + job_id=_TEST_PIPELINE_JOB_ID, + parameter_values=_TEST_PIPELINE_PARAMETER_VALUES, + enable_caching=True, + ) + + job.run( + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + job.wait() + + expected_runtime_config_dict = { + "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, + "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES, + } + runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + json_format.ParseDict(expected_runtime_config_dict, runtime_config) + + pipeline_spec = job_spec_json.get("pipelineSpec") or job_spec_json + + # Construct expected request + expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, + pipeline_spec={ + "components": {}, + "pipelineInfo": pipeline_spec["pipelineInfo"], + "root": pipeline_spec["root"], + "schemaVersion": "2.1.0", + }, + runtime_config=runtime_config, + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=_TEST_PARENT, + pipeline_job=expected_gapic_pipeline_job, + pipeline_job_id=_TEST_PIPELINE_JOB_ID, + timeout=180.0, + ) + + # mock_pipeline_service_get.assert_called_with( + # name=_TEST_PIPELINE_JOB_NAME, retry=base._DEFAULT_RETRY + # ) + + # assert job._gca_resource == make_pipeline_job( + # gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + # ) + @pytest.mark.parametrize( "job_spec_json", [_TEST_PIPELINE_SPEC_LEGACY, _TEST_PIPELINE_JOB_LEGACY], @@ -356,6 +434,7 @@ def test_run_call_pipeline_service_create_legacy( service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, sync=sync, + create_request_timeout=None, ) if not sync: @@ -388,6 +467,7 @@ def test_run_call_pipeline_service_create_legacy( parent=_TEST_PARENT, pipeline_job=expected_gapic_pipeline_job, pipeline_job_id=_TEST_PIPELINE_JOB_ID, + timeout=None, ) mock_pipeline_service_get.assert_called_with( @@ -430,6 +510,7 @@ def test_run_call_pipeline_service_create_tfx( service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, sync=sync, + create_request_timeout=None, ) if not sync: @@ -463,6 +544,7 @@ def test_run_call_pipeline_service_create_tfx( parent=_TEST_PARENT, pipeline_job=expected_gapic_pipeline_job, pipeline_job_id=_TEST_PIPELINE_JOB_ID, + timeout=None, ) mock_pipeline_service_get.assert_called_with( @@ -499,7 +581,11 @@ def test_submit_call_pipeline_service_pipeline_job_create( enable_caching=True, ) - job.submit(service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK) + job.submit( + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + create_request_timeout=None, + ) expected_runtime_config_dict = { "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, @@ -528,6 +614,7 @@ def test_submit_call_pipeline_service_pipeline_job_create( parent=_TEST_PARENT, pipeline_job=expected_gapic_pipeline_job, pipeline_job_id=_TEST_PIPELINE_JOB_ID, + timeout=None, ) assert not mock_pipeline_service_get.called @@ -602,7 +689,11 @@ def test_submit_call_pipeline_service_pipeline_job_create_legacy( enable_caching=True, ) - job.submit(service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK) + job.submit( + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + create_request_timeout=None, + ) expected_runtime_config_dict = { "parameters": {"string_param": {"stringValue": "hello"}}, @@ -631,6 +722,7 @@ def test_submit_call_pipeline_service_pipeline_job_create_legacy( parent=_TEST_PARENT, pipeline_job=expected_gapic_pipeline_job, pipeline_job_id=_TEST_PIPELINE_JOB_ID, + timeout=None, ) assert not mock_pipeline_service_get.called diff --git a/tests/unit/aiplatform/test_tensorboard.py b/tests/unit/aiplatform/test_tensorboard.py index 99ec1d2348..f7df67e56e 100644 --- a/tests/unit/aiplatform/test_tensorboard.py +++ b/tests/unit/aiplatform/test_tensorboard.py @@ -364,6 +364,7 @@ def test_create_tensorboard_with_default_encryption_key( tensorboard.Tensorboard.create( display_name=_TEST_DISPLAY_NAME, + create_request_timeout=None, ) expected_tensorboard = gca_tensorboard.Tensorboard( @@ -375,6 +376,7 @@ def test_create_tensorboard_with_default_encryption_key( parent=_TEST_PARENT, tensorboard=expected_tensorboard, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) @pytest.mark.usefixtures("get_tensorboard_mock") @@ -387,6 +389,7 @@ def test_create_tensorboard(self, create_tensorboard_mock): tensorboard.Tensorboard.create( display_name=_TEST_DISPLAY_NAME, encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + create_request_timeout=None, ) expected_tensorboard = gca_tensorboard.Tensorboard( @@ -398,6 +401,32 @@ def test_create_tensorboard(self, create_tensorboard_mock): parent=_TEST_PARENT, tensorboard=expected_tensorboard, metadata=_TEST_REQUEST_METADATA, + timeout=None, + ) + + @pytest.mark.usefixtures("get_tensorboard_mock") + def test_create_tensorboard_with_timeout(self, create_tensorboard_mock): + + aiplatform.init( + project=_TEST_PROJECT, + ) + + tensorboard.Tensorboard.create( + display_name=_TEST_DISPLAY_NAME, + encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME, + create_request_timeout=180.0, + ) + + expected_tensorboard = gca_tensorboard.Tensorboard( + display_name=_TEST_DISPLAY_NAME, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + + create_tensorboard_mock.assert_called_once_with( + parent=_TEST_PARENT, + tensorboard=expected_tensorboard, + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, ) @pytest.mark.usefixtures("get_tensorboard_mock") @@ -502,6 +531,7 @@ def test_create_tensorboard_experiment( tensorboard_experiment_id=_TEST_TENSORBOARD_EXPERIMENT_ID, tensorboard_name=_TEST_NAME, display_name=_TEST_DISPLAY_NAME, + create_request_timeout=None, ) expected_tensorboard_experiment = ( @@ -515,12 +545,42 @@ def test_create_tensorboard_experiment( tensorboard_experiment=expected_tensorboard_experiment, tensorboard_experiment_id=_TEST_TENSORBOARD_EXPERIMENT_ID, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) get_tensorboard_experiment_mock.assert_called_once_with( name=_TEST_TENSORBOARD_EXPERIMENT_NAME, retry=base._DEFAULT_RETRY ) + def test_create_tensorboard_experiment_with_timeout( + self, create_tensorboard_experiment_mock, get_tensorboard_experiment_mock + ): + + aiplatform.init( + project=_TEST_PROJECT, + ) + + tensorboard.TensorboardExperiment.create( + tensorboard_experiment_id=_TEST_TENSORBOARD_EXPERIMENT_ID, + tensorboard_name=_TEST_NAME, + display_name=_TEST_DISPLAY_NAME, + create_request_timeout=180.0, + ) + + expected_tensorboard_experiment = ( + gca_tensorboard_experiment.TensorboardExperiment( + display_name=_TEST_DISPLAY_NAME, + ) + ) + + create_tensorboard_experiment_mock.assert_called_once_with( + parent=_TEST_NAME, + tensorboard_experiment=expected_tensorboard_experiment, + tensorboard_experiment_id=_TEST_TENSORBOARD_EXPERIMENT_ID, + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, + ) + @pytest.mark.usefixtures("get_tensorboard_experiment_mock") def test_delete_tensorboard_experiement(self, delete_tensorboard_experiment_mock): aiplatform.init(project=_TEST_PROJECT) @@ -599,6 +659,7 @@ def test_create_tensorboard_run( tensorboard.TensorboardRun.create( tensorboard_run_id=_TEST_TENSORBOARD_RUN_ID, tensorboard_experiment_name=_TEST_TENSORBOARD_EXPERIMENT_NAME, + create_request_timeout=None, ) expected_tensorboard_run = gca_tensorboard_run.TensorboardRun( @@ -610,12 +671,39 @@ def test_create_tensorboard_run( tensorboard_run=expected_tensorboard_run, tensorboard_run_id=_TEST_TENSORBOARD_RUN_ID, metadata=_TEST_REQUEST_METADATA, + timeout=None, ) get_tensorboard_run_mock.assert_called_once_with( name=_TEST_TENSORBOARD_RUN_NAME, retry=base._DEFAULT_RETRY ) + def test_create_tensorboard_run_with_timeout( + self, create_tensorboard_run_mock, get_tensorboard_run_mock + ): + + aiplatform.init( + project=_TEST_PROJECT, + ) + + tensorboard.TensorboardRun.create( + tensorboard_run_id=_TEST_TENSORBOARD_RUN_ID, + tensorboard_experiment_name=_TEST_TENSORBOARD_EXPERIMENT_NAME, + create_request_timeout=180.0, + ) + + expected_tensorboard_run = gca_tensorboard_run.TensorboardRun( + display_name=_TEST_TENSORBOARD_RUN_ID, + ) + + create_tensorboard_run_mock.assert_called_once_with( + parent=_TEST_TENSORBOARD_EXPERIMENT_NAME, + tensorboard_run=expected_tensorboard_run, + tensorboard_run_id=_TEST_TENSORBOARD_RUN_ID, + metadata=_TEST_REQUEST_METADATA, + timeout=180.0, + ) + @pytest.mark.usefixtures("get_tensorboard_run_mock") def test_delete_tensorboard_run(self, delete_tensorboard_run_mock): aiplatform.init(project=_TEST_PROJECT) diff --git a/tests/unit/aiplatform/test_training_jobs.py b/tests/unit/aiplatform/test_training_jobs.py index 5b78ba4aab..06eaa9218a 100644 --- a/tests/unit/aiplatform/test_training_jobs.py +++ b/tests/unit/aiplatform/test_training_jobs.py @@ -899,6 +899,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( timestamp_split_column_name=_TEST_TIMESTAMP_SPLIT_COLUMN_NAME, tensorboard=_TEST_TENSORBOARD_RESOURCE_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1008,6 +1009,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -1079,6 +1081,7 @@ def test_custom_training_tabular_done( timestamp_split_column_name=_TEST_TIMESTAMP_SPLIT_COLUMN_NAME, tensorboard=_TEST_TENSORBOARD_RESOURCE_NAME, sync=False, + create_request_timeout=None, ) assert job.done() is False @@ -1087,6 +1090,166 @@ def test_custom_training_tabular_done( assert job.done() is True + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_tabular_dataset_and_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_python_package_to_gcs, + mock_tabular_dataset, + mock_model_service_get, + sync, + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_BUCKET_NAME, + credentials=_TEST_CREDENTIALS, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + job = training_jobs.CustomTrainingJob( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + script_path=_TEST_LOCAL_SCRIPT_FILE_NAME, + container_uri=_TEST_TRAINING_CONTAINER_IMAGE, + model_serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE, + model_serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + model_serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, + model_instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI, + model_parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, + model_prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, + model_serving_container_command=_TEST_MODEL_SERVING_CONTAINER_COMMAND, + model_serving_container_args=_TEST_MODEL_SERVING_CONTAINER_ARGS, + model_serving_container_environment_variables=_TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES, + model_serving_container_ports=_TEST_MODEL_SERVING_CONTAINER_PORTS, + model_description=_TEST_MODEL_DESCRIPTION, + ) + + model_from_job = job.run( + dataset=mock_tabular_dataset, + base_output_dir=_TEST_BASE_OUTPUT_DIR, + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + args=_TEST_RUN_ARGS, + environment_variables=_TEST_ENVIRONMENT_VARIABLES, + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + training_fraction_split=_TEST_TRAINING_FRACTION_SPLIT, + validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, + test_fraction_split=_TEST_TEST_FRACTION_SPLIT, + timestamp_split_column_name=_TEST_TIMESTAMP_SPLIT_COLUMN_NAME, + tensorboard=_TEST_TENSORBOARD_RESOURCE_NAME, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_args = _TEST_RUN_ARGS + true_env = [ + {"name": key, "value": value} + for key, value in _TEST_ENVIRONMENT_VARIABLES.items() + ] + + true_worker_pool_spec = { + "replica_count": _TEST_REPLICA_COUNT, + "machine_spec": { + "machine_type": _TEST_MACHINE_TYPE, + "accelerator_type": _TEST_ACCELERATOR_TYPE, + "accelerator_count": _TEST_ACCELERATOR_COUNT, + }, + "disk_spec": { + "boot_disk_type": _TEST_BOOT_DISK_TYPE_DEFAULT, + "boot_disk_size_gb": _TEST_BOOT_DISK_SIZE_GB_DEFAULT, + }, + "python_package_spec": { + "executor_image_uri": _TEST_TRAINING_CONTAINER_IMAGE, + "python_module": _TEST_MODULE_NAME, + "package_uris": [_TEST_OUTPUT_PYTHON_PACKAGE_PATH], + "args": true_args, + "env": true_env, + }, + } + + true_timestamp_split = gca_training_pipeline.TimestampSplit( + training_fraction=_TEST_TRAINING_FRACTION_SPLIT, + validation_fraction=_TEST_VALIDATION_FRACTION_SPLIT, + test_fraction=_TEST_TEST_FRACTION_SPLIT, + key=_TEST_TIMESTAMP_SPLIT_COLUMN_NAME, + ) + + env = [ + gca_env_var.EnvVar(name=str(key), value=str(value)) + for key, value in _TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES.items() + ] + + ports = [ + gca_model.Port(container_port=port) + for port in _TEST_MODEL_SERVING_CONTAINER_PORTS + ] + + true_container_spec = gca_model.ModelContainerSpec( + image_uri=_TEST_SERVING_CONTAINER_IMAGE, + predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, + command=_TEST_MODEL_SERVING_CONTAINER_COMMAND, + args=_TEST_MODEL_SERVING_CONTAINER_ARGS, + env=env, + ports=ports, + ) + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + description=_TEST_MODEL_DESCRIPTION, + container_spec=true_container_spec, + predict_schemata=gca_model.PredictSchemata( + instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI, + parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, + prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, + ), + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + timestamp_split=true_timestamp_split, + dataset_id=mock_tabular_dataset.name, + gcs_destination=gca_io.GcsDestination( + output_uri_prefix=_TEST_BASE_OUTPUT_DIR + ), + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + training_task_definition=schema.training_job.definition.custom_task, + training_task_inputs=json_format.ParseDict( + { + "worker_pool_specs": [true_worker_pool_spec], + "base_output_directory": { + "output_uri_prefix": _TEST_BASE_OUTPUT_DIR + }, + "service_account": _TEST_SERVICE_ACCOUNT, + "network": _TEST_NETWORK, + "tensorboard": _TEST_TENSORBOARD_RESOURCE_NAME, + }, + struct_pb2.Value(), + ), + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + labels=_TEST_LABELS, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_service_create_with_bigquery_destination( self, @@ -1133,6 +1296,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( model_display_name=_TEST_MODEL_DISPLAY_NAME, predefined_split_column_name=_TEST_PREDEFINED_SPLIT_COLUMN_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1226,6 +1390,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -1280,6 +1445,7 @@ def test_run_called_twice_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) with pytest.raises(RuntimeError): @@ -1296,6 +1462,7 @@ def test_run_called_twice_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1332,6 +1499,7 @@ def test_run_with_invalid_accelerator_type_raises( accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1363,6 +1531,7 @@ def test_run_with_two_splits_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1396,6 +1565,7 @@ def test_run_with_incomplete_model_info_raises_with_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -1437,6 +1607,7 @@ def test_run_call_pipeline_service_create_with_no_dataset( validation_filter_split=_TEST_VALIDATION_FILTER_SPLIT, test_filter_split=_TEST_TEST_FILTER_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1502,6 +1673,7 @@ def test_run_call_pipeline_service_create_with_no_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -1547,6 +1719,7 @@ def test_run_call_pipeline_service_create_with_enable_web_access( accelerator_count=_TEST_ACCELERATOR_COUNT, enable_web_access=_TEST_ENABLE_WEB_ACCESS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1587,6 +1760,7 @@ def test_run_call_pipeline_service_create_with_scheduling(self, sync, caplog): timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1637,6 +1811,7 @@ def test_run_returns_none_if_no_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) assert model is None @@ -1672,6 +1847,7 @@ def test_get_model_raises_if_no_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1710,6 +1886,7 @@ def test_run_raises_if_pipeline_fails( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1791,6 +1968,7 @@ def test_run_call_pipeline_service_create_distributed_training( accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -1891,6 +2069,7 @@ def test_run_call_pipeline_service_create_distributed_training( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -1941,6 +2120,7 @@ def test_run_call_pipeline_service_create_distributed_training_with_reduction_se reduction_server_machine_type=_TEST_REDUCTION_SERVER_MACHINE_TYPE, reduction_server_container_uri=_TEST_REDUCTION_SERVER_CONTAINER_URI, sync=sync, + create_request_timeout=None, ) if not sync: @@ -2025,6 +2205,7 @@ def test_run_call_pipeline_service_create_distributed_training_with_reduction_se mock_pipeline_service_create_with_no_model_to_upload.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline_with_no_model_upload( @@ -2203,6 +2384,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ validation_filter_split=_TEST_VALIDATION_FILTER_SPLIT, test_filter_split=_TEST_TEST_FILTER_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -2302,6 +2484,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -2356,6 +2539,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_raises_if_anno accelerator_type=_TEST_ACCELERATOR_TYPE, accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, + create_request_timeout=None, ) @pytest.mark.usefixtures( @@ -2459,6 +2643,7 @@ def test_custom_container_training_tabular_done( service_account=_TEST_SERVICE_ACCOUNT, tensorboard=_TEST_TENSORBOARD_RESOURCE_NAME, sync=False, + create_request_timeout=None, ) assert job.done() is False @@ -2514,6 +2699,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( service_account=_TEST_SERVICE_ACCOUNT, tensorboard=_TEST_TENSORBOARD_RESOURCE_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -2610,6 +2796,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -2630,6 +2817,171 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( assert job._has_logged_custom_job + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_tabular_dataset_and_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_tabular_dataset, + mock_model_service_get, + sync, + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_BUCKET_NAME, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + job = training_jobs.CustomContainerTrainingJob( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + container_uri=_TEST_TRAINING_CONTAINER_IMAGE, + command=_TEST_TRAINING_CONTAINER_CMD, + model_serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE, + model_serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + model_serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, + model_instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI, + model_parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, + model_prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, + model_serving_container_command=_TEST_MODEL_SERVING_CONTAINER_COMMAND, + model_serving_container_args=_TEST_MODEL_SERVING_CONTAINER_ARGS, + model_serving_container_environment_variables=_TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES, + model_serving_container_ports=_TEST_MODEL_SERVING_CONTAINER_PORTS, + model_description=_TEST_MODEL_DESCRIPTION, + ) + + model_from_job = job.run( + dataset=mock_tabular_dataset, + base_output_dir=_TEST_BASE_OUTPUT_DIR, + args=_TEST_RUN_ARGS, + environment_variables=_TEST_ENVIRONMENT_VARIABLES, + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + predefined_split_column_name=_TEST_PREDEFINED_SPLIT_COLUMN_NAME, + service_account=_TEST_SERVICE_ACCOUNT, + tensorboard=_TEST_TENSORBOARD_RESOURCE_NAME, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_args = _TEST_RUN_ARGS + true_env = [ + {"name": key, "value": value} + for key, value in _TEST_ENVIRONMENT_VARIABLES.items() + ] + + true_worker_pool_spec = { + "replica_count": _TEST_REPLICA_COUNT, + "machine_spec": { + "machine_type": _TEST_MACHINE_TYPE, + "accelerator_type": _TEST_ACCELERATOR_TYPE, + "accelerator_count": _TEST_ACCELERATOR_COUNT, + }, + "disk_spec": { + "boot_disk_type": _TEST_BOOT_DISK_TYPE_DEFAULT, + "boot_disk_size_gb": _TEST_BOOT_DISK_SIZE_GB_DEFAULT, + }, + "containerSpec": { + "imageUri": _TEST_TRAINING_CONTAINER_IMAGE, + "command": _TEST_TRAINING_CONTAINER_CMD, + "args": true_args, + "env": true_env, + }, + } + + env = [ + gca_env_var.EnvVar(name=str(key), value=str(value)) + for key, value in _TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES.items() + ] + + ports = [ + gca_model.Port(container_port=port) + for port in _TEST_MODEL_SERVING_CONTAINER_PORTS + ] + + true_container_spec = gca_model.ModelContainerSpec( + image_uri=_TEST_SERVING_CONTAINER_IMAGE, + predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, + command=_TEST_MODEL_SERVING_CONTAINER_COMMAND, + args=_TEST_MODEL_SERVING_CONTAINER_ARGS, + env=env, + ports=ports, + ) + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + description=_TEST_MODEL_DESCRIPTION, + container_spec=true_container_spec, + predict_schemata=gca_model.PredictSchemata( + instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI, + parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, + prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, + ), + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + predefined_split=gca_training_pipeline.PredefinedSplit( + key=_TEST_PREDEFINED_SPLIT_COLUMN_NAME + ), + dataset_id=mock_tabular_dataset.name, + gcs_destination=gca_io.GcsDestination( + output_uri_prefix=_TEST_BASE_OUTPUT_DIR + ), + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.custom_task, + training_task_inputs=json_format.ParseDict( + { + "worker_pool_specs": [true_worker_pool_spec], + "base_output_directory": { + "output_uri_prefix": _TEST_BASE_OUTPUT_DIR + }, + "service_account": _TEST_SERVICE_ACCOUNT, + "tensorboard": _TEST_TENSORBOARD_RESOURCE_NAME, + }, + struct_pb2.Value(), + ), + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + + # assert job._gca_resource == make_training_pipeline( + # gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + # ) + + # mock_model_service_get.assert_called_once_with( + # name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + # ) + + # assert model_from_job._gca_resource is mock_model_service_get.return_value + + # assert job.get_model()._gca_resource is mock_model_service_get.return_value + + # assert not job.has_failed + + # assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + + # assert job._has_logged_custom_job + @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_service_create_with_bigquery_destination( self, @@ -2674,6 +3026,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( test_fraction_split=_TEST_TEST_FRACTION_SPLIT, timestamp_split_column_name=_TEST_TIMESTAMP_SPLIT_COLUMN_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -2766,6 +3119,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -2820,6 +3174,7 @@ def test_run_called_twice_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) with pytest.raises(RuntimeError): @@ -2836,6 +3191,7 @@ def test_run_called_twice_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -2872,6 +3228,7 @@ def test_run_with_invalid_accelerator_type_raises( accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -2902,6 +3259,7 @@ def test_run_with_two_split_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -2935,6 +3293,7 @@ def test_run_with_incomplete_model_info_raises_with_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -2967,6 +3326,7 @@ def test_run_call_pipeline_service_create_with_no_dataset( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3020,6 +3380,7 @@ def test_run_call_pipeline_service_create_with_no_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -3064,6 +3425,7 @@ def test_run_call_pipeline_service_create_with_enable_web_access( accelerator_count=_TEST_ACCELERATOR_COUNT, enable_web_access=_TEST_ENABLE_WEB_ACCESS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3103,6 +3465,7 @@ def test_run_call_pipeline_service_create_with_scheduling(self, sync, caplog): timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3147,6 +3510,7 @@ def test_run_returns_none_if_no_model_to_upload( accelerator_type=_TEST_ACCELERATOR_TYPE, accelerator_count=_TEST_ACCELERATOR_COUNT, sync=sync, + create_request_timeout=None, ) assert model is None @@ -3181,6 +3545,7 @@ def test_get_model_raises_if_no_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3218,6 +3583,7 @@ def test_run_raises_if_pipeline_fails( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3294,6 +3660,7 @@ def test_run_call_pipeline_service_create_distributed_training( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3387,6 +3754,7 @@ def test_run_call_pipeline_service_create_distributed_training( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -3435,6 +3803,7 @@ def test_run_call_pipeline_service_create_distributed_training_with_reduction_se reduction_server_machine_type=_TEST_REDUCTION_SERVER_MACHINE_TYPE, reduction_server_container_uri=_TEST_REDUCTION_SERVER_CONTAINER_URI, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3505,6 +3874,7 @@ def test_run_call_pipeline_service_create_distributed_training_with_reduction_se mock_pipeline_service_create_with_no_model_to_upload.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline_with_no_model_upload( @@ -3564,6 +3934,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset( validation_filter_split=_TEST_VALIDATION_FILTER_SPLIT, test_filter_split=_TEST_TEST_FILTER_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -3658,6 +4029,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -3712,6 +4084,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_raises_if_anno accelerator_type=_TEST_ACCELERATOR_TYPE, accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, + create_request_timeout=None, ) @@ -4088,6 +4461,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4189,6 +4563,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -4207,6 +4582,161 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + @pytest.mark.parametrize("sync", [True, False]) + def test_run_call_pipeline_service_create_with_tabular_dataset_with_timeout( + self, + mock_pipeline_service_create, + mock_pipeline_service_get, + mock_tabular_dataset, + mock_model_service_get, + sync, + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_BUCKET_NAME, + encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME, + ) + + job = training_jobs.CustomPythonPackageTrainingJob( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + python_package_gcs_uri=_TEST_OUTPUT_PYTHON_PACKAGE_PATH, + python_module_name=_TEST_PYTHON_MODULE_NAME, + container_uri=_TEST_TRAINING_CONTAINER_IMAGE, + model_serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE, + model_serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + model_serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, + model_serving_container_command=_TEST_MODEL_SERVING_CONTAINER_COMMAND, + model_serving_container_args=_TEST_MODEL_SERVING_CONTAINER_ARGS, + model_serving_container_environment_variables=_TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES, + model_serving_container_ports=_TEST_MODEL_SERVING_CONTAINER_PORTS, + model_description=_TEST_MODEL_DESCRIPTION, + model_instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI, + model_parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, + model_prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, + ) + + model_from_job = job.run( + dataset=mock_tabular_dataset, + model_display_name=_TEST_MODEL_DISPLAY_NAME, + model_labels=_TEST_MODEL_LABELS, + base_output_dir=_TEST_BASE_OUTPUT_DIR, + service_account=_TEST_SERVICE_ACCOUNT, + network=_TEST_NETWORK, + args=_TEST_RUN_ARGS, + environment_variables=_TEST_ENVIRONMENT_VARIABLES, + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + training_fraction_split=_TEST_TRAINING_FRACTION_SPLIT, + validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, + test_fraction_split=_TEST_TEST_FRACTION_SPLIT, + sync=sync, + create_request_timeout=180.0, + ) + + if not sync: + model_from_job.wait() + + true_args = _TEST_RUN_ARGS + true_env = [ + {"name": key, "value": value} + for key, value in _TEST_ENVIRONMENT_VARIABLES.items() + ] + + true_worker_pool_spec = { + "replica_count": _TEST_REPLICA_COUNT, + "machine_spec": { + "machine_type": _TEST_MACHINE_TYPE, + "accelerator_type": _TEST_ACCELERATOR_TYPE, + "accelerator_count": _TEST_ACCELERATOR_COUNT, + }, + "disk_spec": { + "boot_disk_type": _TEST_BOOT_DISK_TYPE_DEFAULT, + "boot_disk_size_gb": _TEST_BOOT_DISK_SIZE_GB_DEFAULT, + }, + "python_package_spec": { + "executor_image_uri": _TEST_TRAINING_CONTAINER_IMAGE, + "python_module": _TEST_PYTHON_MODULE_NAME, + "package_uris": [_TEST_OUTPUT_PYTHON_PACKAGE_PATH], + "args": true_args, + "env": true_env, + }, + } + + true_fraction_split = gca_training_pipeline.FractionSplit( + training_fraction=_TEST_TRAINING_FRACTION_SPLIT, + validation_fraction=_TEST_VALIDATION_FRACTION_SPLIT, + test_fraction=_TEST_TEST_FRACTION_SPLIT, + ) + + env = [ + gca_env_var.EnvVar(name=str(key), value=str(value)) + for key, value in _TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES.items() + ] + + ports = [ + gca_model.Port(container_port=port) + for port in _TEST_MODEL_SERVING_CONTAINER_PORTS + ] + + true_container_spec = gca_model.ModelContainerSpec( + image_uri=_TEST_SERVING_CONTAINER_IMAGE, + predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE, + health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE, + command=_TEST_MODEL_SERVING_CONTAINER_COMMAND, + args=_TEST_MODEL_SERVING_CONTAINER_ARGS, + env=env, + ports=ports, + ) + + true_managed_model = gca_model.Model( + display_name=_TEST_MODEL_DISPLAY_NAME, + labels=_TEST_MODEL_LABELS, + description=_TEST_MODEL_DESCRIPTION, + container_spec=true_container_spec, + predict_schemata=gca_model.PredictSchemata( + instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI, + parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI, + prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI, + ), + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + true_input_data_config = gca_training_pipeline.InputDataConfig( + fraction_split=true_fraction_split, + dataset_id=mock_tabular_dataset.name, + gcs_destination=gca_io.GcsDestination( + output_uri_prefix=_TEST_BASE_OUTPUT_DIR + ), + ) + + true_training_pipeline = gca_training_pipeline.TrainingPipeline( + display_name=_TEST_DISPLAY_NAME, + labels=_TEST_LABELS, + training_task_definition=schema.training_job.definition.custom_task, + training_task_inputs=json_format.ParseDict( + { + "worker_pool_specs": [true_worker_pool_spec], + "base_output_directory": { + "output_uri_prefix": _TEST_BASE_OUTPUT_DIR + }, + "service_account": _TEST_SERVICE_ACCOUNT, + "network": _TEST_NETWORK, + }, + struct_pb2.Value(), + ), + model_to_upload=true_managed_model, + input_data_config=true_input_data_config, + encryption_spec=_TEST_DEFAULT_ENCRYPTION_SPEC, + ) + + mock_pipeline_service_create.assert_called_once_with( + parent=initializer.global_config.common_location_path(), + training_pipeline=true_training_pipeline, + timeout=180.0, + ) + @pytest.mark.parametrize("sync", [True, False]) def test_run_call_pipeline_service_create_with_tabular_dataset_without_model_display_name_nor_model_labels( self, @@ -4251,6 +4781,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_without_model_dis accelerator_count=_TEST_ACCELERATOR_COUNT, predefined_split_column_name=_TEST_PREDEFINED_SPLIT_COLUMN_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4341,6 +4872,7 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_without_model_dis mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -4404,6 +4936,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( test_fraction_split=_TEST_TEST_FRACTION_SPLIT, timestamp_split_column_name=_TEST_TIMESTAMP_SPLIT_COLUMN_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4497,6 +5030,7 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -4549,6 +5083,7 @@ def test_run_called_twice_raises( accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) with pytest.raises(RuntimeError): @@ -4562,6 +5097,7 @@ def test_run_called_twice_raises( accelerator_count=_TEST_ACCELERATOR_COUNT, model_display_name=_TEST_MODEL_DISPLAY_NAME, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4602,6 +5138,7 @@ def test_run_with_invalid_accelerator_type_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -4634,6 +5171,7 @@ def test_run_with_two_split_raises( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -4668,6 +5206,7 @@ def test_run_with_incomplete_model_info_raises_with_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) @pytest.mark.parametrize("sync", [True, False]) @@ -4701,6 +5240,7 @@ def test_run_call_pipeline_service_create_with_no_dataset( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4755,6 +5295,7 @@ def test_run_call_pipeline_service_create_with_no_dataset( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -4800,6 +5341,7 @@ def test_run_call_pipeline_service_create_with_enable_web_access( accelerator_count=_TEST_ACCELERATOR_COUNT, enable_web_access=_TEST_ENABLE_WEB_ACCESS, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4840,6 +5382,7 @@ def test_run_call_pipeline_service_create_with_scheduling(self, sync, caplog): timeout=_TEST_TIMEOUT, restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART, sync=sync, + create_request_timeout=None, ) if not sync: @@ -4890,6 +5433,7 @@ def test_run_returns_none_if_no_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) assert model is None @@ -4925,6 +5469,7 @@ def test_get_model_raises_if_no_model_to_upload( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -5039,6 +5584,7 @@ def test_run_call_pipeline_service_create_distributed_training( validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT, test_fraction_split=_TEST_TEST_FRACTION_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -5134,6 +5680,7 @@ def test_run_call_pipeline_service_create_distributed_training( mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline( @@ -5183,6 +5730,7 @@ def test_run_call_pipeline_service_create_distributed_training_with_reduction_se reduction_server_machine_type=_TEST_REDUCTION_SERVER_MACHINE_TYPE, reduction_server_container_uri=_TEST_REDUCTION_SERVER_CONTAINER_URI, sync=sync, + create_request_timeout=None, ) if not sync: @@ -5255,6 +5803,7 @@ def test_run_call_pipeline_service_create_distributed_training_with_reduction_se mock_pipeline_service_create_with_no_model_to_upload.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline_with_no_model_upload( @@ -5313,6 +5862,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ validation_filter_split=_TEST_VALIDATION_FILTER_SPLIT, test_filter_split=_TEST_TEST_FILTER_SPLIT, sync=sync, + create_request_timeout=None, ) if not sync: @@ -5408,6 +5958,7 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ mock_pipeline_service_create.assert_called_once_with( parent=initializer.global_config.common_location_path(), training_pipeline=true_training_pipeline, + timeout=None, ) assert job._gca_resource == make_training_pipeline(