scripts/hf_artifacts.py

"""HF Hub artifact sync helpers for run directories and ETL caches."""

from __future__ import annotations

import argparse
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download


# ---------------------------------------------------------------------------
# Model-repo helpers — run outputs (checkpoints, metrics, embeddings)
# ---------------------------------------------------------------------------


def push(repo: str, preset: str, source: Path, message: str = "snapshot") -> None:
    """Push one preset/run directory into the artifact repo namespace."""

    api = HfApi()
    api.create_repo(repo, repo_type="model", exist_ok=True, private=True)
    folder_path = source / preset
    if not folder_path.exists():
        raise FileNotFoundError(f"Missing preset artifact directory: {folder_path}")
    api.upload_folder(
        repo_id=repo,
        repo_type="model",
        folder_path=str(folder_path),
        path_in_repo=f"runs/{preset}",
        commit_message=message,
        ignore_patterns=["*.lmdb", "*.lmdb-lock"],
    )


def pull(repo: str, preset: str, dest: Path) -> None:
    """Pull one preset/run directory from the artifact repo."""

    dest.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=repo,
        repo_type="model",
        local_dir=str(dest),
        allow_patterns=[f"runs/{preset}/**"],
    )
    source = dest / "runs" / preset
    if not source.exists():
        return
    final_dir = dest / preset
    final_dir.mkdir(parents=True, exist_ok=True)
    for child in source.iterdir():
        target = final_dir / child.name
        if target.exists():
            if target.is_dir():
                continue
            target.unlink()
        child.rename(target)


# ---------------------------------------------------------------------------
# Dataset-repo helpers — ETL caches (LMDB, registry, classes) keyed by
# processed_signature.  Closes the gap documented in CLAUDE.md §6
# "Active gaps in HF coverage" by promoting the LMDB out of the model repo.
# ---------------------------------------------------------------------------


_ETL_ARTIFACTS = ("tensors.lmdb", "registry.csv", "classes.json", "etl_cache_manifest.json")
_QUANT_ARTIFACTS = (
    "transport_plan.pt",
    "alignment_offsets.pt",
    "tile_repo.pt",
    "read_data.pt",
    "counts.csv",
    "state.json",
    "summary.json",
)


def _etl_path_in_repo(processed_signature: str) -> str:
    return f"lmdb/{processed_signature}"


def sync_etl_to_dataset(
    repo: str,
    processed_signature: str,
    processed_root: Path,
    *,
    message: str = "etl-cache",
) -> dict[str, str]:
    """Upload an ETL output (`tensors.lmdb`, `registry.csv`, `classes.json`) to a dataset repo.

    `processed_signature` is the canonical cache key emitted by
    `deepgenopix matrix --json` (e.g. `px12_ps12_fl500-500_vf0p1_tf0p1_seed42_mf1`).
    The function uploads each present file individually so partial caches still
    push and tail files do not block earlier wins.
    """

    api = HfApi()
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    processed_root = Path(processed_root)
    uploaded: dict[str, str] = {}
    path_prefix = _etl_path_in_repo(processed_signature)
    for artifact in _ETL_ARTIFACTS:
        candidate = processed_root / artifact
        if not candidate.exists():
            continue
        api.upload_file(
            repo_id=repo,
            repo_type="dataset",
            path_or_fileobj=str(candidate),
            path_in_repo=f"{path_prefix}/{artifact}",
            commit_message=f"{message}: {artifact}",
        )
        uploaded[artifact] = f"{path_prefix}/{artifact}"
    return uploaded


