Prevent indefinite OCR hangs with hard tesseract batch timeout

smoke_signal_tab.py

@@ -771,6 +771,10 @@ try:
     SS_TESSERACT_MIN_CONF_KEEP = max(0.0, min(1.0, float(os.environ.get("SS_TESSERACT_MIN_CONF_KEEP", "0.58"))))
 except Exception:
     SS_TESSERACT_MIN_CONF_KEEP = 0.58
+try:
+    SS_TESSERACT_BATCH_HARD_TIMEOUT_SEC = max(15, int(os.environ.get("SS_TESSERACT_BATCH_HARD_TIMEOUT_SEC", "90")))
+except Exception:
+    SS_TESSERACT_BATCH_HARD_TIMEOUT_SEC = 90
 
 
 def _load_surya_runtime():
@@ -1004,14 +1008,47 @@ def _run_tesseract_batch(images):
         return [_ocr_single(img) for img in images]
 
     outputs = [None] * len(images)
-    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-        futures = {executor.submit(_ocr_single, img): idx for idx, img in enumerate(images)}
-        for future in concurrent.futures.as_completed(futures):
-            idx = futures[future]
-            try:
-                outputs[idx] = future.result()
-            except Exception as e:
-                outputs[idx] = {"regions": [], "confidence": 0.0, "method": f"error-tesseract-future ({e})"}
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+    futures = {executor.submit(_ocr_single, img): idx for idx, img in enumerate(images)}
+    pending = set(futures.keys())
+    batch_timeout = max(
+        SS_TESSERACT_BATCH_HARD_TIMEOUT_SEC,
+        int((len(images) / max(max_workers, 1)) * SS_TESSERACT_TIMEOUT_SEC * 2 + 15),
+    )
+    deadline = time.monotonic() + batch_timeout
+
+    try:
+        while pending and time.monotonic() < deadline:
+            just_done, pending = concurrent.futures.wait(
+                pending,
+                timeout=0.75,
+                return_when=concurrent.futures.FIRST_COMPLETED,
+            )
+            if not just_done:
+                continue
+            for future in just_done:
+                idx = futures[future]
+                try:
+                    outputs[idx] = future.result()
+                except Exception as e:
+                    outputs[idx] = {"regions": [], "confidence": 0.0, "method": f"error-tesseract-future ({e})"}
+
+        if pending:
+            for future in pending:
+                idx = futures[future]
+                outputs[idx] = {
+                    "regions": [],
+                    "confidence": 0.0,
+                    "method": f"error-tesseract-batch-timeout ({batch_timeout}s)",
+                }
+                future.cancel()
+    finally:
+        # Avoid blocking the OCR run forever if one worker hangs in external OCR process.
+        executor.shutdown(wait=False, cancel_futures=True)
+
+    for i, out in enumerate(outputs):
+        if out is None:
+            outputs[i] = {"regions": [], "confidence": 0.0, "method": "error-tesseract-missing-output"}
 
     return outputs
 
@@ -1230,6 +1267,7 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
                             )
                     else:
                         fallback_preds = _run_tesseract_batch(batch_images)
+                        timeout_count = sum(1 for pred in fallback_preds if "batch-timeout" in str(pred.get("method", "")))
                         for item, pred in zip(batch_items_with_images, fallback_preds):
                             ocr_lookup[item["page_num"]] = pred
                             ocr_processed_pages += 1
@@ -1237,11 +1275,19 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
                                 (total_pages_planned + ocr_processed_pages, total_work_units),
                                 desc=f"Tesseract {book_id} p{item['page_num']}",
                             )
+                        if timeout_count:
+                            log.append(
+                                log_line(
+                                    f"  ⚠ {book_id} batch {batch_pages[0]}-{batch_pages[-1]}: "
+                                    f"{timeout_count}/{len(fallback_preds)} tesseract timeouts"
+                                )
+                            )
                 except Exception as e:
                     # If Surya batch fails, try Tesseract for this batch before giving up.
                     if surya is not None:
                         log.append(log_line(f"  ⚠ {book_id} batch {batch_pages[0]}-{batch_pages[-1]} Surya error: {e}; retrying with Tesseract"))
                         fallback_preds = _run_tesseract_batch(batch_images)
+                        timeout_count = sum(1 for pred in fallback_preds if "batch-timeout" in str(pred.get("method", "")))
                         for item, pred in zip(batch_items_with_images, fallback_preds):
                             ocr_lookup[item["page_num"]] = pred
                             ocr_processed_pages += 1
@@ -1249,6 +1295,13 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
                                 (total_pages_planned + ocr_processed_pages, total_work_units),
                                 desc=f"Tesseract retry {book_id} p{item['page_num']}",
                             )
+                        if timeout_count:
+                            log.append(
+                                log_line(
+                                    f"  ⚠ {book_id} batch {batch_pages[0]}-{batch_pages[-1]} retry: "
+                                    f"{timeout_count}/{len(fallback_preds)} tesseract timeouts"
+                                )
+                            )
                     else:
                         for item in batch_items_with_images:
                             ocr_lookup[item["page_num"]] = {