v5 refactor: add submission.py (with W&B logging)

neurogolf_solver/submission.py

@@ -0,0 +1,134 @@
+#!/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_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
+        if tn in excluded_tasks:
+            print(f"Task {tn:3d}: EXCLUDED (officially)")
+            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
+
+            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'))
+
+    # 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]])
+
+    # Calculate estimated LB score
+    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())
+    est_lb = total_score + unsolved_count * 1.0
+
+    return {
+        'n_files': n_files,
+        'total_size': total_size,
+        'zip_path': zip_path,
+        'zip_size': zip_size,
+        'csv_path': csv_path,
+        'est_lb': est_lb,
+        'total_score': total_score,
+        'unsolved_count': unsolved_count,
+    }
+
+
+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} active 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")
+    zip_ok = 'OK' if submission_info['zip_size'] <= MAX_FILESIZE else 'OVER!'
+    print(f"ZIP size: {submission_info['zip_size'] / 1024:.1f} KB / {MAX_FILESIZE / 1024:.0f} KB limit {zip_ok}")
+    print(f"Estimated LB score: {submission_info['est_lb']:.1f} "
+          f"(solved: {submission_info['total_score']:.1f} + unsolved: {submission_info['unsolved_count']}x1.0)")
+    print(f"Written: {submission_info['zip_path']} | {submission_info['csv_path']}")