# crvs_sync.py -- resumable, rate-limited, interrupt-safe Hugging Face folder sync. # Identical across NB01-NB05 so the cadence rules are enforced in exactly one place. # * push at most once per PUSH_INTERVAL_S (default 30 min) # * push immediately when a stage finishes -> sync.stage_done("name") # * push immediately when execution is stopped -> SIGINT / SIGTERM / atexit # * one upload_folder call per flush, behind a token bucket, backing off on 429 # * resume by pulling the run folder back on startup import os, json, time, random, threading, atexit, signal from pathlib import Path from datetime import datetime, timezone class TokenBucket: # capacity = requests per hour, refilled continuously def __init__(self, per_hour=120): self.capacity = float(per_hour); self.tokens = float(per_hour) self.rate = per_hour / 3600.0; self.t = time.monotonic() self.lock = threading.Lock() def take(self, n=1, block=True, timeout=1200): deadline = time.monotonic() + timeout while True: with self.lock: now = time.monotonic() self.tokens = min(self.capacity, self.tokens + (now - self.t) * self.rate) self.t = now if self.tokens >= n: self.tokens -= n; return True need = (n - self.tokens) / self.rate if not block or time.monotonic() + need > deadline: return False time.sleep(min(need, 5.0)) class HFSync: def __init__(self, repo_id, local_dir, token, repo_type="dataset", private=False, run_id="run", push_interval_s=1800, max_req_hour=120, retry_max=6, verbose=True): from huggingface_hub import HfApi self.api = HfApi(token=token); self.token = token self.repo_id = repo_id; self.repo_type = repo_type; self.private = private self.run_id = run_id self.local = Path(local_dir); self.local.mkdir(parents=True, exist_ok=True) self.interval = push_interval_s self.bucket = TokenBucket(max_req_hour) self.retry_max = retry_max; self.verbose = verbose self._last_push = 0.0 self._flag = threading.Event(); self._stop = threading.Event() self._lock = threading.Lock() self._pushes = 0; self._failures = 0 self.history = self.local / "history.jsonl" self.state_path = self.local / "state.json" self._ensure_repo(); self._install_handlers() self._thread = threading.Thread(target=self._loop, daemon=True, name="hf-uploader") self._thread.start() self.log("sync_started", repo=self.repo_id, private=self.private) def _ensure_repo(self): from huggingface_hub import create_repo create_repo(self.repo_id, repo_type=self.repo_type, private=self.private, exist_ok=True, token=self.token) if not self.private: try: self.api.update_repo_visibility(self.repo_id, private=False, repo_type=self.repo_type, token=self.token) except Exception: pass @property def url(self): kind = "datasets/" if self.repo_type == "dataset" else "" return "https://huggingface.co/" + kind + self.repo_id def log(self, event, **kw): rec = {"ts": datetime.now(timezone.utc).isoformat(), "run": self.run_id, "event": event} rec.update(kw) try: with open(self.history, "a") as f: f.write(json.dumps(rec, default=str) + "\n") except Exception: pass if self.verbose and event not in ("heartbeat",): print(" [" + event + "] " + " ".join(f"{k}={v}" for k, v in kw.items())) def save_state(self, state): tmp = self.state_path.with_suffix(".tmp") tmp.write_text(json.dumps(state, indent=2, default=str)); tmp.replace(self.state_path) def load_state(self, default=None): if self.state_path.exists(): try: return json.loads(self.state_path.read_text()) except Exception: pass return default if default is not None else {} def pull(self, allow_patterns=None, into=None): from huggingface_hub import snapshot_download try: self.bucket.take(1) p = snapshot_download(self.repo_id, repo_type=self.repo_type, token=self.token, local_dir=str(into or self.local), allow_patterns=allow_patterns) self.log("resume_pull_ok", path=str(p)); return True except Exception as e: self.log("resume_pull_empty", err=type(e).__name__); return False def stage_done(self, name, **kw): self.log("stage_done", stage=name, **kw); self._flag.set() def _do_upload(self, msg): from huggingface_hub import upload_folder for attempt in range(self.retry_max): if not self.bucket.take(1, block=True, timeout=1800): self.log("rate_limited_giveup"); return False try: upload_folder(folder_path=str(self.local), repo_id=self.repo_id, repo_type=self.repo_type, token=self.token, commit_message=msg, ignore_patterns=["*.tmp", "**/__pycache__/**", ".git*", "*.lock", ".cache/**"]) self._pushes += 1; self._last_push = time.time() self.log("push_ok", n=self._pushes, msg=msg); return True except Exception as e: self._failures += 1 wait = min(300, (2 ** attempt) * 5) * (0.7 + 0.6 * random.random()) self.log("push_retry", attempt=attempt + 1, err=f"{type(e).__name__}: {e}", sleep=round(wait, 1)) time.sleep(wait) self.log("push_failed_permanently", msg=msg); return False def flush(self, final=False, msg=None): with self._lock: stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M") m = msg or ((self.run_id + " final") if final else (self.run_id + " @ " + stamp + "Z")) ok = self._do_upload(m); self._flag.clear(); return ok def _loop(self): while not self._stop.is_set(): self._stop.wait(20) if self._stop.is_set(): break due = (time.time() - self._last_push) >= self.interval want = self._flag.is_set() if due or want: try: tag = "stage" if want else "periodic" self.flush(msg=self.run_id + " " + tag + " @ " + datetime.now(timezone.utc).strftime("%H:%M") + "Z") except Exception as e: self.log("loop_error", err=str(e)) def _install_handlers(self): def handler(signum, frame): self.log("interrupt", signal=int(signum)) try: self.flush(final=True, msg=self.run_id + " interrupted (sig " + str(signum) + ")") finally: if signum == signal.SIGINT: raise KeyboardInterrupt for sig in (signal.SIGINT, signal.SIGTERM): try: signal.signal(sig, handler) except Exception: pass atexit.register(self.close) def close(self): if self._stop.is_set(): return self.log("closing"); self._stop.set() try: self.flush(final=True) except Exception: pass