scripts/deploy_hf_space.py

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path


ROOT_DIR = Path(__file__).resolve().parent.parent
DEFAULT_BUNDLE_DIR = ROOT_DIR / "outputs" / "huggingface-space"


def load_env_file(path: Path) -> None:
    if not path.exists():
        return
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


def require_huggingface_hub():
    try:
        from huggingface_hub import HfApi, create_repo, upload_folder
    except ImportError as exc:
        raise SystemExit(
            "huggingface_hub is not installed. Run: .\\.venv\\Scripts\\python.exe -m pip install huggingface_hub"
        ) from exc
    return HfApi, create_repo, upload_folder


def run_export(bundle_dir: Path, force_export: bool) -> None:
    if bundle_dir.exists() and not force_export:
        return
    command = [
        sys.executable,
        str(ROOT_DIR / "scripts" / "export_hf_space.py"),
        "--out",
        str(bundle_dir),
    ]
    if force_export:
        command.append("--force")
    subprocess.run(command, cwd=ROOT_DIR, check=True)


def get_token(explicit_token: str | None) -> str:
    load_env_file(ROOT_DIR / ".env")
    token = explicit_token or os.getenv("HF_TOKEN") or os.getenv("HF_API_TOKEN") or os.getenv("HUGGINGFACE_API_TOKEN")
    if not token:
        raise SystemExit(
            "Missing Hugging Face token. Set HF_TOKEN, HF_API_TOKEN, or HUGGINGFACE_API_TOKEN, "
            "or pass --token. Create a write token at https://huggingface.co/settings/tokens."
        )
    return token


def token_available(explicit_token: str | None = None) -> bool:
    load_env_file(ROOT_DIR / ".env")
    return bool(explicit_token or os.getenv("HF_TOKEN") or os.getenv("HF_API_TOKEN") or os.getenv("HUGGINGFACE_API_TOKEN"))


def split_repo_id(repo_id: str) -> tuple[str, str]:
    if "/" not in repo_id:
        raise ValueError("Hugging Face Space repo id must look like username/space-name")
    owner, space_name = repo_id.split("/", 1)
    if not owner or not space_name:
        raise ValueError("Hugging Face Space repo id must include both owner and space name")
    return owner, space_name


def worker_url_for_repo(repo_id: str) -> str:
    owner, space_name = split_repo_id(repo_id)
    return f"https://{owner}-{space_name}.hf.space".lower()


def deploy_hf_space(
    repo_id: str,
    bundle_dir: Path = DEFAULT_BUNDLE_DIR,
    token: str | None = None,
    private: bool = False,
    force_export: bool = False,
    commit_message: str = "Deploy Arabic Audio Reader worker",
) -> dict[str, object]:
    bundle_dir = bundle_dir.resolve()
    run_export(bundle_dir, force_export)
    hf_token = get_token(token)
    _hf_api, create_repo, upload_folder = require_huggingface_hub()

    create_repo(
        repo_id=repo_id,
        repo_type="space",
        space_sdk="docker",
        private=private,
        exist_ok=True,
        token=hf_token,
    )
    commit_info = upload_folder(
        repo_id=repo_id,
        repo_type="space",
        folder_path=str(bundle_dir),
        path_in_repo=".",
        commit_message=commit_message,
        token=hf_token,
    )
    worker_url = worker_url_for_repo(repo_id)
    return {
        "repoId": repo_id,
        "workerUrl": worker_url,
        "bundleDir": str(bundle_dir),
        "commitUrl": getattr(commit_info, "commit_url", None),
        "nextCommand": (
            f"python scripts\\deployment_handoff.py {worker_url} "
            "--origin https://arabic-translator-mu.vercel.app --code 1234"
        ),
    }


def main_with_args(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Create or update the Hugging Face Docker Space worker.")
    parser.add_argument(
        "repo_id",
        help="Hugging Face Space repo id, for example username/arabic-audio-reader-worker.",
    )
    parser.add_argument("--bundle-dir", type=Path, default=DEFAULT_BUNDLE_DIR)
    parser.add_argument("--token", help="Hugging Face write token. Prefer setting HF_TOKEN instead.")
    parser.add_argument("--private", action="store_true", help="Create the Space as private.")
    parser.add_argument("--force-export", action="store_true", help="Rebuild the worker bundle before upload.")
    parser.add_argument("--commit-message", default="Deploy Arabic Audio Reader worker")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args(argv)

    result = deploy_hf_space(
        repo_id=args.repo_id,
        bundle_dir=args.bundle_dir,
        token=args.token,
        private=args.private,
        force_export=args.force_export,
        commit_message=args.commit_message,
    )
    if args.json:
        import json

        print(json.dumps(result, indent=2))
    else:
        print(f"Uploaded worker bundle to https://huggingface.co/spaces/{args.repo_id}")
        print(f"Worker URL: {result['workerUrl']}")
        print("Next:")
        print(result["nextCommand"])
    return 0


def main() -> int:
    return main_with_args()


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