submit_empathy_uv.py

#!/usr/bin/env python3
"""Submit the empathy-only LoRA retrain to HF Jobs (uv), like last time.

Two explicit steps so nothing spends credits by accident:

  1) Upload the new v2 dataset (cheap, backs up the old templated file first):
        python submit_empathy_uv.py --upload-data

  2) Submit the GPU training job (THIS spends credits):
        python submit_empathy_uv.py --submit

Run with no flags to see the plan (dry run). Requires HF_TOKEN in the env.
"""

import argparse, json, os, sys
from pathlib import Path
from huggingface_hub import HfApi

MODEL_REPO   = "Raiff1982/codette-llama-3.1-8b-merged"
DATASET_REPO = "Raiff1982/codette-training-data"
OUTPUT_REPO  = "Raiff1982/codette-lora-adapters"
LOCAL_DATA   = Path("dataset_engine/empathy_reasoning.jsonl")   # generated by empathy_dataset_v2.py
HUB_DATA     = "empathy_reasoning.jsonl"
HUB_BACKUP   = "empathy_reasoning_templated_v1.jsonl"
SCRIPT       = "training/train_empathy_uv.py"
FLAVOR       = "a10g-large"
TIMEOUT      = "1h"


def _token() -> str:
    tok = os.environ.get("HF_TOKEN")
    if not tok:
        try:
            from huggingface_hub import get_token
            tok = get_token()
        except Exception:
            tok = None
    if not tok:
        print("ERROR: no HF token (set HF_TOKEN or run `hf auth login`).")
        sys.exit(1)
    return tok


def upload_data(api: HfApi, token: str) -> None:
    if not LOCAL_DATA.exists():
        print(f"ERROR: {LOCAL_DATA} not found. Generate it first:\n"
              f"  python -m dataset_engine.empathy_dataset_v2 --out {LOCAL_DATA}")
        sys.exit(1)

    # Back up the existing templated v1 once (don't clobber an existing backup).
    try:
        existing = {f for f in api.list_repo_files(DATASET_REPO, repo_type="dataset", token=token)}
    except Exception as e:
        print(f"Could not list dataset repo files: {e}")
        existing = set()

    if HUB_DATA in existing and HUB_BACKUP not in existing:
        from huggingface_hub import hf_hub_download
        old = hf_hub_download(DATASET_REPO, HUB_DATA, repo_type="dataset", token=token)
        api.upload_file(path_or_fileobj=old, path_in_repo=HUB_BACKUP,
                        repo_id=DATASET_REPO, repo_type="dataset", token=token)
        print(f"Backed up templated v1 -> {DATASET_REPO}/{HUB_BACKUP}")
    elif HUB_BACKUP in existing:
        print(f"Backup already exists ({HUB_BACKUP}); not overwriting it.")

    n = sum(1 for _ in LOCAL_DATA.open(encoding="utf-8"))
    api.upload_file(path_or_fileobj=str(LOCAL_DATA), path_in_repo=HUB_DATA,
                    repo_id=DATASET_REPO, repo_type="dataset", token=token)
    print(f"Uploaded v2 dataset ({n} examples) -> {DATASET_REPO}/{HUB_DATA}")


def submit(api: HfApi, token: str) -> None:
    print(f"Submitting uv job: script={SCRIPT} flavor={FLAVOR} timeout={TIMEOUT}")
    job = api.run_uv_job(
        script=SCRIPT,
        flavor=FLAVOR,
        timeout=TIMEOUT,
        env={"HF_TOKEN": token},
        token=token,
    )
    print("\n[SUCCESS] Job submitted")
    print(f"  Job ID: {job.id}")
    print(f"  Status: {job.status}")
    print(f"  URL:    {job.url}")
    print(f"  Adapter will upload to: {OUTPUT_REPO}/empathy_v2 (verify, then promote to /empathy)")
    Path("empathy_job_info.json").write_text(json.dumps(
        {"job_id": job.id, "status": str(job.status), "url": job.url,
         "output_repo": OUTPUT_REPO, "adapter_path": "empathy_v2",
         "flavor": FLAVOR, "script": SCRIPT}, indent=2))
    print("  Saved empathy_job_info.json")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--upload-data", action="store_true", help="Back up v1 + upload v2 dataset")
    ap.add_argument("--submit", action="store_true", help="Submit the GPU training job (spends credits)")
    args = ap.parse_args()

    token = _token()
    api = HfApi(token=token)

    if args.upload_data:
        upload_data(api, token)
    if args.submit:
        submit(api, token)
    if not (args.upload_data or args.submit):
        print("DRY RUN — nothing submitted. Plan:")
        print(f"  Base model : {MODEL_REPO}")
        print(f"  Dataset    : {DATASET_REPO}/{HUB_DATA}  (local: {LOCAL_DATA})")
        print(f"  Script     : {SCRIPT}")
        print(f"  Flavor     : {FLAVOR}  (GPU — spends credits)")
        print(f"  Output     : {OUTPUT_REPO}/empathy_v2")
        print("\nStep 1 (cheap):  python submit_empathy_uv.py --upload-data")
        print("Step 2 (credits): python submit_empathy_uv.py --submit")


if __name__ == "__main__":
    main()