deploy_to_hf.py

#!/usr/bin/env python
"""Deploy the BrainRL OpenEnv server to a Hugging Face Space.

Pushes the project as a Docker-SDK Space (the README.md frontmatter already
declares ``sdk: docker`` and ``app_port: 8000``). The local trainer/client can
then point at the Space URL to interact with the deployed environment.

Examples
--------

    huggingface-cli login   # or set HF_TOKEN

    python deploy_to_hf.py \\
      --repo-id mohith202/transformer \\
      --include-manifest

    # Public Space:
    python deploy_to_hf.py --repo-id you/brainrl-region-selection --public

After deploy, the Space URL is printed. Wire your client to it via:

    import os
    os.environ["BRAINRL_API_URL"] = "https://<user>-<space-name>.hf.space"
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parent

DEFAULT_INCLUDE = [
    "README.md",
    "blog.md",
    "Dockerfile",
    "openenv.yaml",
    "pyproject.toml",
    "requirements.txt",
    "Makefile",
    "__init__.py",
    "client.py",
    "models.py",
    "data_loader.py",
    "data_split.py",
    "stimulus_loader.py",
    "prompts.py",
    "rewards.py",
    "metrics.py",
    "baselines.py",
    "evaluate.py",
    "inference.py",
    "train_grpo.py",
    "prepare_parcels.py",
    "plotting.py",
    "hf_data.py",
    "hf_jobs.py",
    "deploy_to_hf.py",
    "configs/**",
    "notebooks/**",
    "server/**",
    # server/** already covers prediction_ui.py, but list it explicitly so a
    # selective deploy that overrides DEFAULT_INCLUDE still picks it up.
    "server/prediction_ui.py",
]

DEFAULT_IGNORE = [
    "**/__pycache__/**",
    "**/*.pyc",
    "*.egg-info",
    "*.egg-info/**",
    "outputs/**",
    "build/**",
    "dist/**",
    ".git/**",
    ".venv/**",
    "venv/**",
    "*.nii",
    "*.nii.gz",
    "*.npy",
    "*.npz",
]


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Deploy BrainRL to a HF Space")
    parser.add_argument(
        "--repo-id",
        type=str,
        required=True,
        help="<user_or_org>/<space_name> (e.g. mohith/brainrl-region-selection)",
    )
    parser.add_argument(
        "--token",
        type=str,
        default=None,
        help="HF token (defaults to $HF_TOKEN / huggingface-cli login).",
    )
    parser.add_argument(
        "--public",
        action="store_true",
        help="Create the Space as public (default: private).",
    )
    parser.add_argument(
        "--commit-message",
        type=str,
        default="Deploy BrainRL OpenEnv server",
    )
    parser.add_argument(
        "--include-manifest",
        action="store_true",
        help="Include configs/parcel_candidates.json (run prepare_parcels.py first).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would be uploaded without actually pushing.",
    )
    return parser


def _resolve_token(arg_token: str | None) -> str | None:
    return arg_token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def _check_manifest(include_manifest: bool) -> None:
    manifest = PROJECT_ROOT / "configs" / "parcel_candidates.json"
    if include_manifest and not manifest.exists():
        raise SystemExit(
            "configs/parcel_candidates.json is missing. Run `make prepare` first or "
            "drop --include-manifest."
        )
    if not include_manifest and manifest.exists():
        print(
            f"[note] {manifest} exists locally but will NOT be uploaded "
            "(pass --include-manifest to ship it)."
        )


def _build_ignore(include_manifest: bool) -> list[str]:
    ignore = list(DEFAULT_IGNORE)
    if not include_manifest:
        ignore.append("configs/parcel_candidates.json")
    return ignore


def _build_allow(include_manifest: bool) -> list[str]:
    allow = list(DEFAULT_INCLUDE)
    if include_manifest:
        allow.append("configs/parcel_candidates.json")
    return allow


def main() -> None:
    args = build_parser().parse_args()

    try:
        from huggingface_hub import HfApi, create_repo, upload_folder
    except ImportError as exc:  # pragma: no cover
        raise SystemExit(
            "huggingface_hub is required. Install with `pip install -e .[deploy]` "
            "or `pip install huggingface_hub`."
        ) from exc

    token = _resolve_token(args.token)
    _check_manifest(args.include_manifest)

    allow = _build_allow(args.include_manifest)
    ignore = _build_ignore(args.include_manifest)

    if args.dry_run:
        print(f"[dry-run] repo_id={args.repo_id} private={not args.public}")
        print("[dry-run] would upload (allow patterns):")
        for entry in allow:
            print(f"   + {entry}")
        print("[dry-run] would ignore patterns:")
        for entry in ignore:
            print(f"   - {entry}")
        return

    print(f"[1/3] Ensuring HF Space {args.repo_id} exists (sdk=docker)...")
    create_repo(
        repo_id=args.repo_id,
        repo_type="space",
        space_sdk="docker",
        private=not args.public,
        exist_ok=True,
        token=token,
    )

    print(f"[2/3] Uploading project from {PROJECT_ROOT}...")
    upload_folder(
        folder_path=str(PROJECT_ROOT),
        repo_id=args.repo_id,
        repo_type="space",
        token=token,
        commit_message=args.commit_message,
        allow_patterns=allow,
        ignore_patterns=ignore,
    )

    space_url = f"https://huggingface.co/spaces/{args.repo_id}"
    api_url = f"https://{args.repo_id.replace('/', '-')}.hf.space"
    print(f"[3/3] Done. Space:        {space_url}")
    print(f"       Live API URL:     {api_url}")
    print(f"       Set BRAINRL_API_URL={api_url} for clients/inference.")
    print("       First boot may take a few minutes while HF builds the Docker image.")


if __name__ == "__main__":
    sys.exit(main())