deploy_space.py

#!/usr/bin/env python3
"""
创建/更新 Hugging Face Docker Space，并自动写入运行所需的 Secrets / Variables。

用法：
  export HF_TOKEN=hf_xxx
  export HF_SPACE_ID=username/sub2api
  export DATABASE_HOST=...
  export DATABASE_PORT=5432
  export DATABASE_USER=...
  export DATABASE_PASSWORD=...
  export DATABASE_DBNAME=postgres
  export DATABASE_SSLMODE=require

  uv run --with huggingface_hub==1.9.0 python deploy_space.py "$HF_SPACE_ID" .
"""

from __future__ import annotations

import json
import os
import secrets
import string
import sys
import time
from pathlib import Path
import shutil
import tempfile

from huggingface_hub import HfApi


def getenv_first(*names: str, default: str | None = None) -> str | None:
    for name in names:
        value = os.getenv(name)
        if value:
            return value
    return default


def bool_env(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "y", "on"}


def random_password(length: int = 24) -> str:
    alphabet = string.ascii_letters + string.digits + "-_"
    return "".join(secrets.choice(alphabet) for _ in range(length))


def collect_runtime_config() -> tuple[dict[str, str], dict[str, str], dict[str, str]]:
    secrets_map = {
        "DATABASE_HOST": getenv_first("DATABASE_HOST", "SUPABASE_DB_HOST"),
        "DATABASE_PORT": getenv_first("DATABASE_PORT", "SUPABASE_DB_PORT", default="5432"),
        "DATABASE_USER": getenv_first("DATABASE_USER", "SUPABASE_DB_USER"),
        "DATABASE_PASSWORD": getenv_first("DATABASE_PASSWORD", "SUPABASE_DB_PASSWORD"),
        "DATABASE_DBNAME": getenv_first("DATABASE_DBNAME", "SUPABASE_DBNAME", default="postgres"),
        "DATABASE_SSLMODE": getenv_first("DATABASE_SSLMODE", default="require"),
        "ADMIN_EMAIL": getenv_first("ADMIN_EMAIL", default="admin@sub2api.local"),
        "ADMIN_PASSWORD": getenv_first("ADMIN_PASSWORD"),
        "JWT_SECRET": getenv_first("JWT_SECRET"),
    }

    generated: dict[str, str] = {}

    if not secrets_map["ADMIN_PASSWORD"]:
        secrets_map["ADMIN_PASSWORD"] = random_password(20)
        generated["ADMIN_PASSWORD"] = secrets_map["ADMIN_PASSWORD"]

    if not secrets_map["JWT_SECRET"]:
        secrets_map["JWT_SECRET"] = secrets.token_hex(32)
        generated["JWT_SECRET"] = secrets_map["JWT_SECRET"]

    variables_map = {
        "TZ": getenv_first("TZ", default="Asia/Shanghai"),
        "RUN_MODE": getenv_first("RUN_MODE", default="standard"),
        "SERVER_MODE": getenv_first("SERVER_MODE", default="release"),
        "AUTO_SETUP": getenv_first("AUTO_SETUP", default="true"),
    }

    missing = {
        key: value
        for key, value in secrets_map.items()
        if key
        in {
            "DATABASE_HOST",
            "DATABASE_PORT",
            "DATABASE_USER",
            "DATABASE_PASSWORD",
            "DATABASE_DBNAME",
        }
        and not value
    }
    if missing:
        raise SystemExit(
            "Missing required environment values: "
            + ", ".join(sorted(missing))
        )

    return secrets_map, variables_map, generated


