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"
+                    f"Error processing {image_file.name} "
-                    "\tIt will be retried."
+                    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:
+        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()
-    return ExitCode.PARTIAL if warnings else ExitCode.SUCCESS
+                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:

copienator/commands/submit_batches.py

@@ -19,11 +19,7 @@ from copienator import (
 )
 
 
-def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
+def run(workspace: EvaluationWorkspace, *, client=None, overwrite: bool = False) -> ExitCode:
-    if client is None:
-        if not config.API_KEY:
-            raise CliError("GEMINI_API_KEY is not configured")
-        client = genai.Client(api_key=config.API_KEY)
     batches = (
         (
             "flash",
@@ -45,16 +41,32 @@ def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
     }
     if workspace.batch_jobs_file.is_file():
         previous = read_json(workspace.batch_jobs_file)
-        if isinstance(previous, dict) and isinstance(previous.get("jobs"), dict):
+        if not isinstance(previous, dict) or not isinstance(previous.get("jobs"), dict):
+            raise CliError(f"Invalid batch manifest: {workspace.batch_jobs_file}")
         manifest["jobs"] = previous["jobs"]
     started = 0
+    skipped = 0
     for tier, file_path, model_id, display_name in batches:
+        existing = manifest["jobs"].get(tier)
+        if tier in manifest["jobs"]:
+            if (not isinstance(existing, dict)
+                    or not isinstance(existing.get("name"), str)
+                    or not existing["name"].strip()):
+                raise CliError(f"Invalid batch job in {workspace.batch_jobs_file}: {tier}")
+            if not overwrite:
+                print(f"Skipping {tier}: already submitted as {existing['name']}.")
+                skipped += 1
+                continue
         if not file_path.is_file():
             print(f"Skipping {model_id}: {file_path.name} does not exist.")
             continue
         if file_path.stat().st_size == 0:
             print(f"Skipping {model_id}: {file_path.name} is empty.")
             continue
+        if client is None:
+            if not config.API_KEY:
+                raise CliError("GEMINI_API_KEY is not configured")
+            client = genai.Client(api_key=config.API_KEY)
         print(f"Uploading {file_path.name} for model {model_id}...")
         uploaded = client.files.upload(
             file=str(file_path),
@@ -77,19 +89,26 @@ def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
         }
         atomic_write_json(workspace.batch_jobs_file, manifest)
         print(f"Started batch job: {job.name}")
-    if not started:
+    if not started and not skipped:
         print("No non-empty batch request files were found.")
         return ExitCode.PARTIAL
     return ExitCode.SUCCESS
 
 
 def build_parser() -> argparse.ArgumentParser:
-    return evaluation_parser("Upload correction JSONL files and start Gemini batches")
+    parser = evaluation_parser("Upload correction JSONL files and start Gemini batches")
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Submit new jobs even for recorded tiers, replacing their saved job IDs",
+    )
+    return parser
 
 
 def main(argv: Sequence[str] | None = None) -> int:
     parser = build_parser()
-    return execute(parser, argv, lambda args: run(workspace_from_args(args)))
+    return execute(parser, argv, lambda args: run(
+        workspace_from_args(args), overwrite=args.overwrite,
+    ))
 
 
 if __name__ == "__main__":