training/auto_upload.py

"""
Auto-upload artifacts to HuggingFace Hub after the training pipeline finishes.

Called from ``entrypoint.sh``. Runs AT THE END of the training container's
life, so this is the one chance we have to get the training output off the
ephemeral Space filesystem and onto persistent HF storage.

What gets uploaded:

    Model repo  (HF model type)  — chane35/permanence-trained
        training/artifacts/grpo/adapter/       → root of model repo
                                                 (LoRA + tokenizer config)

    Artifact repo (HF dataset type) — chane35/permanence-artifacts
        training/artifacts/pipeline_summary.json
        training/artifacts/sft/status.json + metrics.json
        training/artifacts/gate/status.json + predictions.jsonl
        training/artifacts/grpo/status.json + metrics.json + training_log.json
        training/artifacts/eval/status.json + results.json + comparison.csv
        training/artifacts/grpo/adapter/       → grpo_adapter/   (duplicate,
                                                  so the dataset repo is
                                                  self-contained for forensics)
        results/training_curves.png  (+ any other PNG in results/)
        results/training_summary.txt

Token resolution order:
    1. ``HF_TOKEN`` env var
    2. ``HUGGINGFACE_TOKEN`` env var
    3. ``~/.cache/huggingface/token`` (written by ``huggingface-cli login``)
    4. Skip gracefully if none available — no hard fail
"""
from __future__ import annotations

import os
import sys
import traceback
from pathlib import Path
from typing import List, Optional


MODEL_REPO = os.environ.get("PERMANENCE_MODEL_REPO", "chane35/permanence-trained")
DATASET_REPO = os.environ.get("PERMANENCE_ARTIFACTS_REPO", "chane35/permanence-artifacts")

ARTIFACTS_DIR = Path("training/artifacts")
RESULTS_DIR = Path("results")


def _resolve_token() -> Optional[str]:
    """Find an HF token from env or the huggingface_hub cache."""
    for var in ("HF_TOKEN", "HUGGINGFACE_TOKEN", "HUGGING_FACE_HUB_TOKEN"):
        v = os.environ.get(var)
        if v:
            return v
    token = get_token()
    if token:
        return token
    return None


def _upload_file_if_exists(api, path: Path, repo_id: str, repo_type: str, path_in_repo: Optional[str] = None) -> bool:
    if not path.exists() or not path.is_file():
        return False
    try:
        api.upload_file(
            path_or_fileobj=str(path),
            path_in_repo=path_in_repo or path.name,
            repo_id=repo_id,
            repo_type=repo_type,
        )
        print(f"[auto_upload] ✓ {path} → {repo_id}:{path_in_repo or path.name}")
        return True
    except Exception as exc:
        print(f"[auto_upload] ✗ failed to upload {path}: {exc}")
        return False


def _upload_folder_if_exists(api, folder: Path, repo_id: str, repo_type: str, path_in_repo: str = "") -> bool:
    if not folder.exists() or not folder.is_dir():
        return False
    try:
        api.upload_folder(
            folder_path=str(folder),
            path_in_repo=path_in_repo,
            repo_id=repo_id,
            repo_type=repo_type,
            ignore_patterns=["*.tmp", "*.lock", "__pycache__/*"],
        )
        print(f"[auto_upload] ✓ {folder}/ → {repo_id}:{path_in_repo or '/'}")
        return True
    except Exception as exc:
        print(f"[auto_upload] ✗ failed to upload {folder}/: {exc}")
        return False


