ppai/image-classification: add exponential retries to download file (GoogleCloudPlatform#6668)

David Cavazos · web-flow · commit a5c2f3ecc9be · 2021-09-08T15:58:26.000Z
## Description Fixes GoogleCloudPlatform#6463 Even though the `Connection reset by peer` is a transient error and we can't do much to solve it on our side, we retry the request with exponential backoff and hopefully it'll work at some point. This will hopefully decrease the number of flaky test runs. ## Checklist - [ ] I have followed [Sample Guidelines from AUTHORING_GUIDE.MD](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md) - [ ] README is updated to include [all relevant information](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md#readme-file) - [ ] **Tests** pass: `nox -s py-3.6` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] **Lint** pass: `nox -s lint` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] These samples need a new **API enabled** in testing projects to pass (let us know which ones) - [ ] These samples need a new/updated **env vars** in testing projects set to pass (let us know which ones) - [ ] Please **merge** this PR for me once it is approved. - [ ] This sample adds a new sample directory, and I updated the [CODEOWNERS file](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/.github/CODEOWNERS) with the codeowners for this sample
diff --git a/people-and-planet-ai/image-classification/predict.py b/people-and-planet-ai/image-classification/predict.py
@@ -21,6 +21,8 @@
 from google.cloud.aiplatform.gapic.schema import predict
 import requests
 
+from train_model import with_retries
+
 
 def run(
     project: str, region: str, model_endpoint_id: str, image_file: str
@@ -43,7 +45,7 @@ def run(
     )
 
     base_url = "https://lilablobssc.blob.core.windows.net/wcs-unzipped"
-    image_bytes = requests.get(f"{base_url}/{image_file}").content
+    image_bytes = with_retries(lambda: requests.get(f"{base_url}/{image_file}").content)
 
     response = client.predict(
         endpoint=client.endpoint_path(
