Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Sync Label Studio data directory and logs with R2.""" | |
| from __future__ import annotations | |
| import argparse | |
| import sys | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from r2_storage import PREFIX_LS, download_prefix, sync_directory | |
| def log(message: str, *, file: object = sys.stdout) -> None: | |
| timestamp = datetime.now(timezone.utc).isoformat(timespec="seconds").replace( | |
| "+00:00", "Z" | |
| ) | |
| print(f"[{timestamp}] [r2-sync] {message}", file=file, flush=True) | |
| 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) | |
| log(f"Pulled {n} objects from r2://…/{PREFIX_LS}/ -> {data_dir}") | |
| if args.direction in ("push", "both"): | |
| if not data_dir.exists(): | |
| log(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") | |
| log(f"Pushed {n} files to r2://…/{PREFIX_LS}/") | |
| if __name__ == "__main__": | |
| main() | |