def pull_etl_from_dataset(
    repo: str,
    processed_signature: str,
    processed_root: Path,
) -> dict[str, str]:
    """Pull the ETL cache for `processed_signature` from the dataset repo, if present.

    Returns a dict of `{artifact_name: local_path}` for each artifact that was
    materialized locally. An empty dict signals a cache miss — the caller is
    expected to run `deepgenopix etl` in that case.
    """

    processed_root = Path(processed_root)
    processed_root.mkdir(parents=True, exist_ok=True)
    path_prefix = _etl_path_in_repo(processed_signature)
    snapshot_dir = processed_root / "_hf_pull_cache"
    snapshot_dir.mkdir(parents=True, exist_ok=True)
    try:
        snapshot_download(
            repo_id=repo,
            repo_type="dataset",
            local_dir=str(snapshot_dir),
            allow_patterns=[f"{path_prefix}/**"],
        )
    except Exception:
        return {}

    pulled: dict[str, str] = {}
    source = snapshot_dir / "lmdb" / processed_signature
    if not source.exists():
        return pulled
    for artifact in _ETL_ARTIFACTS:
        candidate = source / artifact
        if not candidate.exists():
            continue
        destination = processed_root / artifact
        if destination.exists():
            destination.unlink()
        candidate.rename(destination)
        pulled[artifact] = str(destination)
    return pulled


# ---------------------------------------------------------------------------
# Dataset-repo helpers — quantification outputs consumed by variant jobs.
# These artifacts are large run inputs, not model checkpoints, so the canonical
# HF Jobs path keeps them in DATA_DATASET/quant_runs/<output_name>/.
# ---------------------------------------------------------------------------


def _quant_path_in_repo(output_name: str) -> str:
    return f"quant_runs/{output_name}"


def sync_quant_to_dataset(
    repo: str,
    output_name: str,
    quant_run_dir: Path,
    *,
    message: str = "quant-run",
) -> dict[str, str]:
    """Upload quantification artifacts into the dataset repo namespace."""

    api = HfApi()
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    quant_run_dir = Path(quant_run_dir)
    uploaded: dict[str, str] = {}
    path_prefix = _quant_path_in_repo(output_name)
    for artifact in _QUANT_ARTIFACTS:
        candidate = quant_run_dir / artifact
        if not candidate.exists():
            continue
        api.upload_file(
            repo_id=repo,
            repo_type="dataset",
            path_or_fileobj=str(candidate),
            path_in_repo=f"{path_prefix}/{artifact}",
            commit_message=f"{message}: {artifact}",
        )
        uploaded[artifact] = f"{path_prefix}/{artifact}"
    return uploaded


def pull_quant_from_dataset(
    repo: str,
    output_name: str,
    quant_run_dir: Path,
) -> dict[str, str]:
    """Pull quantification artifacts for `output_name` from the dataset repo."""

    quant_run_dir = Path(quant_run_dir)
    quant_run_dir.mkdir(parents=True, exist_ok=True)
    path_prefix = _quant_path_in_repo(output_name)
    snapshot_dir = quant_run_dir / "_hf_pull_cache"
    snapshot_dir.mkdir(parents=True, exist_ok=True)
    try:
        snapshot_download(
            repo_id=repo,
            repo_type="dataset",
            local_dir=str(snapshot_dir),
            allow_patterns=[f"{path_prefix}/**"],
        )
    except Exception:
        return {}

    pulled: dict[str, str] = {}
    source = snapshot_dir / "quant_runs" / output_name
    if not source.exists():
        return pulled
    for artifact in _QUANT_ARTIFACTS:
        candidate = source / artifact
        if not candidate.exists():
            continue
        destination = quant_run_dir / artifact
        if destination.exists():
            destination.unlink()
        candidate.rename(destination)
        pulled[artifact] = str(destination)
    return pulled


