Standardisation 8

submit_batches.py

@@ -1,77 +1,96 @@
-import os
-import sys
+from __future__ import annotations
+
 import argparse
-from pathlib import Path
+from collections.abc import Sequence
+
 from google import genai
 from google.genai import types
 
-def main():
-    parser = argparse.ArgumentParser(description="Upload JSONL files and create Gemini Batch jobs.")
-    parser.add_argument("root_dir", type=str, help="Root directory containing the batch JSONL files")
-    args = parser.parse_args()
+import config
+from copienator import (
+    CliError,
+    EvaluationWorkspace,
+    ExitCode,
+    atomic_write_json,
+    evaluation_parser,
+    execute,
+    read_json,
+    workspace_from_args,
+)
 
-    root_dir = Path(args.root_dir)
 
-    if "GEMINI_API_KEY" not in os.environ:
-        sys.exit("Error: GEMINI_API_KEY environment variable not set.")
-
-    client = genai.Client()
-
-    # Define the batch files and their corresponding models
-    batches_to_create = [
-        {
-            "file_path": root_dir / "batch_requests_flash.jsonl",
-            "model_id": "gemini-3-flash-preview",
-            "display_name": f"flash-correction-{root_dir.name}"
-        },
-        {
-            "file_path": root_dir / "batch_requests_pro.jsonl",
-            "model_id": "gemini-3.1-pro-preview",
-            "display_name": f"pro-correction-{root_dir.name}"
-        }
-    ]
-
-    for batch in batches_to_create:
-        file_path = batch["file_path"]
-        model_id = batch["model_id"]
-        display_name = batch["display_name"]
-
-        # Check if the file exists
-        if not file_path.exists():
+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)
+    batches = (
+        (
+            "flash",
+            workspace.root / "batch_requests_flash.jsonl",
+            config.MODEL_FLASH_ID,
+            f"flash-correction-{workspace.name}",
+        ),
+        (
+            "pro",
+            workspace.root / "batch_requests_pro.jsonl",
+            config.MODEL_PRO_ID,
+            f"pro-correction-{workspace.name}",
+        ),
+    )
+    manifest = {
+        "version": 1,
+        "evaluation": workspace.name,
+        "jobs": {},
+    }
+    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"]
+    started = 0
+    for tier, file_path, model_id, display_name in batches:
+        if not file_path.is_file():
             print(f"Skipping {model_id}: {file_path.name} does not exist.")
             continue
-
-        # Check if the file is empty (e.g., if all tasks went to Flash, Pro might be empty)
         if file_path.stat().st_size == 0:
             print(f"Skipping {model_id}: {file_path.name} is empty.")
             continue
-
-        print(f"Processing {file_path.name} for model {model_id}...")
-
-        # 1. Upload the file to the File API
-        print(f"  Uploading file...")
-        uploaded_file = client.files.upload(
+        print(f"Uploading {file_path.name} for model {model_id}...")
+        uploaded = client.files.upload(
             file=str(file_path),
             config=types.UploadFileConfig(
                 display_name=f"{display_name}-input",
-                mime_type='jsonl'
-            )
+                mime_type="jsonl",
+            ),
         )
-        print(f"  Uploaded successfully! File ID: {uploaded_file.name}")
-
-        # 2. Create the batch job
-        print(f"  Starting batch job...")
-        batch_job = client.batches.create(
+        job = client.batches.create(
             model=model_id,
-            src=uploaded_file.name,
-            config={
-                'display_name': display_name,
-            },
+            src=uploaded.name,
+            config={"display_name": display_name},
         )
-        print(f"  Success! Batch Job Name: {batch_job.name}\n")
+        started += 1
+        manifest["jobs"][tier] = {
+            "name": job.name,
+            "display_name": display_name,
+            "model": model_id,
+            "request_file": file_path.name,
+        }
+        atomic_write_json(workspace.batch_jobs_file, manifest)
+        print(f"Started batch job: {job.name}")
+    if not started:
+        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")
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = build_parser()
+    return execute(parser, argv, lambda args: run(workspace_from_args(args)))
 
-    print("-" * 50)
-    print("All batch jobs have been initiated.")
 
 if __name__ == "__main__":
-    main()
+    raise SystemExit(main())