Fix 4: main.py — add --strict_score, stop if any model unscorable

neurogolf_solver/main.py

@@ -16,6 +16,7 @@ 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 .profiler import score_network
 from .constants import EXCLUDED_TASKS, MAX_ONNX_FILESIZE
 
 try:
@@ -24,27 +25,36 @@ 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 = []
+def check_all_models(output_dir, strict_size, strict_score):
+    """Check all .onnx files for size limit and scoreability.
+    If strict and any fail, exit immediately."""
+    problems = []
     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:
+        if not f.endswith('.onnx'):
+            continue
+        fpath = os.path.join(output_dir, f)
+        fsize = os.path.getsize(fpath)
+
+        if fsize > MAX_ONNX_FILESIZE:
+            problems.append((f, f"OVERSIZED: {fsize:,} bytes ({fsize/1024:.1f} KB) > {MAX_ONNX_FILESIZE:,}"))
+
+        macs, memory, params = score_network(fpath)
+        if macs is None or memory is None or params is None:
+            problems.append((f, "UNSCORABLE: score_network returned (None, None, None) — Kaggle will reject"))
+
+    if problems:
         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"PROBLEMS FOUND: {len(problems)} .onnx files will be REJECTED by Kaggle:")
+        for f, msg in problems:
+            print(f"  {f}: {msg}")
         print(f"{'!'*70}")
-        if strict:
-            print("Stopping (--strict_size is on). Fix oversized models before submitting.")
+        if strict_size or strict_score:
+            print("Stopping. Fix these models before submitting.")
             sys.exit(1)
-    return oversized
+    else:
+        print(f"\nAll .onnx files pass size and score checks.")
+
+    return problems
 
 
 def main():
@@ -57,7 +67,8 @@ 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)')
+    parser.add_argument('--strict_size', type=bool, default=True, help='Stop if any .onnx > 1.44MB (default: True)')
+    parser.add_argument('--strict_score', type=bool, default=True, help='Stop if any model unscorable (default: True)')
     args = parser.parse_args()
 
     providers = get_providers(args.device)
@@ -72,7 +83,8 @@ 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)")
+    print(f"Strict size: {args.strict_size} | Strict score: {args.strict_score}")
+    print(f"Max .onnx file size: {MAX_ONNX_FILESIZE:,} bytes")
 
     # Load tasks
     if args.kaggle:
@@ -106,7 +118,7 @@ def main():
     elapsed = time.time() - t0
 
     # Check all output files BEFORE generating submission
-    check_output_sizes(args.output_dir, args.strict_size)
+    check_all_models(args.output_dir, args.strict_size, args.strict_score)
 
     submission_info = generate_submission(args.output_dir, results, costs_dict, task_nums)
     print_summary(results, submission_info, elapsed)