def upload_bundle_if_present(api: HfApi, repo_id: str, folder_path: Path) -> None:
    bundle_path = folder_path / "bundle.tar.gz"
    if not bundle_path.exists():
        return

    if not hasattr(api, "upload_large_folder"):
        api.upload_file(
            repo_id=repo_id,
            repo_type="space",
            path_in_repo="bundle.tar.gz",
            path_or_fileobj=str(bundle_path),
            commit_message="Upload runtime bundle",
        )
        return

    with tempfile.TemporaryDirectory(prefix="hf-space-bundle-upload-") as tmp_dir:
        upload_dir = Path(tmp_dir)
        shutil.copy2(bundle_path, upload_dir / "bundle.tar.gz")

        original_create_repo = api.create_repo

        def create_repo_with_space_sdk(*args, **kwargs):
            kwargs.setdefault("space_sdk", "docker")
            kwargs.setdefault("exist_ok", True)
            return original_create_repo(*args, **kwargs)

        api.create_repo = create_repo_with_space_sdk
        try:
            api.upload_large_folder(
                repo_id=repo_id,
                folder_path=str(upload_dir),
                repo_type="space",
                revision="main",
                num_workers=1,
                print_report=True,
                print_report_every=10,
            )
        finally:
            api.create_repo = original_create_repo


def upload_space(api: HfApi, repo_id: str, folder_path: Path, private: bool) -> None:
    api.create_repo(
        repo_id=repo_id,
        repo_type="space",
        space_sdk="docker",
        private=private,
        exist_ok=True,
    )

    api.upload_folder(
        repo_id=repo_id,
        repo_type="space",
        folder_path=str(folder_path),
        ignore_patterns=[
            ".git",
            ".DS_Store",
            "__pycache__",
            ".venv",
            "*.pyc",
            ".env",
            "bundle.tar.gz",
        ],
        commit_message="Deploy Hugging Face Docker Space wrapper",
    )
    upload_bundle_if_present(api, repo_id, folder_path)


def sync_runtime_config(
    api: HfApi,
    repo_id: str,
    secrets_map: dict[str, str],
    variables_map: dict[str, str],
) -> None:
    for key, value in secrets_map.items():
        if value:
            api.add_space_secret(repo_id=repo_id, key=key, value=value)

    for key, value in variables_map.items():
        if value:
            api.add_space_variable(repo_id=repo_id, key=key, value=value)


def wait_for_runtime(api: HfApi, repo_id: str, attempts: int = 30) -> None:
    for attempt in range(1, attempts + 1):
        runtime = api.get_space_runtime(repo_id=repo_id)
        stage = getattr(runtime, "stage", None)
        hardware = getattr(runtime, "hardware", None)
        requested_hardware = getattr(runtime, "requested_hardware", None)
        print(
            json.dumps(
                {
                    "attempt": attempt,
                    "stage": stage,
                    "hardware": hardware,
                    "requested_hardware": requested_hardware,
                },
                ensure_ascii=False,
            )
        )
        if stage in {"RUNNING", "RUNNING_BUILDING"}:
            return
        time.sleep(10)


def main() -> int:
    if len(sys.argv) < 2:
        print("Usage: deploy_space.py <username/space-name> [folder]", file=sys.stderr)
        return 2

    repo_id = sys.argv[1]
    folder = Path(sys.argv[2] if len(sys.argv) >= 3 else ".").resolve()
    token = getenv_first("HF_TOKEN", "HUGGINGFACE_TOKEN")
    if not token:
        raise SystemExit("HF_TOKEN is required")

    private = bool_env("HF_SPACE_PRIVATE", default=False)
    api = HfApi(token=token)

    secrets_map, variables_map, generated = collect_runtime_config()

    print(f"Uploading {folder} -> https://huggingface.co/spaces/{repo_id}")
    upload_space(api, repo_id, folder, private=private)
    sync_runtime_config(api, repo_id, secrets_map, variables_map)

    print("Generated secrets (please save them):")
    if generated:
        for key, value in generated.items():
            print(f"{key}={value}")
    else:
        print("(none)")

    print("Waiting for runtime status...")
    wait_for_runtime(api, repo_id)
    print(f"Space URL: https://huggingface.co/spaces/{repo_id}")
    print(f"Embed URL: https://{repo_id.replace('/', '-')}.hf.space")
    return 0


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