Merge pull request #855 from Labelbox/ms/retry-connection-timeouts

retry connection timeouts

Author: msokoloff1

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+# Version 3.37.0 (2023-xx-xx)
+## Fixed
+* Retry connection timeouts
+
 # Version 3.36.1 (2023-01-24)
 ### Fixed
 * `confidence` is now optional for TextEntity

labelbox/client.py

@@ -97,7 +97,8 @@ def __init__(self,
         self._data_row_metadata_ontology = None
 
     @retry.Retry(predicate=retry.if_exception_type(
-        labelbox.exceptions.InternalServerError))
+        labelbox.exceptions.InternalServerError,
+        labelbox.exceptions.TimeoutError))
     def execute(self,
                 query=None,
                 params=None,

tests/integration/test_client_errors.py

@@ -2,6 +2,7 @@
 import os
 import time
 import pytest
+from google.api_core.exceptions import RetryError
 
 from labelbox import Project, Dataset, User
 import labelbox.client
@@ -43,12 +44,12 @@ def test_semantic_error(client):
 
 
 def test_timeout_error(client, project):
-    with pytest.raises(labelbox.exceptions.TimeoutError) as excinfo:
-        query_str = """query getOntology { 
-        project (where: {id: %s}) { 
-            ontology { 
-                normalized 
-                } 
+    with pytest.raises(RetryError) as excinfo:
+        query_str = """query getOntology {
+        project (where: {id: %s}) {
+            ontology {
+                normalized
+                }
             }
         } """ % (project.uid)