scripts/release_challenge_data.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import re
from dataclasses import dataclass
from pathlib import Path

import h5py
from huggingface_hub import CommitOperationAdd, HfApi


DEFAULT_DEV_REPO_ID = "galaxy-sbi-challenge/sbi-galaxy-dev"
DEFAULT_BLINDED_FEATURES_REPO_ID = (
    "galaxy-sbi-challenge/sbi-galaxy-blinded-features"
)
DEFAULT_PRIVATE_TARGETS_REPO_ID = (
    "galaxy-sbi-challenge/sbi-galaxy-blinded-targets"
)
RELEASE_PREFIX = "SBI_DATA_RELEASE"
VERSION_PATTERN = re.compile(r"^v(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)$")
RELEASE_VERSION_PATTERN = re.compile(r"\bversion=(v\d+\.\d+\.\d+)\b")


@dataclass(frozen=True)
class ReleaseTarget:
    label: str
    repo_id: str
    files: tuple[Path, ...]


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Upload challenge data files to main and preserve releases through "
            "standardized commit messages plus HDF5 dataset_version attributes."
        )
    )
    parser.add_argument("--dev-repo-id", default=DEFAULT_DEV_REPO_ID)
    parser.add_argument(
        "--blinded-features-repo-id",
        default=DEFAULT_BLINDED_FEATURES_REPO_ID,
    )
    parser.add_argument(
        "--private-targets-repo-id",
        default=DEFAULT_PRIVATE_TARGETS_REPO_ID,
    )
    parser.add_argument("--repo-type", default="dataset")
    parser.add_argument("--revision", default="main", help="Branch to upload to.")
    parser.add_argument(
        "--version",
        help="Release version. Defaults to the next patch version across repos.",
    )
    parser.add_argument(
        "--bump",
        choices=("major", "minor", "patch"),
        default="patch",
        help="Version component to increment when --version is omitted.",
    )
    parser.add_argument("--dev", nargs="*", type=Path, default=())
    parser.add_argument("--blinded-features", nargs="*", type=Path, default=())
    parser.add_argument("--private-targets", nargs="*", type=Path, default=())
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the planned release without changing files or Hugging Face.",
    )
    args = parser.parse_args()

    targets = [
        ReleaseTarget("dev", args.dev_repo_id, tuple(_validate_files(args.dev))),
        ReleaseTarget(
            "blinded-features",
            args.blinded_features_repo_id,
            tuple(_validate_files(args.blinded_features)),
        ),
        ReleaseTarget(
            "private-targets",
            args.private_targets_repo_id,
            tuple(_validate_files(args.private_targets)),
        ),
    ]
    targets = [target for target in targets if target.files]
    if not targets:
        raise SystemExit(
            "Provide at least one file via --dev, --blinded-features, "
            "or --private-targets."
        )

    api = HfApi()
    existing_versions = existing_release_versions(
        api=api,
        repo_ids=[target.repo_id for target in targets],
        repo_type=args.repo_type,
    )
    version = args.version or next_version(existing_versions, args.bump)
    if not VERSION_PATTERN.fullmatch(version):
        raise SystemExit(f"Version must look like vMAJOR.MINOR.PATCH; got {version!r}.")
    if version in existing_versions:
        raise SystemExit(f"Release version {version!r} already exists.")

    planned_challenges = sorted(
        {
            challenge_id_from_filename(file_path.name)
            for target in targets
            for file_path in target.files
        }
    )

    print(f"release version: {version}")
    print(f"upload revision: {args.revision}")
    print(f"challenges: {', '.join(planned_challenges)}")
    for target in targets:
        print(f"{target.label}: {target.repo_id}")
        for file_path in target.files:
            print(f"  {file_path} -> {file_path.name}")

    if args.dry_run:
        return

    for target in targets:
        for file_path in target.files:
            write_hdf5_release_attrs(file_path, version)

        commit_info = upload_files(
            api=api,
            target=target,
            repo_type=args.repo_type,
            revision=args.revision,
            version=version,
            challenges=planned_challenges,
        )
        print(f"released {target.repo_id}@{commit_info.oid}")


def _validate_files(paths: list[Path] | tuple[Path, ...]) -> list[Path]:
    files: list[Path] = []
    for path in paths:
        if not path.exists():
            raise SystemExit(f"File does not exist: {path}")
        if not path.is_file():
            raise SystemExit(f"Path is not a file: {path}")
        files.append(path)
    return files


def existing_release_versions(
    *,
    api: HfApi,
    repo_ids: list[str],
    repo_type: str,
) -> set[str]:
    versions: set[str] = set()
    for repo_id in repo_ids:
        for commit in api.list_repo_commits(repo_id, repo_type=repo_type):
            version = release_version(commit.title) or release_version(commit.message)
            if version:
                versions.add(version)
    return versions


def next_version(existing_versions: set[str], bump: str) -> str:
    parsed = []
    for version in existing_versions:
        match = VERSION_PATTERN.fullmatch(version)
        if match:
            parsed.append(tuple(int(match.group(name)) for name in ("major", "minor", "patch")))

    major, minor, patch = max(parsed) if parsed else (1, 0, -1)
    if bump == "major":
        major, minor, patch = major + 1, 0, 0
    elif bump == "minor":
        minor, patch = minor + 1, 0
    else:
        patch += 1
    return f"v{major}.{minor}.{patch}"


def release_version(message: str | None) -> str | None:
    if not message or RELEASE_PREFIX not in message:
        return None
    match = RELEASE_VERSION_PATTERN.search(message)
    if match is None:
        return None
    return match.group(1)


def upload_files(
    *,
    api: HfApi,
    target: ReleaseTarget,
    repo_type: str,
    revision: str,
    version: str,
    challenges: list[str],
):
    operations = [
        CommitOperationAdd(path_or_fileobj=file_path, path_in_repo=file_path.name)
        for file_path in target.files
    ]
    message = (
        f"{RELEASE_PREFIX} version={version} "
        f"challenges={','.join(challenges)} files="
        f"{','.join(file_path.name for file_path in target.files)}"
    )
    return api.create_commit(
        repo_id=target.repo_id,
        repo_type=repo_type,
        revision=revision,
        operations=operations,
        commit_message=message,
    )


def write_hdf5_release_attrs(path: Path, version: str) -> None:
    with h5py.File(path, "a") as hdf:
        existing = hdf.attrs.get("dataset_version")
        if existing is not None:
            existing = decode_attr(existing)
            if existing != version:
                print(f"updating {path} dataset_version: {existing!r} -> {version!r}")
        hdf.attrs["dataset_version"] = version
        hdf.attrs["challenge_id"] = challenge_id_from_filename(path.name)


def decode_attr(value: object) -> object:
    if isinstance(value, bytes):
        return value.decode("utf-8")
    return value


def challenge_id_from_filename(filename: str) -> str:
    name = Path(filename).name
    if name.startswith("SBI_DC_"):
        remainder = name.removeprefix("SBI_DC_")
        if "_" in remainder:
            return remainder.split("_", maxsplit=1)[0]
    return name.removesuffix(".h5").removesuffix(".hdf5")


if __name__ == "__main__":
    main()