scripts/deploy.py

"""Create (or update) the Space and upload this project.

Usage:
  python scripts/deploy.py [--repo wang2226/steering-showcase] [--message "..."] [--hardware zero-a10g]

Only files tracked by git are uploaded, so untracked material and local secrets (.env.local, the
virtualenv, local/) can never reach the Space, and the working tree must be clean, so what the Space
runs is always a commit; its hash goes into the Space's commit message. Files the Space has that this
project no longer has are deleted in the same commit. ZeroGPU is requested when the Space is created;
pass --hardware to change it later, since a request restarts the Space. Needs a write token: HF_TOKEN in
the environment, `hf auth login`, or HF_TOKEN=... in a git-ignored .env.local.
"""

from __future__ import annotations

import argparse
import os
import subprocess
from pathlib import Path

from huggingface_hub import HfApi

ROOT = Path(__file__).resolve().parents[1]
#: Tracked, but meaningful only on GitHub.
SKIP_PREFIXES = (".github/",)
#: The Space's own files, which this project does not track.
KEEP_REMOTE = {".gitattributes"}


def token_from_env_file() -> str | None:
    path = ROOT / ".env.local"
    if not path.exists():
        return None
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip().startswith("HF_TOKEN="):
            return line.split("=", 1)[1].strip() or None
    return None


def git(*args: str) -> str:
    return subprocess.run(["git", *args], cwd=ROOT, check=True, capture_output=True).stdout.decode()


def tracked_files() -> list[str]:
    out = git("ls-files", "-z")
    return sorted(f for f in out.split("\0") if f and not f.startswith(SKIP_PREFIXES) and (ROOT / f).is_file())


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="wang2226/steering-showcase")
    ap.add_argument("--private", action="store_true")
    ap.add_argument("--hardware", default="", help="request this hardware, e.g. zero-a10g (restarts the Space)")
    ap.add_argument("--message", default="Deploy the representation steering demo")
    ap.add_argument("--dry-run", action="store_true", help="list what would be uploaded and deleted")
    args = ap.parse_args()

    dirty = [line for line in git("status", "--porcelain", "--untracked-files=no").splitlines() if line.strip()]
    if dirty and not args.dry_run:
        raise SystemExit("commit or stash these first, so the Space runs a commit:\n" + "\n".join(dirty))
    head = git("rev-parse", "--short", "HEAD").strip()

    files = tracked_files()
    api = HfApi(token=os.environ.get("HF_TOKEN") or token_from_env_file())
    exists = api.repo_exists(args.repo, repo_type="space")
    remote = set(api.list_repo_files(args.repo, repo_type="space")) if exists else set()
    stale = sorted(remote - set(files) - KEEP_REMOTE)
    print(f"upload {len(files)} files at {head}; delete {len(stale)} the Space no longer needs: {stale}")
    if args.dry_run:
        print("\n".join(files))
        return

    if not exists:
        # Request the hardware at creation: free accounts may not create Gradio Spaces on cpu-basic, but may
        # host ZeroGPU Spaces.
        url = api.create_repo(args.repo, repo_type="space", space_sdk="gradio", private=args.private,
                              space_hardware=args.hardware or "zero-a10g")
        print("created:", url)
    elif args.hardware:
        try:
            api.request_space_hardware(args.repo, args.hardware)
        except Exception as exc:  # e.g. the account is not eligible for ZeroGPU
            print(f"could not set hardware {args.hardware!r}: {exc}\n  -> choose it in the Space settings")
    info = api.upload_folder(repo_id=args.repo, repo_type="space", folder_path=str(ROOT), allow_patterns=files,
                             delete_patterns=stale or None, commit_message=f"{args.message} ({head})")
    print("uploaded:", getattr(info, "commit_url", info))
    runtime = api.get_space_runtime(args.repo)
    print("runtime:", runtime.stage, "hardware:", runtime.hardware)


if __name__ == "__main__":
    main()