hf_data.py

#!/usr/bin/env python
"""Hugging Face data sync helpers for BrainRL.

The environment is intentionally lightweight: it needs configs, the frozen
parcel manifest, participant metadata, and optional word annotations. This
module keeps those artifacts versioned in a HF Dataset repo so Colab, HF Jobs,
and HF Spaces can all run from the same data revision.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_EXPORT_DIR = PROJECT_ROOT / "hf_data_bundle"
DEFAULT_CACHE_DIR = Path(os.getenv("BRAINRL_DATA_DIR", "/tmp/brainrl-data")).expanduser()
REQUIRED_CONFIG_FILES = (
    "subset_config.yaml",
    "region_priors.json",
    "participant_run_info.json",
)
OPTIONAL_CONFIG_FILES = (
    "parcel_candidates.json",
)


def _resolve_token(raw_token: str | None = None) -> str | None:
    return raw_token or os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")


def _copy_file(src: Path, dst: Path, *, required: bool) -> bool:
    if not src.exists():
        if required:
            raise FileNotFoundError(f"Required BrainRL data file is missing: {src}")
        return False
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)
    return True


def export_data_bundle(
    output_dir: Path,
    *,
    config_dir: Path | None = None,
    annotation_dir: Path | None = None,
    include_annotations: bool = True,
) -> dict[str, Any]:
    """Copy the portable BrainRL data subset into ``output_dir``."""

    config_root = (config_dir or PROJECT_ROOT / "configs").expanduser()
    annotation_root = (
        annotation_dir
        or Path(os.getenv("BRAINRL_STIMULUS_DIR", PROJECT_ROOT.parent / "data" / "annotation"))
    ).expanduser()
    output_dir = output_dir.expanduser()
    if output_dir.exists():
        shutil.rmtree(output_dir)
    (output_dir / "configs").mkdir(parents=True, exist_ok=True)

    copied: list[str] = []
    for name in REQUIRED_CONFIG_FILES:
        if _copy_file(config_root / name, output_dir / "configs" / name, required=True):
            copied.append(f"configs/{name}")
    for name in OPTIONAL_CONFIG_FILES:
        if _copy_file(config_root / name, output_dir / "configs" / name, required=False):
            copied.append(f"configs/{name}")

    annotation_count = 0
    if include_annotations and annotation_root.exists():
        out_annotation = output_dir / "annotation"
        out_annotation.mkdir(parents=True, exist_ok=True)
        for csv_path in sorted(annotation_root.glob("*.csv")):
            shutil.copy2(csv_path, out_annotation / csv_path.name)
            annotation_count += 1
            copied.append(f"annotation/{csv_path.name}")

    metadata = {
        "format": "brainrl-hf-data-v1",
        "config_files": copied,
        "annotation_csv_count": annotation_count,
        "source_config_dir": str(config_root),
        "source_annotation_dir": str(annotation_root) if annotation_root.exists() else None,
    }
    with (output_dir / "metadata.json").open("w", encoding="utf-8") as handle:
        json.dump(metadata, handle, indent=2)
    return metadata


def _file_sha256(path: Path) -> str:
    """Stream-hash a file with sha256 so we can pin a parcel manifest revision."""

    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def parcel_manifest_summary(parcel_manifest: Path) -> dict[str, Any]:
    """Extract checksum + budget metadata from the parcel manifest, if present.

    Used by the Space ``/health`` endpoint so callers can confirm the running
    Space is using the exact same manifest revision as the trainer.
    """

    if not parcel_manifest.exists():
        return {"present": False}

    payload = json.loads(parcel_manifest.read_text(encoding="utf-8"))
    candidates = payload.get("candidates")
    return {
        "present": True,
        "path": str(parcel_manifest),
        "sha256": _file_sha256(parcel_manifest),
        "size_bytes": parcel_manifest.stat().st_size,
        "selection_budget": payload.get("selection_budget"),
        "max_candidates": payload.get("max_candidates"),
        "candidate_count": len(candidates) if isinstance(candidates, list) else None,
    }


def validate_data_root(root: Path) -> dict[str, Any]:
    """Validate a downloaded/exported HF data root and return a summary.

    The summary is consumed by the Space's ``/brainrl/data_status`` and
    ``/health`` endpoints so a notebook caller can confirm the running Space
    matches the data revision the trainer just pushed.
    """

    root = root.expanduser()
    config_dir = root / "configs"
    missing = [name for name in REQUIRED_CONFIG_FILES if not (config_dir / name).exists()]
    if missing:
        raise FileNotFoundError(
            f"BrainRL data root {root} is missing required config files: {missing}"
        )

    json_files = [config_dir / "region_priors.json", config_dir / "participant_run_info.json"]
    parcel_manifest = config_dir / "parcel_candidates.json"
    if parcel_manifest.exists():
        json_files.append(parcel_manifest)
    for json_path in json_files:
        with json_path.open("r", encoding="utf-8") as handle:
            json.load(handle)

    annotation_dir = root / "annotation"
    return {
        "root": str(root),
        "config_dir": str(config_dir),
        "has_parcel_manifest": parcel_manifest.exists(),
        "parcel_manifest": parcel_manifest_summary(parcel_manifest),
        "annotation_csv_count": len(list(annotation_dir.glob("*.csv")))
        if annotation_dir.exists()
        else 0,
        "data_repo": os.getenv("BRAINRL_DATA_REPO"),
        "data_revision": os.getenv("BRAINRL_DATA_REVISION") or None,
    }


def download_dataset_repo(
    repo_id: str,
    *,
    output_dir: Path = DEFAULT_CACHE_DIR,
    revision: str | None = None,
    token: str | None = None,
) -> dict[str, Any]:
    """Download a HF Dataset repo into ``output_dir`` and validate it."""

    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:  # pragma: no cover
        raise SystemExit(
            "huggingface_hub is required for HF data sync. Install with "
            "`pip install -e .[deploy]` or `pip install huggingface_hub`."
        ) from exc

    output_dir = output_dir.expanduser()
    output_dir.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        revision=revision,
        local_dir=str(output_dir),
        token=_resolve_token(token),
        allow_patterns=["configs/**", "annotation/**", "metadata.json", "README.md"],
    )
    return validate_data_root(output_dir)


def upload_dataset_repo(
    repo_id: str,
    *,
    bundle_dir: Path,
    private: bool,
    token: str | None = None,
    commit_message: str = "Upload BrainRL config data",
) -> None:
    """Create/update the HF Dataset repo with a prepared data bundle."""

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

    token = _resolve_token(token)
    create_repo(
        repo_id=repo_id,
        repo_type="dataset",
        private=private,
        exist_ok=True,
        token=token,
    )
    upload_folder(
        folder_path=str(bundle_dir.expanduser()),
        repo_id=repo_id,
        repo_type="dataset",
        token=token,
        commit_message=commit_message,
    )


def sync_data_from_env() -> dict[str, Any] | None:
    """Download HF data when ``BRAINRL_DATA_REPO`` is set.

    The function also sets ``BRAINRL_CONFIG_DIR`` and ``BRAINRL_STIMULUS_DIR``
    for downstream loaders if the downloaded files are present.
    """

    repo_id = os.getenv("BRAINRL_DATA_REPO")
    if not repo_id:
        return None

    revision = os.getenv("BRAINRL_DATA_REVISION") or None
    output_dir = Path(os.getenv("BRAINRL_DATA_DIR", str(DEFAULT_CACHE_DIR))).expanduser()
    summary = download_dataset_repo(repo_id, output_dir=output_dir, revision=revision)
    os.environ.setdefault("BRAINRL_CONFIG_DIR", summary["config_dir"])
    annotation_dir = output_dir / "annotation"
    if annotation_dir.exists():
        os.environ.setdefault("BRAINRL_STIMULUS_DIR", str(annotation_dir))
    return summary


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Manage BrainRL HF Dataset artifacts")
    sub = parser.add_subparsers(dest="command", required=True)

    export = sub.add_parser("export", help="Export local configs/annotations to a bundle")
    export.add_argument("--output-dir", type=str, default=str(DEFAULT_EXPORT_DIR))
    export.add_argument("--config-dir", type=str, default=None)
    export.add_argument("--annotation-dir", type=str, default=None)
    export.add_argument("--no-annotations", action="store_true")

    upload = sub.add_parser("upload", help="Export then upload to a HF Dataset repo")
    upload.add_argument("--repo-id", type=str, required=True)
    upload.add_argument("--bundle-dir", type=str, default=str(DEFAULT_EXPORT_DIR))
    upload.add_argument("--config-dir", type=str, default=None)
    upload.add_argument("--annotation-dir", type=str, default=None)
    upload.add_argument("--no-annotations", action="store_true")
    upload.add_argument("--public", action="store_true")
    upload.add_argument("--token", type=str, default=None)
    upload.add_argument("--commit-message", type=str, default="Upload BrainRL config data")

    download = sub.add_parser("download", help="Download/validate a HF Dataset repo")
    download.add_argument("--repo-id", type=str, required=True)
    download.add_argument("--output-dir", type=str, default=str(DEFAULT_CACHE_DIR))
    download.add_argument("--revision", type=str, default=None)
    download.add_argument("--token", type=str, default=None)

    validate = sub.add_parser("validate", help="Validate a local data root")
    validate.add_argument("--data-root", type=str, required=True)

    return parser


def main() -> None:
    args = build_parser().parse_args()
    if args.command == "export":
        summary = export_data_bundle(
            Path(args.output_dir),
            config_dir=Path(args.config_dir) if args.config_dir else None,
            annotation_dir=Path(args.annotation_dir) if args.annotation_dir else None,
            include_annotations=not bool(args.no_annotations),
        )
        print(json.dumps(summary, indent=2))
        return

    if args.command == "upload":
        bundle_dir = Path(args.bundle_dir)
        export_data_bundle(
            bundle_dir,
            config_dir=Path(args.config_dir) if args.config_dir else None,
            annotation_dir=Path(args.annotation_dir) if args.annotation_dir else None,
            include_annotations=not bool(args.no_annotations),
        )
        upload_dataset_repo(
            args.repo_id,
            bundle_dir=bundle_dir,
            private=not bool(args.public),
            token=args.token,
            commit_message=args.commit_message,
        )
        print(f"Uploaded BrainRL data bundle to dataset repo {args.repo_id}")
        return

    if args.command == "download":
        summary = download_dataset_repo(
            args.repo_id,
            output_dir=Path(args.output_dir),
            revision=args.revision,
            token=args.token,
        )
        print(json.dumps(summary, indent=2))
        return

    if args.command == "validate":
        print(json.dumps(validate_data_root(Path(args.data_root)), indent=2))
        return


if __name__ == "__main__":
    main()