Timeout for label detection

copienator/commands/gemini_for_labels.py

@@ -4,9 +4,10 @@ import argparse
 import re
 import time
 import typing
+from threading import Lock
 from collections import defaultdict
 from collections.abc import Callable, Sequence
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, wait
 from pathlib import Path
 
 from google import genai
@@ -271,6 +272,10 @@ def generate_request(
     ]
 
     generate_content_config = types.GenerateContentConfig(
+        http_options=types.HttpOptions(
+            timeout=120_000,
+            retry_options=types.HttpRetryOptions(attempts=1),
+        ),
         automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
         temperature=1.0,
         top_p=0.95,
@@ -282,9 +287,17 @@ def generate_request(
     return (contents, generate_content_config)
 
 TARGET_INTERVAL = 3.5
+MAX_ATTEMPTS = 3
+PROGRESS_INTERVAL = 30
 Sleep = Callable[[float], None]
 
 
+class IncompleteCopyError(Exception):
+    def __init__(self, images: list[Path]):
+        self.images = images
+        super().__init__(f"Label detection incomplete for {images[0].name}")
+
+
 def selected_images(
     workspace: EvaluationWorkspace,
     targets: list[Path],
@@ -395,11 +408,12 @@ def process_copy_group(
     overwrite: bool,
     sleep: Sleep = time.sleep,
     target_interval: float = TARGET_INTERVAL,
+    progress: Callable[[str, str, float], None] | None = None,
 ) -> int:
     """Process one student's image parts sequentially to preserve context."""
     accumulated_labels: list[str] = []
     generated = 0
-    for image_file in files:
+    for image_index, image_file in enumerate(files):
         started = time.monotonic()
         output_json = workspace.copies_dir / f"{image_file.stem}.json"
         if output_json.exists() and not overwrite:
@@ -414,9 +428,13 @@ def process_copy_group(
         attempt = 0
         label_retry_count = 0
         wrong_labels: list[str] = []
-        while True:
+        while attempt < MAX_ATTEMPTS:
+            if progress:
+                progress(group_key, f"{image_file.name}, attempt {attempt + 1}/{MAX_ATTEMPTS}", started)
             if attempt > 0:
+                print(f"[{group_key}] Retrying {image_file.name} in {10 * attempt}s...")
                 sleep(10 * attempt)
+            print(f"[{group_key}] {image_file.name}: attempt {attempt + 1}/{MAX_ATTEMPTS} (120s timeout).")
             try:
                 text_context_file, text_context = closest_text_context(
                     workspace, wrong_labels
@@ -486,15 +504,23 @@ def process_copy_group(
                 atomic_write_json(output_json, annotation.model_dump())
                 accumulated_labels.extend(box.label for box in annotation.list)
                 generated += 1
+                print(f"[{group_key}] Saved {image_file.name} after {time.monotonic() - started:.0f}s.")
                 break
             except KeyboardInterrupt:
                 raise
             except Exception as exc:  # noqa: BLE001 - remote API retry boundary
                 print(
-                    f"Error processing {image_file.name}: {exc}\n"
-                    "\tIt will be retried."
+                    f"Error processing {image_file.name} "
+                    f"(attempt {attempt + 1}/{MAX_ATTEMPTS}): {type(exc).__name__}: {exc}"
                 )
                 attempt += 1
+        else:
+            # Later pages need the context of this page; defer them until resume.
+            incomplete = [path for path in files[image_index:]
+                          if overwrite or not (workspace.copies_dir / f"{path.stem}.json").exists()]
+            print(f"[{group_key}] Failed after {MAX_ATTEMPTS} attempts; "
+                  "deferring remaining pages of this copy until the next run.")
+            raise IncompleteCopyError(incomplete)
         sleep(max(0.0, target_interval - (time.monotonic() - started)))
     return generated
 
@@ -528,8 +554,17 @@ def run(
         client = genai.Client(api_key=api_key)
 
     groups = group_images(images)
+    active: dict[str, tuple[str, float]] = {}
+    progress_lock = Lock()
+    incomplete: list[Path] = []
+    run_started = time.monotonic()
+
+    def report_progress(group: str, description: str, started: float) -> None:
+        with progress_lock:
+            active[group] = (description, started)
+
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        futures = [
+        futures = {
             executor.submit(
                 process_copy_group,
                 workspace,
@@ -542,12 +577,33 @@ def run(
                 valid_names=valid_names,
                 overwrite=overwrite,
                 sleep=sleep,
-            )
+                progress=report_progress,
+            ): group_key
             for group_key, files in groups.items()
-        ]
-        for future in futures:
-            future.result()
-    return ExitCode.PARTIAL if warnings else ExitCode.SUCCESS
+        }
+        pending = set(futures)
+        while pending:
+            done, pending = wait(pending, timeout=PROGRESS_INTERVAL)
+            for future in done:
+                with progress_lock:
+                    active.pop(futures[future], None)
+                try:
+                    future.result()
+                except IncompleteCopyError as exc:
+                    incomplete.extend(exc.images)
+            if pending:
+                now = time.monotonic()
+                with progress_lock:
+                    snapshot = list(active.items())
+                print(f"Label detection: {now - run_started:.0f}s elapsed; "
+                      f"{len(pending)} copies outstanding ({len(pending) - len(snapshot)} queued).")
+                for group, (description, started) in snapshot:
+                    print(f"  [{group}] {description}: {now - started:.0f}s elapsed for image.")
+    if incomplete:
+        print("Incomplete images (rerun label detection to resume):")
+        for path in sorted(incomplete):
+            print(f"  {path.name}")
+    return ExitCode.PARTIAL if warnings or incomplete else ExitCode.SUCCESS
 
 
 def build_parser() -> argparse.ArgumentParser: