solver_registry: check file size after onnx.save, before validate — skip oversized immediately

neurogolf_solver/solvers/solver_registry.py

@@ -15,7 +15,7 @@ from .mode import s_mode_fill
 from .conv import solve_conv_fixed, solve_conv_variable, solve_conv_diffshape, solve_conv_var_diff
 from ..data_loader import get_exs, fixed_shapes
 from ..validators import validate
-from ..constants import EXCLUDED_TASKS
+from ..constants import EXCLUDED_TASKS, MAX_ONNX_FILESIZE
 
 # Analytical solvers registry — order matters (cheaper first)
 ANALYTICAL_SOLVERS = [
@@ -45,6 +45,14 @@ ANALYTICAL_SOLVERS = [
 ]
 
 
+def _check_size(path):
+    """Return True if file is within 1.44MB limit."""
+    try:
+        return os.path.getsize(path) <= MAX_ONNX_FILESIZE
+    except OSError:
+        return False
+
+
 def solve_task(tn, td, outdir, providers, conv_budget=30.0, excluded_tasks=None):
     """Solve a single ARC-AGI task.
     
@@ -67,6 +75,8 @@ def solve_task(tn, td, outdir, providers, conv_budget=30.0, excluded_tasks=None)
             if model is None:
                 continue
             onnx.save(model, path)
+            if not _check_size(path):
+                continue  # oversized, skip to next solver
             if validate(path, td, providers):
                 return True, sname, os.path.getsize(path), time.time() - t_start, path
         except:
@@ -84,12 +94,14 @@ def solve_task(tn, td, outdir, providers, conv_budget=30.0, excluded_tasks=None)
         if fixed_in:
             result = solve_conv_fixed(td, path, providers, time_budget=conv_time / 2)
             if result is not None:
-                sname, model = result
-                return True, sname, os.path.getsize(path), time.time() - t_start, path
+                if _check_size(path):
+                    sname, model = result
+                    return True, sname, os.path.getsize(path), time.time() - t_start, path
         result = solve_conv_variable(td, path, providers, time_budget=conv_time)
         if result is not None:
-            sname, model = result
-            return True, sname, os.path.getsize(path), time.time() - t_start, path
+            if _check_size(path):
+                sname, model = result
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
     else:
         sp = fixed_shapes(td)
         if sp is not None:
@@ -97,12 +109,14 @@ def solve_task(tn, td, outdir, providers, conv_budget=30.0, excluded_tasks=None)
             if OH <= IH and OW <= IW:
                 result = solve_conv_diffshape(td, path, providers, time_budget=conv_time)
                 if result is not None:
-                    sname, model = result
-                    return True, sname, os.path.getsize(path), time.time() - t_start, path
+                    if _check_size(path):
+                        sname, model = result
+                        return True, sname, os.path.getsize(path), time.time() - t_start, path
 
         result = solve_conv_var_diff(td, path, providers, time_budget=conv_time)
         if result is not None:
-            sname, model = result
-            return True, sname, os.path.getsize(path), time.time() - t_start, path
+            if _check_size(path):
+                sname, model = result
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
 
     return False, None, None, time.time() - t_start, path