def push_dataset_folder(
    repo: str,
    source: Path,
    path_in_repo: str,
    *,
    message: str = "dataset-artifacts",
) -> None:
    """Upload a local folder into a dataset repo.

    This is used by Lane B jobs whose artifact shapes differ by corpus or sweep
    but all belong in dataset namespaces rather than the model artifact repo.
    """

    source = Path(source)
    if not source.exists():
        raise FileNotFoundError(f"Missing dataset artifact folder: {source}")
    api = HfApi()
    api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
    api.upload_folder(
        repo_id=repo,
        repo_type="dataset",
        folder_path=str(source),
        path_in_repo=path_in_repo.strip("/"),
        commit_message=message,
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="Push/pull DeepGenopix artifacts via the HF Hub")
    subparsers = parser.add_subparsers(dest="action", required=True)

    run_push = subparsers.add_parser("push", help="Push a run folder into the model repo")
    run_push.add_argument("--repo", required=True)
    run_push.add_argument("--preset", required=True)
    run_push.add_argument("--source", type=Path, default=Path("data/output"))
    run_push.add_argument("--message", default="snapshot")

    run_pull = subparsers.add_parser("pull", help="Pull a run folder from the model repo")
    run_pull.add_argument("--repo", required=True)
    run_pull.add_argument("--preset", required=True)
    run_pull.add_argument("--dest", type=Path, default=Path("data/output"))

    etl_push = subparsers.add_parser("push-etl", help="Push an ETL cache into a dataset repo")
    etl_push.add_argument("--repo", required=True)
    etl_push.add_argument("--signature", required=True)
    etl_push.add_argument("--processed-root", type=Path, required=True)
    etl_push.add_argument("--message", default="etl-cache")

    etl_pull = subparsers.add_parser("pull-etl", help="Pull an ETL cache from a dataset repo")
    etl_pull.add_argument("--repo", required=True)
    etl_pull.add_argument("--signature", required=True)
    etl_pull.add_argument("--processed-root", type=Path, required=True)

    quant_push = subparsers.add_parser("push-quant", help="Push quant artifacts into a dataset repo")
    quant_push.add_argument("--repo", required=True)
    quant_push.add_argument("--output-name", required=True)
    quant_push.add_argument("--quant-run-dir", type=Path, required=True)
    quant_push.add_argument("--message", default="quant-run")

    quant_pull = subparsers.add_parser("pull-quant", help="Pull quant artifacts from a dataset repo")
    quant_pull.add_argument("--repo", required=True)
    quant_pull.add_argument("--output-name", required=True)
    quant_pull.add_argument("--quant-run-dir", type=Path, required=True)

    dataset_folder_push = subparsers.add_parser(
        "push-dataset-folder", help="Push an arbitrary folder into a dataset repo"
    )
    dataset_folder_push.add_argument("--repo", required=True)
    dataset_folder_push.add_argument("--source", type=Path, required=True)
    dataset_folder_push.add_argument("--path-in-repo", required=True)
    dataset_folder_push.add_argument("--message", default="dataset-artifacts")

    args = parser.parse_args()
    if args.action == "push":
        push(args.repo, args.preset, args.source, args.message)
    elif args.action == "pull":
        pull(args.repo, args.preset, args.dest)
    elif args.action == "push-etl":
        result = sync_etl_to_dataset(
            args.repo, args.signature, args.processed_root, message=args.message
        )
        for artifact, remote_path in result.items():
            print(f"uploaded {artifact} → {remote_path}")
    elif args.action == "pull-etl":
        result = pull_etl_from_dataset(args.repo, args.signature, args.processed_root)
        if not result:
            print(f"cache miss for signature {args.signature}")
        else:
            for artifact, local_path in result.items():
                print(f"pulled {artifact} → {local_path}")
    elif args.action == "push-quant":
        result = sync_quant_to_dataset(
            args.repo, args.output_name, args.quant_run_dir, message=args.message
        )
        for artifact, remote_path in result.items():
            print(f"uploaded {artifact} → {remote_path}")
    elif args.action == "pull-quant":
        result = pull_quant_from_dataset(args.repo, args.output_name, args.quant_run_dir)
        if not result:
            print(f"cache miss for quant run {args.output_name}")
        else:
            for artifact, local_path in result.items():
                print(f"pulled {artifact} → {local_path}")
    elif args.action == "push-dataset-folder":
        push_dataset_folder(
            args.repo,
            args.source,
            args.path_in_repo,
            message=args.message,
        )
        print(f"uploaded {args.source} → {args.repo}/{args.path_in_repo.strip('/')}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())