more retry logic ig

neggles · Jan 29, 2024 · 8528cf1 · 8528cf1
1 parent 4d28dd1
commit 8528cf1
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 7 deletions.
diff --git a/src/neurosis/dataset/mongo/aspect.py b/src/neurosis/dataset/mongo/aspect.py
@@ -233,6 +233,8 @@ def __init__(
         s3_bucket: Optional[str] = None,
         s3fs_kwargs: dict = {},
         pma_schema: Optional[Schema] = None,
+        retries: int = 3,
+        retry_delay: int = 5,
         num_workers: int = 0,
         prefetch_factor: int = 2,
         pin_memory: bool = True,
@@ -259,7 +261,10 @@ def __init__(
             s3_bucket=s3_bucket,
             s3fs_kwargs=s3fs_kwargs,
             pma_schema=pma_schema,
+            retries=retries,
+            retry_delay=retry_delay,
         )
+
         self.num_workers = num_workers
         self.pin_memory = pin_memory
         self.prefetch_factor = prefetch_factor

diff --git a/src/neurosis/dataset/mongo/base.py b/src/neurosis/dataset/mongo/base.py
@@ -146,10 +146,12 @@ def _get_image(self, path: str) -> Image.Image:
         while attempts < self.retries and image is None:
             try:
                 image = self.fs.cat(path)
-            except ConnectionError:
+            except Exception as e:
                 logger.exception(f"Caught connection error on {path}, retry in {self.retry_delay}s")
-                sleep(self.retry_delay)
-            attempts += 1
+                sleep(self.retry_delay + attempts)
+                attempts += 1
+                if attempts >= self.retries:
+                    raise FileNotFoundError(f"Failed to load image from {path}") from e
 
         if not isinstance(image, bytes):
             raise FileNotFoundError(f"Failed to load image from {path}")

diff --git a/src/neurosis/dataset/mongo/nobucket.py b/src/neurosis/dataset/mongo/nobucket.py
@@ -150,7 +150,8 @@ def __init__(
         s3_bucket: Optional[str] = None,
         s3fs_kwargs: dict = {},
         pma_schema: Optional[Schema] = None,
-        seed: Optional[int] = None,
+        retries: int = 3,
+        retry_delay: int = 5,
         num_workers: int = 0,
         prefetch_factor: int = 2,
         pin_memory: bool = True,
@@ -176,8 +177,10 @@ def __init__(
             s3_bucket=s3_bucket,
             s3fs_kwargs=s3fs_kwargs,
             pma_schema=pma_schema,
+            retries=retries,
+            retry_delay=retry_delay,
         )
-        self.seed = seed
+
         self.num_workers = num_workers
         self.pin_memory = pin_memory
         self.prefetch_factor = prefetch_factor

diff --git a/src/neurosis/dataset/mongo/nocaption.py b/src/neurosis/dataset/mongo/nocaption.py
@@ -95,7 +95,8 @@ def __init__(
         s3_bucket: Optional[str] = None,
         s3fs_kwargs: dict = {},
         pma_schema: Optional[Schema] = None,
-        seed: Optional[int] = None,
+        retries: int = 3,
+        retry_delay: int = 5,
         num_workers: int = 0,
         prefetch_factor: int = 2,
         pin_memory: bool = True,
@@ -115,8 +116,10 @@ def __init__(
             s3_bucket=s3_bucket,
             s3fs_kwargs=s3fs_kwargs,
             pma_schema=pma_schema,
+            retries=retries,
+            retry_delay=retry_delay,
         )
-        self.seed = seed
+
         self.num_workers = num_workers
         self.pin_memory = pin_memory
         self.prefetch_factor = prefetch_factor