captchaforge 0.2.33

#!/usr/bin/env python3
"""End-to-end VLM training pipeline DRY RUN.

Runs the full pipeline (capture → label → train → distill → promote)
in a mode that does NOT need GPU / teacher API / Ollama. Each stage
produces a real on-disk artifact (synthetic for the stages that
would normally hit external services). Lets CI prove the pipeline
plumbing stays connected end-to-end, separately from the heavy
real-training run.

Usage:
    python3 training/scripts/dry_run.py --out /tmp/captchaforge-vlm-dry-run

The script exits non-zero on any stage failure. Stage outputs are
under <out>/:
    capture/    — synthetic captured TrainingSamples (5 entries)
    labeled.jsonl — teacher-label simulation (5 entries)
    7b_lora.bin — synthetic 7B-LoRA checkpoint (small text file)
    2b_distilled.bin — synthetic 2B-distilled checkpoint
    manifest.json — pipeline summary

The same stage layout the real pipeline produces, so a real run
swaps in the heavy-compute stages without touching the orchestration.
"""

import argparse
import hashlib
import json
import sys
import time
from pathlib import Path


def stage_capture(out_root: Path, n: int = 5) -> Path:
    """Emit n synthetic TrainingSample JSONL rows."""
    capture_dir = out_root / "capture"
    capture_dir.mkdir(parents=True, exist_ok=True)
    path = capture_dir / "synthetic.jsonl"
    rows = []
    for i in range(n):
        rows.append({
            "solver": "VlmCaptchaSolver",
            "vendor": "hcaptcha",
            "detected_kind": "HCaptcha",
            "url": f"https://example.com/page-{i}",
            "outcome": "failure" if i % 2 == 0 else "success",
            "confidence": 0.5,
            "time_ms": 1234,
            "screenshot_b64": None,
            "dom_snapshot": None,
            "verified_outcome": None,
            "captured_at_unix": 1_700_000_000 + i,
        })
    with path.open("w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    print(f"[capture] {n} synthetic samples → {path}")
    return path


def stage_label(capture_path: Path, out_root: Path) -> Path:
    """Apply a deterministic 'teacher' label to every capture row."""
    labeled = out_root / "labeled.jsonl"
    n = 0
    with capture_path.open() as src, labeled.open("w") as dst:
        for line in src:
            row = json.loads(line)
            # Deterministic label: hash-derived answer.
            label = hashlib.sha256(row["url"].encode()).hexdigest()[:8]
            row["teacher_label"] = label
            dst.write(json.dumps(row) + "\n")
            n += 1
    print(f"[label] {n} rows labelled → {labeled}")
    return labeled


def stage_train_lora(labeled: Path, out_root: Path) -> Path:
    """Synthesise a 7B-LoRA checkpoint without touching a GPU."""
    ckpt = out_root / "7b_lora.bin"
    with labeled.open() as f:
        body = f.read()
    digest = hashlib.sha256(body.encode()).hexdigest()
    payload = (
        "captchaforge-vlm-7b-lora dry-run\n"
        f"trained_at: {int(time.time())}\n"
        f"data_digest: {digest}\n"
        f"data_lines: {body.count(chr(10))}\n"
    )
    ckpt.write_text(payload)
    print(f"[train_lora] synthetic 7B checkpoint → {ckpt}")
    return ckpt


def stage_distill(teacher_ckpt: Path, out_root: Path) -> Path:
    """Synthesise a 2B-distilled checkpoint from the 7B teacher."""
    ckpt = out_root / "2b_distilled.bin"
    teacher_digest = hashlib.sha256(teacher_ckpt.read_bytes()).hexdigest()
    payload = (
        "captchaforge-vlm-2b dry-run\n"
        f"distilled_at: {int(time.time())}\n"
        f"teacher_digest: {teacher_digest}\n"
    )
    ckpt.write_text(payload)
    print(f"[distill] synthetic 2B checkpoint → {ckpt}")
    return ckpt


def stage_promote(distilled_ckpt: Path, out_root: Path) -> Path:
    """Write the Ollama-promotion manifest (no actual ollama push)."""
    manifest = out_root / "ollama_manifest.json"
    manifest.write_text(json.dumps({
        "model": "captchaforge/captchaforge-vlm:2b-dry-run",
        "source": str(distilled_ckpt),
        "would_push_to": "ollama.com/captchaforge/captchaforge-vlm",
        "dry_run": True,
    }, indent=2))
    print(f"[promote] manifest → {manifest}")
    return manifest


def write_summary(
    out_root: Path,
    capture: Path,
    labeled: Path,
    ckpt7b: Path,
    ckpt2b: Path,
    manifest: Path,
) -> Path:
    summary = out_root / "manifest.json"
    summary.write_text(json.dumps({
        "pipeline": "captchaforge VLM training (dry-run)",
        "stages": {
            "capture": str(capture),
            "label": str(labeled),
            "train_lora": str(ckpt7b),
            "distill": str(ckpt2b),
            "promote": str(manifest),
        },
        "completed_at_unix": int(time.time()),
    }, indent=2) + "\n")
    print(f"[summary] {summary}")
    return summary


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", required=True, help="Output root directory")
    args = parser.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    print(f"captchaforge VLM pipeline dry-run → {out}")
    capture = stage_capture(out)
    labeled = stage_label(capture, out)
    ckpt7b = stage_train_lora(labeled, out)
    ckpt2b = stage_distill(ckpt7b, out)
    manifest = stage_promote(ckpt2b, out)
    summary = write_summary(out, capture, labeled, ckpt7b, ckpt2b, manifest)
    print(f"DONE — every stage produced an artifact under {out}")
    # Sanity assertions before exit.
    for required in (capture, labeled, ckpt7b, ckpt2b, manifest, summary):
        if not required.is_file():
            print(f"ERROR: stage artifact missing: {required}", file=sys.stderr)
            return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())