Retry failed ZeroGPU allocations

app.py

@@ -4,6 +4,7 @@ import secrets
 import shutil
 import sys
 import threading
+import time
 import urllib.parse
 import urllib.request
 import uuid
@@ -125,7 +126,11 @@ DEFAULT_SECOND_SCHEDULER = "karras"
 DEFAULT_SECOND_STEPS = 18
 DEFAULT_SECOND_CFG = 5
 DEFAULT_DENOISE = .5
-GPU_ABORT_ERROR = "GPU task aborted"
+GPU_ATTEMPTS = 3
+GPU_RETRY_ERRORS = (
+    "GPU task aborted",
+    "uncorrectable ECC error",
+)
 
 jobs = {}
 images = {}
@@ -394,6 +399,22 @@ def generate_core(body):
         return image[0], first_seed, second_seed
 
 
+def generate_gpu(body):
+    for attempt in range(GPU_ATTEMPTS):
+        try:
+            return generate_core(body)
+        except Exception as error:
+            if (
+                attempt == GPU_ATTEMPTS - 1
+                or not any(text in str(error) for text in GPU_RETRY_ERRORS)
+            ):
+                raise
+
+            delay = 2**attempt
+            log(f"GPU allocation failed, retrying in {delay}s")
+            time.sleep(delay)
+
+
 def load_model(name):
     if name not in state["models"]:
         stage_model("checkpoints", name)
@@ -522,7 +543,7 @@ def generate_image(request, tag="ST"):
         if request.upscale and request.second_model:
             load_chain(request.second_model, request.second_loras)
 
-        image, first_seed, second_seed = generate_core(
+        image, first_seed, second_seed = generate_gpu(
             request.model_dump_json()
         )
 
@@ -643,18 +664,12 @@ def api_generate(body, p=""):
 
 def queued_generate(request):
     client = Client(LOCAL_URL, verbose=False)
-
-    for attempt in range(2):
-        try:
-            result = client.predict(
-                request.model_dump_json(),
-                PASSWORD,
-                api_name="/generate",
-            )
-            return Image.open(result).copy()
-        except Exception as error:
-            if attempt or GPU_ABORT_ERROR not in str(error):
-                raise
+    result = client.predict(
+        request.model_dump_json(),
+        PASSWORD,
+        api_name="/generate",
+    )
+    return Image.open(result).copy()
 
 
 def linked_value(workflow, value):