scripts/sync_label_studio_r2.py

#!/usr/bin/env python3
"""Sync Label Studio data directory and logs with R2."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from r2_storage import PREFIX_LS, download_prefix, sync_directory


def main() -> None:
    parser = argparse.ArgumentParser(description="Sync Label Studio state with R2")
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path("/label-studio/data"),
    )
    parser.add_argument("direction", choices=["pull", "push", "both"])
    args = parser.parse_args()
    data_dir: Path = args.data_dir

    if args.direction in ("pull", "both"):
        n = download_prefix(PREFIX_LS, data_dir)
        print(f"Pulled {n} objects from r2://…/{PREFIX_LS}/ -> {data_dir}")

    if args.direction in ("push", "both"):
        if not data_dir.exists():
            print(f"No data dir at {data_dir}", file=sys.stderr)
            sys.exit(1)
        n = sync_directory(data_dir, PREFIX_LS)
        logs = data_dir.parent / "logs"
        if logs.is_dir():
            n += sync_directory(logs, f"{PREFIX_LS}/logs")
        print(f"Pushed {n} files to r2://…/{PREFIX_LS}/")


if __name__ == "__main__":
    main()