From 55f9a296771e5e2fafa24c75fa9f37527dce75be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Miquel?= Date: Mon, 21 Sep 2026 15:40:33 +0200 Subject: [PATCH] Timeout for label detection --- copienator/commands/gemini_for_labels.py | 78 ++++++++++++++++++++---- copienator/commands/submit_batches.py | 39 +++++++++--- 2 files changed, 96 insertions(+), 21 deletions(-) diff --git a/copienator/commands/gemini_for_labels.py b/copienator/commands/gemini_for_labels.py index 28c44ff..e1ceeca 100644 --- a/copienator/commands/gemini_for_labels.py +++ b/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: diff --git a/copienator/commands/submit_batches.py b/copienator/commands/submit_batches.py index bcec9f9..b0293bc 100644 --- a/copienator/commands/submit_batches.py +++ b/copienator/commands/submit_batches.py @@ -19,11 +19,7 @@ from copienator import ( ) -def run(workspace: EvaluationWorkspace, *, client=None) -> 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) +def run(workspace: EvaluationWorkspace, *, client=None, overwrite: bool = False) -> ExitCode: 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): - manifest["jobs"] = previous["jobs"] + 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__":