Move own-solver/neurogolf_solver/submission.py to own-solver/

own-solver/neurogolf_solver/submission.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+"""Submission file generation and task running with W&B logging."""
+
+import os
+import csv
+import io
+import math
+import zipfile
+from collections import Counter
+from .profiler import score_network
+from .constants import MAX_ONNX_FILESIZE, EXCLUDED_TASKS
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+def run_tasks(task_nums, tasks, output_dir, providers, conv_budget, excluded_tasks, use_wandb):
+    """Run all tasks and collect results.
+    
+    Returns: (results, costs_dict, total_score)
+    """
+    from .solvers.solver_registry import solve_task
+
+    results = {}
+    costs_dict = {}
+    total_score = 0
+
+    for tn in task_nums:
+        if tn not in tasks:
+            continue
+
+        td = tasks[tn]['data']
+        ok, sname, sz, t_task, model_path = solve_task(
+            tn, td, output_dir, providers, conv_budget, excluded_tasks
+        )
+
+        if ok:
+            macs, memory, params = score_network(model_path)
+            if macs is None:
+                macs, memory, params = 0, 0, 0
+            cost = macs + memory + params
+            score = max(1.0, 25.0 - math.log(max(1, cost)))
+            total_score += score
+
+            # Check per-file size limit
+            if sz and sz > MAX_ONNX_FILESIZE:
+                print(f"Task {tn:3d}: {sname:25s} OVER SIZE LIMIT ({sz:,} > {MAX_ONNX_FILESIZE:,})")
+                continue
+
+            results[tn] = (sname, t_task, sz)
+            costs_dict[tn] = cost
+            print(f"Task {tn:3d}: {sname:25s} {score:7.3f} {cost:>12} {t_task:7.3f}s  ({sz:>8,} bytes)")
+        else:
+            score = 0
+            cost = 0
+            print(f"Task {tn:3d}: UNSOLVED  {t_task:7.3f}s")
+
+        if use_wandb and wandb is not None:
+            wandb.log({
+                "task_id": tn,
+                "solver": sname if ok else "unsolved",
+                "onnx_bytes": sz if ok else 0,
+                "task_time_sec": t_task,
+                "cost": cost,
+                "score": score,
+            })
+
+    return results, costs_dict, total_score
+
+
+def generate_submission(output_dir, results, costs_dict, active_tasks):
+    """Generate submission.zip and submission.csv.
+    
+    Returns dict with submission info.
+    """
+    n_files = len([f for f in os.listdir(output_dir) if f.endswith('.onnx')])
+    total_size = sum(os.path.getsize(os.path.join(output_dir, f))
+                     for f in os.listdir(output_dir) if f.endswith('.onnx'))
+
+    # Check per-file size limits
+    oversized = []
+    for f in os.listdir(output_dir):
+        if f.endswith('.onnx'):
+            fsize = os.path.getsize(os.path.join(output_dir, f))
+            if fsize > MAX_ONNX_FILESIZE:
+                oversized.append((f, fsize))
+
+    # Create submission.zip
+    parent_dir = os.path.dirname(output_dir) or '/kaggle/working/'
+    zip_path = os.path.join(parent_dir, 'submission.zip')
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for f in sorted(os.listdir(output_dir)):
+            if f.endswith('.onnx'):
+                zf.write(os.path.join(output_dir, f), f)
+    zip_bytes = buf.getvalue()
+    with open(zip_path, 'wb') as f:
+        f.write(zip_bytes)
+    zip_size = len(zip_bytes)
+
+    # Create submission.csv
+    csv_path = os.path.join(parent_dir, 'submission.csv')
+    with open(csv_path, 'w', newline='') as f:
+        w = csv.writer(f)
+        w.writerow(['task_id', 'total_cost'])
+        for tn in sorted(costs_dict.keys()):
+            w.writerow([f'task{tn:03d}', costs_dict[tn]])
+
+    unsolved_count = len(active_tasks) - len(results)
+    total_score = sum(max(1.0, 25.0 - math.log(max(1, cost))) for cost in costs_dict.values())
+    total_cost = sum(costs_dict.values())
+
+    return {
+        'n_files': n_files,
+        'total_size': total_size,
+        'zip_path': zip_path,
+        'zip_size': zip_size,
+        'csv_path': csv_path,
+        'total_score': total_score,
+        'total_cost': total_cost,
+        'unsolved_count': unsolved_count,
+        'oversized': oversized,
+    }
+
+
+def print_summary(results, submission_info, elapsed):
+    """Print summary statistics."""
+    active_count = submission_info['unsolved_count'] + len(results)
+
+    print(f"\n{'=' * 70}")
+    print(f"Solved: {len(results)}/{active_count} tasks in {elapsed:.0f}s")
+    solver_names = [v[0] for v in results.values()]
+    sc = Counter(solver_names)
+    for s, c in sc.most_common():
+        print(f"  {s}: {c}")
+
+    print(f"\n{submission_info['n_files']} ONNX files, {submission_info['total_size'] / 1024:.1f} KB uncompressed")
+    print(f"ZIP size: {submission_info['zip_size'] / 1024:.1f} KB")
+
+    if submission_info['oversized']:
+        print(f"WARNING: {len(submission_info['oversized'])} files exceed 1.44MB limit:")
+        for f, sz in submission_info['oversized']:
+            print(f"  {f}: {sz / 1024:.1f} KB")
+
+    print(f"\nEstimated LB score: {submission_info['total_score']:.1f}")
+    print(f"Total cost: {submission_info['total_cost']:,}")
+    print(f"Solved: {len(results)} | Unsolved: {submission_info['unsolved_count']}")
+    print(f"Written: {submission_info['zip_path']} | {submission_info['csv_path']}")