def upload() -> None:
    from huggingface_hub import HfApi

    token = _resolve_token()
    if not token:
        print("[auto_upload] No HF token available — skipping upload. Artifacts remain in the container only.")
        return

    api = HfApi(token=token)
    print(f"[auto_upload] Uploading artifacts")
    print(f"[auto_upload] Model repo:     {MODEL_REPO}")
    print(f"[auto_upload] Artifacts repo: {DATASET_REPO}")

    # ── Model repo — the trained GRPO adapter ──────────────────────────
    grpo_adapter = ARTIFACTS_DIR / "grpo" / "adapter"
    try:
        api.create_repo(MODEL_REPO, repo_type="model", exist_ok=True)
        ok = _upload_folder_if_exists(api, grpo_adapter, MODEL_REPO, "model")
        if not ok:
            print(f"[auto_upload] No GRPO adapter found at {grpo_adapter}. (Pipeline may have aborted before stage 3 finished.)")
            # Fall back to uploading the SFT adapter so at least SOMETHING
            # trained is preserved.
            sft_adapter = ARTIFACTS_DIR / "sft" / "adapter"
            if sft_adapter.exists():
                print(f"[auto_upload] Uploading SFT adapter as fallback")
                _upload_folder_if_exists(api, sft_adapter, MODEL_REPO, "model")
    except Exception as exc:
        print(f"[auto_upload] Model repo upload failed: {exc}")
        traceback.print_exc()

    # ── Artifacts repo — every structured output, for reproducibility ──
    try:
        api.create_repo(DATASET_REPO, repo_type="dataset", exist_ok=True)

        # Top-level pipeline summary (single file)
        _upload_file_if_exists(api, ARTIFACTS_DIR / "pipeline_summary.json", DATASET_REPO, "dataset")

        # Per-stage artifacts (JSON / JSONL / CSV)
        stage_files = [
            ("sft/status.json", "sft/status.json"),
            ("sft/metrics.json", "sft/metrics.json"),
            ("gate/status.json", "gate/status.json"),
            ("gate/predictions.jsonl", "gate/predictions.jsonl"),
            ("grpo/status.json", "grpo/status.json"),
            ("grpo/metrics.json", "grpo/metrics.json"),
            ("grpo/training_log.json", "grpo/training_log.json"),
            ("eval/status.json", "eval/status.json"),
            ("eval/results.json", "eval/results.json"),
            ("eval/comparison.csv", "eval/comparison.csv"),
        ]
        for rel_src, rel_dst in stage_files:
            _upload_file_if_exists(api, ARTIFACTS_DIR / rel_src, DATASET_REPO, "dataset", rel_dst)

        # Adapter weights (duplicated here so the dataset repo is self-contained)
        _upload_folder_if_exists(api, grpo_adapter, DATASET_REPO, "dataset", "grpo_adapter")
        _upload_folder_if_exists(api, ARTIFACTS_DIR / "sft" / "adapter", DATASET_REPO, "dataset", "sft_adapter")

        # Curves + human-readable summaries
        if RESULTS_DIR.exists():
            for png in RESULTS_DIR.glob("*.png"):
                _upload_file_if_exists(api, png, DATASET_REPO, "dataset", f"curves/{png.name}")
            for txt in RESULTS_DIR.glob("*.txt"):
                _upload_file_if_exists(api, txt, DATASET_REPO, "dataset", txt.name)
            for json_file in RESULTS_DIR.glob("*.json"):
                _upload_file_if_exists(api, json_file, DATASET_REPO, "dataset", json_file.name)

        # Legacy training_log.json at permanence_output root, in case anything
        # still writes there (backward compat).
        _upload_file_if_exists(api, Path("permanence_output") / "training_log.json", DATASET_REPO, "dataset", "legacy_training_log.json")

        print(f"[auto_upload] ✓ Artifacts pushed to {DATASET_REPO}")
    except Exception as exc:
        print(f"[auto_upload] Artifact repo upload failed: {exc}")
        traceback.print_exc()


if __name__ == "__main__":
    try:
        upload()
    except Exception as exc:
        # Never block the entrypoint on upload errors — we can still retrieve
        # manually via ``hf download`` if something survived.
        print(f"[auto_upload] FATAL: {exc}")
        traceback.print_exc()
        sys.exit(0)