diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index ad211ba7fd..a1ee4d348d 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -39,6 +39,7 @@ import proto +from google.api_core import retry from google.api_core import operation from google.auth import credentials as auth_credentials from google.cloud.aiplatform import initializer @@ -48,6 +49,9 @@ logging.basicConfig(level=logging.INFO, stream=sys.stdout) +# This is the default retry callback to be used with get methods. +_DEFAULT_RETRY = retry.Retry() + class Logger: """Logging wrapper class with high level helper methods.""" @@ -532,7 +536,9 @@ def _get_gca_resource(self, resource_name: str) -> proto.Message: location=self.location, ) - return getattr(self.api_client, self._getter_method)(name=resource_name) + return getattr(self.api_client, self._getter_method)( + name=resource_name, retry=_DEFAULT_RETRY + ) def _sync_gca_resource(self): """Sync GAPIC service representation of client class resource.""" diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index b225909eee..3ebcaa5112 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -94,7 +94,7 @@ def __init__( ) self._gca_resource = getattr(self.api_client, self._getter_method)( - name=full_resource_name + name=full_resource_name, retry=base._DEFAULT_RETRY ) @property diff --git a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py index 5aad3225c7..9f758987bc 100644 --- a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py @@ -3,6 +3,7 @@ from unittest import mock from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import schema @@ -301,7 +302,9 @@ def test_run_call_pipeline_service_create( assert job._gca_resource is mock_pipeline_service_get.return_value - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 diff --git a/tests/unit/aiplatform/test_automl_image_training_jobs.py b/tests/unit/aiplatform/test_automl_image_training_jobs.py index 330fe3d0ea..338edf7adc 100644 --- a/tests/unit/aiplatform/test_automl_image_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_image_training_jobs.py @@ -6,7 +6,7 @@ from google.protobuf import struct_pb2 from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models @@ -309,7 +309,9 @@ def test_run_call_pipeline_service_create( training_pipeline=true_training_pipeline, ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + mock_model_service_get.assert_called_once_with( + name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + ) assert job._gca_resource is mock_pipeline_service_get.return_value 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 diff --git a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py index 41614b738f..8b00792944 100644 --- a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py @@ -3,7 +3,7 @@ from unittest import mock from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import schema @@ -367,7 +367,9 @@ def test_run_call_pipeline_service_create( assert job._gca_resource is mock_pipeline_service_get.return_value - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -446,7 +448,9 @@ def test_run_call_pipeline_service_create_with_export_eval_data_items( assert job._gca_resource is mock_pipeline_service_get.return_value - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 diff --git a/tests/unit/aiplatform/test_automl_text_training_jobs.py b/tests/unit/aiplatform/test_automl_text_training_jobs.py index 20220a1247..76d6f789b4 100644 --- a/tests/unit/aiplatform/test_automl_text_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_text_training_jobs.py @@ -3,7 +3,7 @@ from unittest import mock from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models @@ -370,7 +370,9 @@ def test_run_call_pipeline_service_create_classification( training_pipeline=true_training_pipeline, ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + mock_model_service_get.assert_called_once_with( + name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + ) assert job._gca_resource is mock_pipeline_service_get.return_value 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 @@ -437,7 +439,9 @@ def test_run_call_pipeline_service_create_extraction( training_pipeline=true_training_pipeline, ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + mock_model_service_get.assert_called_once_with( + name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + ) assert job._gca_resource is mock_pipeline_service_get.return_value 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 @@ -505,7 +509,9 @@ def test_run_call_pipeline_service_create_sentiment( training_pipeline=true_training_pipeline, ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + mock_model_service_get.assert_called_once_with( + name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + ) assert job._gca_resource is mock_pipeline_service_get.return_value 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 diff --git a/tests/unit/aiplatform/test_automl_video_training_jobs.py b/tests/unit/aiplatform/test_automl_video_training_jobs.py index 7326050ae4..58ea230313 100644 --- a/tests/unit/aiplatform/test_automl_video_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_video_training_jobs.py @@ -6,7 +6,7 @@ from google.protobuf import struct_pb2 from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models @@ -271,7 +271,9 @@ def test_init_aiplatform_with_encryption_key_name_and_create_training_job( training_pipeline=true_training_pipeline, ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + mock_model_service_get.assert_called_once_with( + name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + ) assert job._gca_resource is mock_pipeline_service_get.return_value 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 @@ -538,7 +540,9 @@ def test_run_call_pipeline_service_create( training_pipeline=true_training_pipeline, ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + mock_model_service_get.assert_called_once_with( + name=_TEST_MODEL_NAME, retry=base._DEFAULT_RETRY + ) assert job._gca_resource is mock_pipeline_service_get.return_value 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 diff --git a/tests/unit/aiplatform/test_custom_job.py b/tests/unit/aiplatform/test_custom_job.py index 5fc3243451..ea362ecda8 100644 --- a/tests/unit/aiplatform/test_custom_job.py +++ b/tests/unit/aiplatform/test_custom_job.py @@ -29,6 +29,7 @@ from test_training_jobs import mock_python_package_to_gcs # noqa: F401 from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform.compat.types import custom_job as gca_custom_job_compat from google.cloud.aiplatform.compat.types import ( custom_job_v1beta1 as gca_custom_job_v1beta1, @@ -447,7 +448,9 @@ def test_get_custom_job(self, get_custom_job_mock): job = aiplatform.CustomJob.get(_TEST_CUSTOM_JOB_NAME) - get_custom_job_mock.assert_called_once_with(name=_TEST_CUSTOM_JOB_NAME) + get_custom_job_mock.assert_called_once_with( + name=_TEST_CUSTOM_JOB_NAME, retry=base._DEFAULT_RETRY + ) assert ( job._gca_resource.state == gca_job_state_compat.JobState.JOB_STATE_PENDING ) diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index 7864ca0d35..f80beac240 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -28,13 +28,13 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform -from google.cloud import bigquery -from google.cloud import storage - +from google.cloud.aiplatform import base from google.cloud.aiplatform import compat from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import schema +from google.cloud import bigquery +from google.cloud import storage from google.cloud.aiplatform_v1.services.dataset_service import ( client as dataset_service_client, @@ -474,7 +474,9 @@ def teardown_method(self): def test_init_dataset(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT) datasets._Dataset(dataset_name=_TEST_NAME) - get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_dataset_with_id_only_with_project_and_location( self, get_dataset_mock @@ -483,21 +485,27 @@ def test_init_dataset_with_id_only_with_project_and_location( datasets._Dataset( dataset_name=_TEST_ID, project=_TEST_PROJECT, location=_TEST_LOCATION ) - get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_dataset_with_project_and_location(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT) datasets._Dataset( dataset_name=_TEST_NAME, project=_TEST_PROJECT, location=_TEST_LOCATION ) - get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_dataset_with_alt_project_and_location(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT) datasets._Dataset( dataset_name=_TEST_NAME, project=_TEST_ALT_PROJECT, location=_TEST_LOCATION ) - get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_dataset_with_alt_location(self, get_dataset_tabular_gcs_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_ALT_LOCATION) @@ -511,7 +519,9 @@ def test_init_dataset_with_alt_location(self, get_dataset_tabular_gcs_mock): assert _TEST_ALT_LOCATION != _TEST_LOCATION - get_dataset_tabular_gcs_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_tabular_gcs_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_dataset_with_project_and_alt_location(self): aiplatform.init(project=_TEST_PROJECT) @@ -525,7 +535,9 @@ def test_init_dataset_with_project_and_alt_location(self): def test_init_dataset_with_id_only(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) datasets._Dataset(dataset_name=_TEST_ID) - get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_dataset_without_name_mock") @patch.dict( @@ -541,7 +553,9 @@ def test_init_dataset_with_id_only_without_project_or_location(self): def test_init_dataset_with_location_override(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) datasets._Dataset(dataset_name=_TEST_ID, location=_TEST_ALT_LOCATION) - get_dataset_mock.assert_called_once_with(name=_TEST_ALT_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_ALT_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_dataset_mock") def test_init_dataset_with_invalid_name(self): @@ -764,7 +778,9 @@ def test_create_then_import( metadata=_TEST_REQUEST_METADATA, ) - get_dataset_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) import_data_mock.assert_called_once_with( name=_TEST_NAME, import_configs=[expected_import_config] @@ -798,7 +814,9 @@ def teardown_method(self): def test_init_dataset_image(self, get_dataset_image_mock): aiplatform.init(project=_TEST_PROJECT) datasets.ImageDataset(dataset_name=_TEST_NAME) - get_dataset_image_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_image_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_dataset_tabular_bq_mock") def test_init_dataset_non_image(self): @@ -934,7 +952,9 @@ def test_create_then_import( metadata=_TEST_REQUEST_METADATA, ) - get_dataset_image_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_image_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) expected_import_config = gca_dataset.ImportDataConfig( gcs_source=gca_io.GcsSource(uris=[_TEST_SOURCE_URI_GCS]), @@ -989,7 +1009,9 @@ def teardown_method(self): def test_init_dataset_tabular(self, get_dataset_tabular_bq_mock): datasets.TabularDataset(dataset_name=_TEST_NAME) - get_dataset_tabular_bq_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_tabular_bq_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_dataset_image_mock") def test_init_dataset_non_tabular(self): @@ -1236,7 +1258,9 @@ def teardown_method(self): def test_init_dataset_text(self, get_dataset_text_mock): aiplatform.init(project=_TEST_PROJECT) datasets.TextDataset(dataset_name=_TEST_NAME) - get_dataset_text_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_text_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_dataset_image_mock") def test_init_dataset_non_text(self): @@ -1409,7 +1433,9 @@ def test_create_then_import( metadata=_TEST_REQUEST_METADATA, ) - get_dataset_text_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_text_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) expected_import_config = gca_dataset.ImportDataConfig( gcs_source=gca_io.GcsSource(uris=[_TEST_SOURCE_URI_GCS]), @@ -1463,7 +1489,9 @@ def teardown_method(self): def test_init_dataset_video(self, get_dataset_video_mock): aiplatform.init(project=_TEST_PROJECT) datasets.VideoDataset(dataset_name=_TEST_NAME) - get_dataset_video_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_video_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_dataset_tabular_bq_mock") def test_init_dataset_non_video(self): @@ -1599,7 +1627,9 @@ def test_create_then_import( metadata=_TEST_REQUEST_METADATA, ) - get_dataset_video_mock.assert_called_once_with(name=_TEST_NAME) + get_dataset_video_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) expected_import_config = gca_dataset.ImportDataConfig( gcs_source=gca_io.GcsSource(uris=[_TEST_SOURCE_URI_GCS]), diff --git a/tests/unit/aiplatform/test_end_to_end.py b/tests/unit/aiplatform/test_end_to_end.py index 10ba0c3b0c..37b6d57649 100644 --- a/tests/unit/aiplatform/test_end_to_end.py +++ b/tests/unit/aiplatform/test_end_to_end.py @@ -19,11 +19,12 @@ from importlib import reload -from google.cloud.aiplatform.utils import source_utils from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models from google.cloud.aiplatform import schema +from google.cloud.aiplatform.utils import source_utils from google.cloud.aiplatform_v1.types import ( dataset as gca_dataset, @@ -280,7 +281,7 @@ def test_dataset_create_to_model_predict( ) mock_model_service_get.assert_called_once_with( - name=test_training_jobs._TEST_MODEL_NAME + name=test_training_jobs._TEST_MODEL_NAME, retry=base._DEFAULT_RETRY ) assert model_from_job._gca_resource is mock_model_service_get.return_value diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index 00fe5093cf..4ec527e31e 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/tests/unit/aiplatform/test_endpoints.py @@ -25,7 +25,7 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import explain from google.cloud.aiplatform import models @@ -408,7 +408,9 @@ def test_lazy_constructor_calls_get_on_property_access(self, get_endpoint_mock): assert not get_endpoint_mock.called ep.display_name # Retrieve a property that requires a call to Endpoint getter - get_endpoint_mock.assert_called_with(name=_TEST_ENDPOINT_NAME) + get_endpoint_mock.assert_called_with( + name=_TEST_ENDPOINT_NAME, retry=base._DEFAULT_RETRY + ) def test_lazy_constructor_with_custom_project(self, get_endpoint_mock): ep = models.Endpoint(endpoint_name=_TEST_ID, project=_TEST_PROJECT_2) @@ -418,7 +420,9 @@ def test_lazy_constructor_with_custom_project(self, get_endpoint_mock): assert not get_endpoint_mock.called ep.name # Retrieve a property that requires a call to Endpoint getter - get_endpoint_mock.assert_called_with(name=test_endpoint_resource_name) + get_endpoint_mock.assert_called_with( + name=test_endpoint_resource_name, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_endpoint_mock") def test_constructor_with_conflicting_location(self): @@ -447,7 +451,7 @@ def test_lazy_constructor_with_custom_location( ep.network # Accessing a property that requires calling getter get_endpoint_alt_location_mock.assert_called_with( - name=test_endpoint_resource_name + name=test_endpoint_resource_name, retry=base._DEFAULT_RETRY ) def test_constructor_with_custom_credentials(self, create_endpoint_client_mock): diff --git a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py index 1ac9630882..7b4f434736 100644 --- a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py +++ b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py @@ -25,6 +25,7 @@ from google.rpc import status_pb2 from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform import hyperparameter_tuning as hpt from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec_compat, @@ -598,7 +599,7 @@ def test_get_hyperparameter_tuning_job(self, get_hyperparameter_tuning_job_mock) ) get_hyperparameter_tuning_job_mock.assert_called_once_with( - name=_TEST_HYPERPARAMETERTUNING_JOB_NAME + name=_TEST_HYPERPARAMETERTUNING_JOB_NAME, retry=base._DEFAULT_RETRY ) assert ( job._gca_resource.state == gca_job_state_compat.JobState.JOB_STATE_PENDING diff --git a/tests/unit/aiplatform/test_jobs.py b/tests/unit/aiplatform/test_jobs.py index 462295c6be..9292ed27c3 100644 --- a/tests/unit/aiplatform/test_jobs.py +++ b/tests/unit/aiplatform/test_jobs.py @@ -27,7 +27,7 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import jobs @@ -381,7 +381,7 @@ def test_init_batch_prediction_job(self, get_batch_prediction_job_mock): batch_prediction_job_name=_TEST_BATCH_PREDICTION_JOB_NAME ) get_batch_prediction_job_mock.assert_called_once_with( - name=_TEST_BATCH_PREDICTION_JOB_NAME + name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=base._DEFAULT_RETRY ) def test_batch_prediction_job_status(self, get_batch_prediction_job_mock): @@ -396,7 +396,7 @@ def test_batch_prediction_job_status(self, get_batch_prediction_job_mock): assert bp_job_state == _TEST_JOB_STATE_RUNNING get_batch_prediction_job_mock.assert_called_with( - name=_TEST_BATCH_PREDICTION_JOB_NAME + name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=base._DEFAULT_RETRY ) @pytest.mark.usefixtures("get_batch_prediction_job_gcs_output_mock") diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index 70b8ea6b5c..5d0020507c 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -25,6 +25,7 @@ from google.auth import credentials from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import metadata @@ -390,8 +391,12 @@ def test_init_experiment_with_existing_metadataStore_and_context( project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT ) - get_metadata_store_mock.assert_called_once_with(name=_TEST_METADATASTORE) - get_context_mock.assert_called_once_with(name=_TEST_CONTEXT_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_METADATASTORE, retry=base._DEFAULT_RETRY + ) + get_context_mock.assert_called_once_with( + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY + ) def test_init_experiment_with_credentials( self, get_metadata_store_mock, get_context_mock @@ -410,8 +415,12 @@ def test_init_experiment_with_credentials( == creds ) - get_metadata_store_mock.assert_called_once_with(name=_TEST_METADATASTORE) - get_context_mock.assert_called_once_with(name=_TEST_CONTEXT_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_METADATASTORE, retry=base._DEFAULT_RETRY + ) + get_context_mock.assert_called_once_with( + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY + ) def test_init_and_get_metadata_store_with_credentials( self, get_metadata_store_mock @@ -451,8 +460,12 @@ def test_init_experiment_with_existing_description( experiment_description=_TEST_EXPERIMENT_DESCRIPTION, ) - get_metadata_store_mock.assert_called_once_with(name=_TEST_METADATASTORE) - get_context_mock.assert_called_once_with(name=_TEST_CONTEXT_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_METADATASTORE, retry=base._DEFAULT_RETRY + ) + get_context_mock.assert_called_once_with( + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_metadata_store_mock") @pytest.mark.usefixtures("get_context_mock") @@ -521,13 +534,17 @@ def test_start_run_with_existing_execution_and_artifact( ) aiplatform.start_run(_TEST_RUN) - get_execution_mock.assert_called_once_with(name=_TEST_EXECUTION_NAME) + get_execution_mock.assert_called_once_with( + name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY + ) add_context_artifacts_and_executions_mock.assert_called_once_with( context=_TEST_CONTEXT_NAME, artifacts=None, executions=[_TEST_EXECUTION_NAME], ) - get_artifact_mock.assert_called_once_with(name=_TEST_ARTIFACT_NAME) + get_artifact_mock.assert_called_once_with( + name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY + ) add_execution_events_mock.assert_called_once_with( execution=_TEST_EXECUTION_NAME, events=[Event(artifact=_TEST_ARTIFACT_NAME, type_=Event.Type.OUTPUT)], diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py index 7443dd695f..64e443bfd0 100644 --- a/tests/unit/aiplatform/test_metadata_resources.py +++ b/tests/unit/aiplatform/test_metadata_resources.py @@ -22,6 +22,7 @@ from google.api_core import exceptions from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import context @@ -366,14 +367,18 @@ def teardown_method(self): def test_init_context(self, get_context_mock): aiplatform.init(project=_TEST_PROJECT) context._Context(resource_name=_TEST_CONTEXT_NAME) - get_context_mock.assert_called_once_with(name=_TEST_CONTEXT_NAME) + get_context_mock.assert_called_once_with( + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY + ) def test_init_context_with_id(self, get_context_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) context._Context( resource_name=_TEST_CONTEXT_ID, metadata_store_id=_TEST_METADATA_STORE ) - get_context_mock.assert_called_once_with(name=_TEST_CONTEXT_NAME) + get_context_mock.assert_called_once_with( + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY + ) def test_get_or_create_context( self, get_context_for_get_or_create_mock, create_context_mock @@ -398,7 +403,7 @@ def test_get_or_create_context( metadata=_TEST_METADATA, ) get_context_for_get_or_create_mock.assert_called_once_with( - name=_TEST_CONTEXT_NAME + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY ) create_context_mock.assert_called_once_with( parent=_TEST_PARENT, context_id=_TEST_CONTEXT_ID, context=expected_context, @@ -542,14 +547,18 @@ def teardown_method(self): def test_init_execution(self, get_execution_mock): aiplatform.init(project=_TEST_PROJECT) execution._Execution(resource_name=_TEST_EXECUTION_NAME) - get_execution_mock.assert_called_once_with(name=_TEST_EXECUTION_NAME) + get_execution_mock.assert_called_once_with( + name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY + ) def test_init_execution_with_id(self, get_execution_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) execution._Execution( resource_name=_TEST_EXECUTION_ID, metadata_store_id=_TEST_METADATA_STORE ) - get_execution_mock.assert_called_once_with(name=_TEST_EXECUTION_NAME) + get_execution_mock.assert_called_once_with( + name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY + ) def test_get_or_create_execution( self, get_execution_for_get_or_create_mock, create_execution_mock @@ -574,7 +583,7 @@ def test_get_or_create_execution( metadata=_TEST_METADATA, ) get_execution_for_get_or_create_mock.assert_called_once_with( - name=_TEST_EXECUTION_NAME + name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY ) create_execution_mock.assert_called_once_with( parent=_TEST_PARENT, @@ -704,14 +713,18 @@ def teardown_method(self): def test_init_artifact(self, get_artifact_mock): aiplatform.init(project=_TEST_PROJECT) artifact._Artifact(resource_name=_TEST_ARTIFACT_NAME) - get_artifact_mock.assert_called_once_with(name=_TEST_ARTIFACT_NAME) + get_artifact_mock.assert_called_once_with( + name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY + ) def test_init_artifact_with_id(self, get_artifact_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) artifact._Artifact( resource_name=_TEST_ARTIFACT_ID, metadata_store_id=_TEST_METADATA_STORE ) - get_artifact_mock.assert_called_once_with(name=_TEST_ARTIFACT_NAME) + get_artifact_mock.assert_called_once_with( + name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY + ) def test_get_or_create_artifact( self, get_artifact_for_get_or_create_mock, create_artifact_mock @@ -736,7 +749,7 @@ def test_get_or_create_artifact( metadata=_TEST_METADATA, ) get_artifact_for_get_or_create_mock.assert_called_once_with( - name=_TEST_ARTIFACT_NAME + name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY ) create_artifact_mock.assert_called_once_with( parent=_TEST_PARENT, diff --git a/tests/unit/aiplatform/test_metadata_store.py b/tests/unit/aiplatform/test_metadata_store.py index a10a86675d..ae3775d2a2 100644 --- a/tests/unit/aiplatform/test_metadata_store.py +++ b/tests/unit/aiplatform/test_metadata_store.py @@ -26,6 +26,7 @@ from google.auth.exceptions import GoogleAuthError from google.cloud import aiplatform +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform_v1 import MetadataServiceClient @@ -143,17 +144,23 @@ def teardown_method(self): def test_init_metadata_store(self, get_metadata_store_mock): aiplatform.init(project=_TEST_PROJECT) metadata_store._MetadataStore(metadata_store_name=_TEST_NAME) - get_metadata_store_mock.assert_called_once_with(name=_TEST_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_metadata_store_with_id(self, get_metadata_store_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) metadata_store._MetadataStore(metadata_store_name=_TEST_ID) - get_metadata_store_mock.assert_called_once_with(name=_TEST_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_metadata_store_with_default_id(self, get_metadata_store_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) metadata_store._MetadataStore() - get_metadata_store_mock.assert_called_once_with(name=_TEST_DEFAULT_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_DEFAULT_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_metadata_store_without_name_mock") @patch.dict( @@ -171,7 +178,9 @@ def test_init_metadata_store_with_location_override(self, get_metadata_store_moc metadata_store._MetadataStore( metadata_store_name=_TEST_ID, location=_TEST_ALT_LOCATION ) - get_metadata_store_mock.assert_called_once_with(name=_TEST_ALT_LOC_NAME) + get_metadata_store_mock.assert_called_once_with( + name=_TEST_ALT_LOC_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_metadata_store_mock") def test_init_metadata_store_with_invalid_name(self): diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index cfd2669723..0c9a649ad8 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/tests/unit/aiplatform/test_models.py @@ -26,7 +26,7 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import models from google.cloud.aiplatform import utils @@ -474,21 +474,27 @@ def test_constructor_creates_client_with_custom_credentials( def test_constructor_gets_model(self, get_model_mock): models.Model(_TEST_ID) - get_model_mock.assert_called_once_with(name=_TEST_MODEL_RESOURCE_NAME) + get_model_mock.assert_called_once_with( + name=_TEST_MODEL_RESOURCE_NAME, retry=base._DEFAULT_RETRY + ) def test_constructor_gets_model_with_custom_project(self, get_model_mock): models.Model(_TEST_ID, project=_TEST_PROJECT_2) test_model_resource_name = model_service_client.ModelServiceClient.model_path( _TEST_PROJECT_2, _TEST_LOCATION, _TEST_ID ) - get_model_mock.assert_called_once_with(name=test_model_resource_name) + get_model_mock.assert_called_once_with( + name=test_model_resource_name, retry=base._DEFAULT_RETRY + ) def test_constructor_gets_model_with_custom_location(self, get_model_mock): models.Model(_TEST_ID, location=_TEST_LOCATION_2) test_model_resource_name = model_service_client.ModelServiceClient.model_path( _TEST_PROJECT, _TEST_LOCATION_2, _TEST_ID ) - get_model_mock.assert_called_once_with(name=test_model_resource_name) + 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_uploads_and_gets_model( @@ -521,7 +527,9 @@ def test_upload_uploads_and_gets_model( model=managed_model, ) - get_model_mock.assert_called_once_with(name=_TEST_MODEL_RESOURCE_NAME) + 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_uploads_and_gets_model_with_labels( @@ -557,7 +565,9 @@ def test_upload_uploads_and_gets_model_with_labels( model=managed_model, ) - get_model_mock.assert_called_once_with(name=_TEST_MODEL_RESOURCE_NAME) + get_model_mock.assert_called_once_with( + name=_TEST_MODEL_RESOURCE_NAME, retry=base._DEFAULT_RETRY + ) def test_upload_raises_with_impartial_explanation_spec(self): @@ -641,7 +651,9 @@ def test_upload_uploads_and_gets_model_with_all_args( parent=initializer.global_config.common_location_path(), model=managed_model, ) - get_model_mock.assert_called_once_with(name=_TEST_MODEL_RESOURCE_NAME) + get_model_mock.assert_called_once_with( + name=_TEST_MODEL_RESOURCE_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_model_with_custom_project_mock") @pytest.mark.parametrize("sync", [True, False]) @@ -687,7 +699,7 @@ def test_upload_uploads_and_gets_model_with_custom_project( ) get_model_with_custom_project_mock.assert_called_once_with( - name=test_model_resource_name + name=test_model_resource_name, retry=base._DEFAULT_RETRY ) assert my_model.uri == _TEST_ARTIFACT_URI @@ -774,7 +786,7 @@ def test_upload_uploads_and_gets_model_with_custom_location( ) get_model_with_custom_location_mock.assert_called_once_with( - name=test_model_resource_name + name=test_model_resource_name, retry=base._DEFAULT_RETRY ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index d6580de24d..f8a62a19c3 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -25,9 +25,10 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform -from google.cloud import storage -from google.cloud.aiplatform import pipeline_jobs +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import pipeline_jobs +from google.cloud import storage from google.protobuf import json_format from google.cloud.aiplatform_v1beta1.services.pipeline_service import ( @@ -266,7 +267,9 @@ def test_run_call_pipeline_service_pipeline_job_create( pipeline_job_id=_TEST_PIPELINE_JOB_ID, ) - mock_pipeline_service_get.assert_called_with(name=_TEST_PIPELINE_JOB_NAME) + 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_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED @@ -325,7 +328,9 @@ def test_run_call_pipeline_service_pipeline_spec_create( pipeline_job_id=_TEST_PIPELINE_JOB_ID, ) - mock_pipeline_service_get.assert_called_with(name=_TEST_PIPELINE_JOB_NAME) + 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_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED @@ -336,7 +341,9 @@ def test_get_pipeline_job(self, mock_pipeline_service_get): aiplatform.init(project=_TEST_PROJECT) job = pipeline_jobs.PipelineJob.get(resource_name=_TEST_PIPELINE_JOB_ID) - mock_pipeline_service_get.assert_called_once_with(name=_TEST_PIPELINE_JOB_NAME) + mock_pipeline_service_get.assert_called_once_with( + name=_TEST_PIPELINE_JOB_NAME, retry=base._DEFAULT_RETRY + ) assert isinstance(job, pipeline_jobs.PipelineJob) @pytest.mark.usefixtures( diff --git a/tests/unit/aiplatform/test_tensorboard.py b/tests/unit/aiplatform/test_tensorboard.py index 60f85683df..ffa4c72850 100644 --- a/tests/unit/aiplatform/test_tensorboard.py +++ b/tests/unit/aiplatform/test_tensorboard.py @@ -28,7 +28,7 @@ from google.auth import credentials as auth_credentials from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import tensorboard @@ -143,7 +143,9 @@ def teardown_method(self): def test_init_tensorboard(self, get_tensorboard_mock): aiplatform.init(project=_TEST_PROJECT) tensorboard.Tensorboard(tensorboard_name=_TEST_NAME) - get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + get_tensorboard_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_tensorboard_with_id_only_with_project_and_location( self, get_tensorboard_mock @@ -152,14 +154,18 @@ def test_init_tensorboard_with_id_only_with_project_and_location( tensorboard.Tensorboard( tensorboard_name=_TEST_ID, project=_TEST_PROJECT, location=_TEST_LOCATION ) - get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + get_tensorboard_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_tensorboard_with_project_and_location(self, get_tensorboard_mock): aiplatform.init(project=_TEST_PROJECT) tensorboard.Tensorboard( tensorboard_name=_TEST_NAME, project=_TEST_PROJECT, location=_TEST_LOCATION ) - get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + get_tensorboard_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_tensorboard_with_alt_project_and_location(self, get_tensorboard_mock): aiplatform.init(project=_TEST_PROJECT) @@ -168,12 +174,16 @@ def test_init_tensorboard_with_alt_project_and_location(self, get_tensorboard_mo project=_TEST_ALT_PROJECT, location=_TEST_LOCATION, ) - get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + get_tensorboard_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_tensorboard_with_alt_location(self, get_tensorboard_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_ALT_LOCATION) tensorboard.Tensorboard(tensorboard_name=_TEST_NAME,) - get_tensorboard_mock.assert_called_once_with(name=_TEST_NAME) + get_tensorboard_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_init_tensorboard_with_project_and_alt_location(self): aiplatform.init(project=_TEST_PROJECT) @@ -197,7 +207,9 @@ def test_init_tensorboard_with_id_only_without_project_or_location(self): def test_init_tensorboard_with_location_override(self, get_tensorboard_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) tensorboard.Tensorboard(tensorboard_name=_TEST_ID, location=_TEST_ALT_LOCATION) - get_tensorboard_mock.assert_called_once_with(name=_TEST_ALT_NAME) + get_tensorboard_mock.assert_called_once_with( + name=_TEST_ALT_NAME, retry=base._DEFAULT_RETRY + ) @pytest.mark.usefixtures("get_tensorboard_mock") def test_init_tensorboard_with_invalid_name(self): diff --git a/tests/unit/aiplatform/test_training_jobs.py b/tests/unit/aiplatform/test_training_jobs.py index cf4b4e7ca8..8efdf1a5c2 100644 --- a/tests/unit/aiplatform/test_training_jobs.py +++ b/tests/unit/aiplatform/test_training_jobs.py @@ -33,15 +33,16 @@ from google.auth import credentials as auth_credentials -from google.cloud.aiplatform import utils -from google.cloud.aiplatform.utils import source_utils -from google.cloud.aiplatform.utils import worker_spec_utils from google.cloud import aiplatform - +from google.cloud.aiplatform import base from google.cloud.aiplatform import datasets from google.cloud.aiplatform import initializer from google.cloud.aiplatform import schema from google.cloud.aiplatform import training_jobs +from google.cloud.aiplatform import utils +from google.cloud.aiplatform.utils import source_utils +from google.cloud.aiplatform.utils import worker_spec_utils + from google.cloud.aiplatform_v1.services.job_service import client as job_service_client from google.cloud.aiplatform_v1.services.model_service import ( @@ -911,7 +912,9 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -1067,7 +1070,9 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -1339,7 +1344,9 @@ def test_run_call_pipeline_service_create_with_no_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -1674,7 +1681,9 @@ def test_run_call_pipeline_service_create_distributed_training( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -1815,7 +1824,9 @@ def test_get_training_job(self, get_training_job_custom_mock): aiplatform.init(project=_TEST_PROJECT) job = training_jobs.CustomTrainingJob.get(resource_name=_TEST_NAME) - get_training_job_custom_mock.assert_called_once_with(name=_TEST_NAME) + get_training_job_custom_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) assert isinstance(job, training_jobs.CustomTrainingJob) @pytest.mark.usefixtures("get_training_job_custom_mock") @@ -1850,7 +1861,9 @@ def test_get_training_job_tabular(self, get_training_job_tabular_mock): def test_get_training_job_with_id_only(self, get_training_job_custom_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) training_jobs.CustomTrainingJob.get(resource_name=_TEST_ID) - get_training_job_custom_mock.assert_called_once_with(name=_TEST_NAME) + get_training_job_custom_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_get_training_job_with_id_only_with_project_and_location( self, get_training_job_custom_mock @@ -1859,7 +1872,9 @@ def test_get_training_job_with_id_only_with_project_and_location( training_jobs.CustomTrainingJob.get( resource_name=_TEST_ID, project=_TEST_PROJECT, location=_TEST_LOCATION ) - get_training_job_custom_mock.assert_called_once_with(name=_TEST_NAME) + get_training_job_custom_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_get_training_job_with_project_and_location( self, get_training_job_custom_mock @@ -1868,7 +1883,9 @@ def test_get_training_job_with_project_and_location( training_jobs.CustomTrainingJob.get( resource_name=_TEST_NAME, project=_TEST_PROJECT, location=_TEST_LOCATION ) - get_training_job_custom_mock.assert_called_once_with(name=_TEST_NAME) + get_training_job_custom_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_get_training_job_with_alt_project_and_location( self, get_training_job_custom_mock @@ -1877,7 +1894,9 @@ def test_get_training_job_with_alt_project_and_location( training_jobs.CustomTrainingJob.get( resource_name=_TEST_NAME, project=_TEST_ALT_PROJECT, location=_TEST_LOCATION ) - get_training_job_custom_mock.assert_called_once_with(name=_TEST_NAME) + get_training_job_custom_mock.assert_called_once_with( + name=_TEST_NAME, retry=base._DEFAULT_RETRY + ) def test_get_training_job_with_project_and_alt_location(self): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) @@ -2073,7 +2092,9 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -2322,7 +2343,9 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -2476,7 +2499,9 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -2726,7 +2751,9 @@ def test_run_call_pipeline_service_create_with_no_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -3042,7 +3069,9 @@ def test_run_call_pipeline_service_create_distributed_training( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -3310,7 +3339,9 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -3833,7 +3864,9 @@ def test_run_call_pipeline_service_create_with_tabular_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -3983,7 +4016,9 @@ def test_run_call_pipeline_service_create_with_tabular_dataset_without_model_dis gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -4137,7 +4172,9 @@ def test_run_call_pipeline_service_create_with_bigquery_destination( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -4391,7 +4428,9 @@ def test_run_call_pipeline_service_create_with_no_dataset( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -4716,7 +4755,9 @@ def test_run_call_pipeline_service_create_distributed_training( gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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 @@ -4987,7 +5028,9 @@ def test_run_call_pipeline_service_create_with_nontabular_dataset_without_model_ gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) - mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME) + 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