main.py: add --strict_size (default True) — stops entire run if any .onnx > 1.44MB

neurogolf_solver/main.py

@@ -9,12 +9,14 @@ Usage:
 """
 
 import argparse
+import os
+import sys
 import time
 import onnxruntime as ort
 from .config import get_providers
 from .data_loader import load_tasks_dir, load_tasks_kaggle
 from .submission import run_tasks, generate_submission, print_summary
-from .constants import EXCLUDED_TASKS
+from .constants import EXCLUDED_TASKS, MAX_ONNX_FILESIZE
 
 try:
     import wandb
@@ -22,6 +24,29 @@ except ImportError:
     wandb = None
 
 
+def check_output_sizes(output_dir, strict):
+    """Check all .onnx files are within 1.44MB limit.
+    If strict=True and any file is over, print error and exit immediately."""
+    oversized = []
+    for f in sorted(os.listdir(output_dir)):
+        if f.endswith('.onnx'):
+            fpath = os.path.join(output_dir, f)
+            fsize = os.path.getsize(fpath)
+            if fsize > MAX_ONNX_FILESIZE:
+                oversized.append((f, fsize))
+    if oversized:
+        print(f"\n{'!'*70}")
+        print(f"FATAL: {len(oversized)} .onnx files exceed 1.44MB limit:")
+        for f, sz in oversized:
+            print(f"  {f}: {sz:,} bytes ({sz/1024:.1f} KB) — OVER by {sz - MAX_ONNX_FILESIZE:,} bytes")
+        print(f"Kaggle WILL reject this submission.")
+        print(f"{'!'*70}")
+        if strict:
+            print("Stopping (--strict_size is on). Fix oversized models before submitting.")
+            sys.exit(1)
+    return oversized
+
+
 def main():
     parser = argparse.ArgumentParser(description='NeuroGolf Solver v5')
     parser.add_argument('--data_dir', default='ARC-AGI/data/training/')
@@ -32,6 +57,7 @@ def main():
     parser.add_argument('--tasks', type=str, default='', help='Comma-separated task numbers')
     parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cpu', 'cuda'])
     parser.add_argument('--use_wandb', action='store_true', help='Enable W&B logging')
+    parser.add_argument('--strict_size', type=bool, default=True, help='Stop run if any .onnx > 1.44MB (default: True)')
     args = parser.parse_args()
 
     providers = get_providers(args.device)
@@ -46,6 +72,7 @@ def main():
 
     ort.set_default_logger_severity(3)
     print(f"Using providers: {providers}")
+    print(f"Strict size check: {args.strict_size} (max {MAX_ONNX_FILESIZE:,} bytes per .onnx)")
 
     # Load tasks
     if args.kaggle:
@@ -78,6 +105,9 @@ def main():
 
     elapsed = time.time() - t0
 
+    # Check all output files BEFORE generating submission
+    check_output_sizes(args.output_dir, args.strict_size)
+
     submission_info = generate_submission(args.output_dir, results, costs_dict, task_nums)
     print_summary(results, submission_info, elapsed)