diff --git a/datalabeling/create_annotation_spec_set_test.py b/datalabeling/create_annotation_spec_set_test.py index bcbfb7cb339b..7972c7330d63 100644 --- a/datalabeling/create_annotation_spec_set_test.py +++ b/datalabeling/create_annotation_spec_set_test.py @@ -17,7 +17,7 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServerError import pytest import create_annotation_spec_set @@ -40,7 +40,7 @@ def cleaner(): def test_create_annotation_spec_set(cleaner, capsys): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): return create_annotation_spec_set.create_annotation_spec_set(PROJECT_ID) diff --git a/datalabeling/create_instruction_test.py b/datalabeling/create_instruction_test.py index 183704b2bd7d..6bc7c875785d 100644 --- a/datalabeling/create_instruction_test.py +++ b/datalabeling/create_instruction_test.py @@ -17,7 +17,7 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServerError import pytest import create_instruction @@ -42,7 +42,7 @@ def cleaner(): def test_create_instruction(cleaner, capsys): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): return create_instruction.create_instruction( PROJECT_ID, 'IMAGE', INSTRUCTION_GCS_URI) diff --git a/datalabeling/import_data_test.py b/datalabeling/import_data_test.py index c5b111748a63..a2b8e5dad7ea 100644 --- a/datalabeling/import_data_test.py +++ b/datalabeling/import_data_test.py @@ -17,7 +17,7 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServerError import pytest import import_data @@ -42,7 +42,7 @@ def dataset(): def test_import_data(capsys, dataset): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): import_data.import_data(dataset.name, 'IMAGE', INPUT_GCS_URI) diff --git a/datalabeling/label_image_test.py b/datalabeling/label_image_test.py index f50606623e0b..d2d5e69784d9 100644 --- a/datalabeling/label_image_test.py +++ b/datalabeling/label_image_test.py @@ -17,7 +17,7 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServerError import pytest import label_image @@ -80,7 +80,7 @@ def test_label_image( capsys, annotation_spec_set, instruction, dataset, cleaner): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): # Start labeling. return label_image.label_image( diff --git a/datalabeling/label_text_test.py b/datalabeling/label_text_test.py index b4bb5194fabb..a189f99a03ba 100644 --- a/datalabeling/label_text_test.py +++ b/datalabeling/label_text_test.py @@ -17,7 +17,7 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServerError import pytest import label_text @@ -80,7 +80,7 @@ def cleaner(): def test_label_text(capsys, annotation_spec_set, instruction, dataset, cleaner): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): # Start labeling. return label_text.label_text( diff --git a/datalabeling/label_video_test.py b/datalabeling/label_video_test.py index fd9bb846de37..e1c6f53b275d 100644 --- a/datalabeling/label_video_test.py +++ b/datalabeling/label_video_test.py @@ -17,7 +17,7 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServerError import pytest import label_video @@ -80,7 +80,7 @@ def test_label_video( capsys, annotation_spec_set, instruction, dataset, cleaner): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): # Start labeling. return label_video.label_video( diff --git a/datalabeling/manage_dataset_test.py b/datalabeling/manage_dataset_test.py index 13de03a73058..6699425ae61b 100644 --- a/datalabeling/manage_dataset_test.py +++ b/datalabeling/manage_dataset_test.py @@ -17,8 +17,8 @@ import os import backoff -from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import RetryError +from google.api_core.exceptions import ServerError import pytest import manage_dataset @@ -60,7 +60,7 @@ def cleaner(): def test_create_dataset(cleaner, capsys): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): return manage_dataset.create_dataset(PROJECT_ID) @@ -75,7 +75,7 @@ def run_sample(): def test_list_dataset(capsys, dataset): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): manage_dataset.list_datasets(PROJECT_ID) @@ -86,7 +86,7 @@ def run_sample(): def test_get_dataset(capsys, dataset): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): manage_dataset.get_dataset(dataset.name) @@ -97,7 +97,7 @@ def run_sample(): def test_delete_dataset(capsys, dataset): @backoff.on_exception( - backoff.expo, DeadlineExceeded, max_time=testing_lib.RETRY_DEADLINE) + backoff.expo, ServerError, max_time=testing_lib.RETRY_DEADLINE) def run_sample(): manage_dataset.delete_dataset(dataset.name)