Download app/hub.py from agent-collaborations/dev-bucket-sync: direct link, hf CLI and curl.
- Browser
- Download file 24.8 kB
-
https://huggingface.co/spaces/agent-collaborations/dev-bucket-sync/resolve/main/app/hub.py
- Command line
-
hf download hf://spaces/agent-collaborations/dev-bucket-sync/app/hub.py
-
curl -L -o hub.py https://huggingface.co/spaces/agent-collaborations/dev-bucket-sync/resolve/main/app/hub.py
24.8 kB
| """Wrapper over huggingface_hub's bucket API.""" | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| import os | |
| import subprocess | |
| import sys | |
| import tempfile | |
| import threading | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Iterable | |
| try: # huggingface_hub 2.x ships its own httpx fork; 1.x uses httpx itself. | |
| import httpx2 as httpx # type: ignore[import-not-found] | |
| except ImportError: | |
| import httpx # httpx is not in requirements.txt: it comes with huggingface_hub | |
| from huggingface_hub import ( | |
| batch_bucket_files, | |
| bucket_info, | |
| create_bucket, | |
| download_bucket_files, | |
| list_bucket_tree, | |
| whoami, | |
| ) | |
| from huggingface_hub.constants import ENDPOINT | |
| from huggingface_hub.errors import ( | |
| EntryNotFoundError, | |
| HfHubHTTPError, | |
| RepositoryNotFoundError, | |
| ) | |
| from huggingface_hub.utils import build_hf_headers, get_session | |
| from app.caller_write import classify as classify_caller_write_error | |
| from app.config import Settings | |
| from app.naming import SourceURI, parse_source_uri | |
| log = logging.getLogger(__name__) | |
| class ListedFile: | |
| rel_path: str | |
| size: int | |
| xet_hash: str | None = None | |
| class HubIdentity: | |
| username: str | |
| orgs: set[str] | |
| email: str | None = None | |
| class OrgMemberRole: | |
| user: str | |
| role: str | |
| class HubUnreachable(Exception): | |
| """The Hub answered 5xx/429 or not at all: transient, safe to retry.""" | |
| def _status(e: Exception) -> int | None: | |
| resp = getattr(e, "response", None) | |
| return getattr(resp, "status_code", None) | |
| def _transient(e: Exception) -> bool: | |
| if isinstance(e, httpx.TransportError): | |
| return True | |
| status = _status(e) | |
| return isinstance(e, HfHubHTTPError) and (status is None or status >= 500 or status == 429) | |
| def _raise_caller_write_error(e: Exception) -> None: | |
| """A Hub call refused for a caller's token raises PermissionError; an | |
| outage raises HubUnreachable; anything else returns for the caller to | |
| re-raise. Uses the caller-write child's classifier, so a refusal is only | |
| ever an explicit 401/403.""" | |
| if classify_caller_write_error(e)["result"] == "forbidden": | |
| raise PermissionError(f"the caller's token was refused (HTTP {_status(e) or '401/403'})") from e | |
| if isinstance(e, ConnectionError) or _transient(e): | |
| raise HubUnreachable(type(e).__name__) from e | |
| # ββ Caller-token bucket writes run in a child process (see app/caller_write.py) | |
| _CALLER_WRITE_CMD = [sys.executable, "-m", "app.caller_write"] | |
| _CALLER_WRITE_CWD = Path(__file__).resolve().parent.parent # the dir holding app/ | |
| CALLER_WRITE_TIMEOUT_S = 60.0 | |
| # Registration is limited to a few per minute per user, so a handful of slots | |
| # is plenty; a request that cannot get one in time gets a 503, not a pile-up. | |
| CALLER_WRITE_SLOTS = threading.BoundedSemaphore(4) | |
| CALLER_WRITE_SLOT_WAIT_S = 20.0 | |
| _ADMIN_TOKEN_VARS = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACEHUB_API_TOKEN") | |
| def _caller_write_env() -> dict[str, str]: | |
| """The Space's environment minus its admin credential: the child acts only | |
| with the token it is handed on stdin.""" | |
| env = {k: v for k, v in os.environ.items() if k not in _ADMIN_TOKEN_VARS} | |
| env["HF_TOKEN_PATH"] = os.devnull # no fallback to a stored token file | |
| env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" | |
| return env | |
| def _run_caller_write(bucket: str, path: str, text: str, token: str) -> dict: | |
| """Run one caller-token write in a child and return its structured result. | |
| Timeouts, crashes and unreadable output come back as "failed".""" | |
| if not CALLER_WRITE_SLOTS.acquire(timeout=CALLER_WRITE_SLOT_WAIT_S): | |
| log.warning("caller write: no free slot (bucket=%s)", bucket) | |
| return {"result": "failed", "type": "NoSlot", "status": None} | |
| try: | |
| try: | |
| proc = subprocess.run( | |
| _CALLER_WRITE_CMD, | |
| input=json.dumps({"token": token, "bucket": bucket, "path": path, "text": text}), | |
| capture_output=True, text=True, timeout=CALLER_WRITE_TIMEOUT_S, | |
| cwd=_CALLER_WRITE_CWD, env=_caller_write_env(), | |
| ) | |
| except subprocess.TimeoutExpired: | |
| log.warning("caller write: timed out after %ss (bucket=%s)", CALLER_WRITE_TIMEOUT_S, bucket) | |
| return {"result": "failed", "type": "Timeout", "status": None} | |
| finally: | |
| CALLER_WRITE_SLOTS.release() | |
| # The child's stderr is never logged: only its size, for diagnosis. | |
| lines = proc.stdout.strip().splitlines() | |
| try: | |
| result = json.loads(lines[-1]) | |
| if result.get("result") not in ("ok", "forbidden", "failed"): | |
| raise ValueError | |
| except (IndexError, ValueError, AttributeError): | |
| log.warning( | |
| "caller write: no result from child (bucket=%s exit=%s stderr_bytes=%d)", | |
| bucket, proc.returncode, len(proc.stderr or ""), | |
| ) | |
| return {"result": "failed", "type": "ChildCrashed", "status": None} | |
| return result | |
| class HubClient: | |
| def __init__(self, settings: Settings): | |
| self._settings = settings | |
| def _token(self) -> str: | |
| return self._settings.resolved_token() | |
| # βββββββββββββββββββββββββ Bucket existence & identity βββββββββββββββββββββββββ | |
| def bucket_exists(self, bucket: str) -> bool: | |
| """Only a 404 means "no such bucket"; a Hub outage raises | |
| HubUnreachable instead of looking like a missing bucket.""" | |
| try: | |
| bucket_info(bucket, token=self._token) | |
| return True | |
| except (HfHubHTTPError, httpx.TransportError) as e: | |
| if isinstance(e, RepositoryNotFoundError) or _status(e) == 404: | |
| return False | |
| if _transient(e): | |
| raise HubUnreachable(str(e)) from e | |
| raise | |
| def create_bucket_as(self, bucket: str, token: str) -> None: | |
| """Create `bucket` with the caller's token so the caller owns it (only | |
| its creator and org admins can then write there). Default visibility. | |
| Raises PermissionError if the token may not create it.""" | |
| try: | |
| create_bucket(bucket, exist_ok=True, token=token) | |
| except (HfHubHTTPError, httpx.TransportError, ConnectionError) as e: | |
| _raise_caller_write_error(e) | |
| raise | |
| def caller_may_write(self, bucket: str, token: str) -> bool | None: | |
| """Ask the Hub for a bucket write token with the caller's token, over | |
| plain HTTP (the request batch_bucket_files makes first). False on | |
| 401/403, True on success, None when it can't tell. Only a pre-check | |
| for clean errors: the write itself remains the ownership proof. The | |
| response carries a Xet access token, so only its status is read.""" | |
| url = f"{ENDPOINT}/api/buckets/{bucket}/xet-write-token" | |
| try: | |
| resp = get_session().get(url, headers=build_hf_headers(token=token), timeout=10) | |
| except Exception as e: | |
| log.warning("caller write pre-check failed (bucket=%s type=%s)", bucket, type(e).__name__) | |
| return None | |
| if resp.status_code in (401, 403): | |
| return False | |
| return True if 200 <= resp.status_code < 300 else None | |
| def write_text_as(self, bucket: str, path: str, text: str, token: str) -> None: | |
| """Write one file with the caller's token. Raises PermissionError when | |
| the Hub refuses (the bucket is not the caller's), HubUnreachable for | |
| anything else. The upload runs in a child process: a refused Xet | |
| upload poisons the process-wide Xet session (app/caller_write.py).""" | |
| if self.caller_may_write(bucket, token) is False: | |
| raise PermissionError(f"the caller's token may not write {bucket}") | |
| result = _run_caller_write(bucket, path, text, token) | |
| if result["result"] == "ok": | |
| return | |
| if result["result"] == "forbidden": | |
| raise PermissionError(f"the caller's token may not write {bucket} (HTTP {result.get('status')})") | |
| log.warning( | |
| "caller write failed (bucket=%s type=%s status=%s)", | |
| bucket, result.get("type"), result.get("status"), | |
| ) | |
| raise HubUnreachable(f"caller write to {bucket} failed") | |
| def bucket_author(self, bucket: str) -> str | None: | |
| """Return the `author` field from BucketInfo (the org name for org buckets). | |
| Not a true creator-of-record; for identity binding we rely on the | |
| whoami-at-registration flow plus a handshake file in the agent's bucket. | |
| """ | |
| try: | |
| info = bucket_info(bucket, token=self._token) | |
| except (RepositoryNotFoundError, HfHubHTTPError) as e: | |
| log.debug("bucket_author(%s) failed: %s", bucket, e) | |
| return None | |
| return getattr(info, "author", None) | |
| def whoami_for_token(self, token: str) -> str: | |
| info = whoami(token=token) | |
| if isinstance(info, dict) and info.get("name"): | |
| return info["name"] | |
| raise ValueError("whoami did not return a `name` field") | |
| def whoami_identity(self, token: str) -> HubIdentity: | |
| """Resolve a caller token to HF identity facts used by human posts. | |
| OAuth tokens expose org membership but not roleInOrg. When the Space's | |
| OAuth app requests the email scope, the email lets the organizer gate | |
| perform a targeted org-member lookup instead of scanning the full org. | |
| """ | |
| try: | |
| info = whoami(token=token) | |
| except (HfHubHTTPError, httpx.TransportError) as e: | |
| if _transient(e): | |
| raise HubUnreachable(str(e)) from e | |
| raise | |
| if not isinstance(info, dict) or not info.get("name"): | |
| raise ValueError("whoami did not return a `name` field") | |
| orgs = { | |
| o["name"] | |
| for o in (info.get("orgs") or []) | |
| if isinstance(o, dict) and o.get("name") | |
| } | |
| email = info.get("email") | |
| if not isinstance(email, str) or not email.strip(): | |
| email = None | |
| return HubIdentity(info["name"], orgs, email.strip() if email else None) | |
| def whoami_user_and_orgs(self, token: str) -> tuple[str, set[str]]: | |
| """Resolve a caller token to (hf_user, org names). Used by existing | |
| identity gates that do not need optional email.""" | |
| identity = self.whoami_identity(token) | |
| return identity.username, identity.orgs | |
| def org_member_role_by_email(self, org: str, email: str) -> OrgMemberRole | None: | |
| """Resolve one org member by email and return its role, if available. | |
| The members endpoint supports an email filter for orgs with a matching | |
| Organization email domain or SSO allowed domain. This is the scalable | |
| organizer check path: one request for the caller instead of listing | |
| every org member. Transport errors propagate so callers can fall back | |
| or fail closed. | |
| """ | |
| session = get_session() | |
| headers = build_hf_headers(token=self._token) | |
| url = f"{ENDPOINT}/api/organizations/{org}/members" | |
| resp = session.get( | |
| url, | |
| headers=headers, | |
| params={"email": email, "limit": 1}, | |
| timeout=10, | |
| ) | |
| resp.raise_for_status() | |
| page = resp.json() | |
| if not isinstance(page, list) or not page: | |
| return None | |
| member = page[0] | |
| if not isinstance(member, dict): | |
| return None | |
| user, role = member.get("user"), member.get("role") | |
| if not isinstance(user, str) or not isinstance(role, str): | |
| return None | |
| return OrgMemberRole(user=user, role=role) | |
| def org_member_roles(self, org: str) -> dict[str, str]: | |
| """Map {lowercased username: org role} for every member of ``org``. | |
| whoami omits a caller's org role for OAuth tokens, so the organizer | |
| gate can't read it from the caller's own token; instead the Space | |
| looks the role up here with its admin token. Lists the members | |
| endpoint, following Link pagination and falling back to offset-style | |
| pagination if needed. Raises on transport error so the caller | |
| can fail closed rather than treat an outage as "not an organizer". | |
| """ | |
| session = get_session() | |
| headers = build_hf_headers(token=self._token) | |
| url = f"{ENDPOINT}/api/organizations/{org}/members" | |
| page_size = 100 | |
| roles: dict[str, str] = {} | |
| offset = 0 | |
| params: dict[str, int] | None = {"limit": page_size, "offset": offset} | |
| while True: | |
| resp = session.get( | |
| url, | |
| headers=headers, | |
| params=params, | |
| timeout=10, | |
| ) | |
| resp.raise_for_status() | |
| page = resp.json() | |
| if not isinstance(page, list) or not page: | |
| break | |
| before = len(roles) | |
| for m in page: | |
| if not isinstance(m, dict): | |
| continue | |
| user, role = m.get("user"), m.get("role") | |
| if user and role: | |
| roles[user.lower()] = role | |
| next_url = resp.links.get("next", {}).get("url") | |
| if next_url: | |
| url = next_url | |
| params = None | |
| continue | |
| if len(page) < page_size or len(roles) == before: | |
| break | |
| offset += page_size | |
| params = {"limit": page_size, "offset": offset} | |
| return roles | |
| # βββββββββββββββββββββββββ Reads βββββββββββββββββββββββββ | |
| def read_bytes(self, uri: SourceURI | str) -> bytes: | |
| parsed = uri if isinstance(uri, SourceURI) else parse_source_uri(uri) | |
| if parsed is None: | |
| raise ValueError(f"invalid source URI: {uri}") | |
| bucket = f"{parsed.org}/{parsed.bucket}" | |
| return self._download_one(bucket, parsed.path) | |
| def read_text(self, uri: SourceURI | str) -> str: | |
| return self.read_bytes(uri).decode("utf-8") | |
| def read_central_bytes(self, target_path: str) -> bytes: | |
| return self._download_one(self._settings.central_bucket, target_path) | |
| def read_central_text(self, target_path: str) -> str: | |
| return self.read_central_bytes(target_path).decode("utf-8") | |
| def read_central_bytes_optional(self, target_path: str) -> bytes | None: | |
| """Read a central-bucket file, distinguishing a genuinely missing file | |
| (returns None) from a transport/HTTP error (propagates). | |
| Unlike ``read_central_bytes`` β which flattens both cases to | |
| ``FileNotFoundError`` β this lets read-modify-write callers fail SAFE on | |
| a storage blip: skip the update rather than overwrite a live file with a | |
| fresh, near-empty one. Mirrors ``read_audit_bytes`` for the central bucket. | |
| """ | |
| return self._download_optional(self._settings.central_bucket, target_path) | |
| def read_text_optional(self, uri: SourceURI | str) -> str | None: | |
| """``read_central_bytes_optional`` for any bucket: None only when the | |
| file is genuinely missing; any other failure propagates.""" | |
| parsed = uri if isinstance(uri, SourceURI) else parse_source_uri(uri) | |
| if parsed is None: | |
| raise ValueError(f"invalid source URI: {uri}") | |
| data = self._download_optional(f"{parsed.org}/{parsed.bucket}", parsed.path) | |
| return None if data is None else data.decode("utf-8") | |
| def _download_optional(self, bucket: str, remote_path: str) -> bytes | None: | |
| with tempfile.TemporaryDirectory() as td: | |
| local = Path(td) / "f" | |
| try: | |
| download_bucket_files( | |
| bucket_id=bucket, | |
| files=[(remote_path, str(local))], | |
| raise_on_missing_files=True, | |
| token=self._token, | |
| ) | |
| except EntryNotFoundError: | |
| return None | |
| return local.read_bytes() | |
| def _download_one(self, bucket: str, remote_path: str) -> bytes: | |
| with tempfile.TemporaryDirectory() as td: | |
| local = Path(td) / "f" | |
| try: | |
| download_bucket_files( | |
| bucket_id=bucket, | |
| files=[(remote_path, str(local))], | |
| raise_on_missing_files=True, | |
| token=self._token, | |
| ) | |
| except (EntryNotFoundError, HfHubHTTPError) as e: | |
| raise FileNotFoundError(f"{bucket}/{remote_path}: {e}") | |
| return local.read_bytes() | |
| def download_many(self, bucket: str, remote_paths: list[str]) -> dict[str, bytes]: | |
| """Batch-download files, returning {remote_path: bytes}. | |
| Missing or failed entries are simply absent from the result β callers | |
| (the read model, the backfill script) treat absence as transient and | |
| retry on a later pass. Chunked so a multi-thousand-file cold fill | |
| doesn't ride on a single oversized call. | |
| """ | |
| out: dict[str, bytes] = {} | |
| chunk_size = 500 | |
| for start in range(0, len(remote_paths), chunk_size): | |
| chunk = remote_paths[start : start + chunk_size] | |
| with tempfile.TemporaryDirectory() as td: | |
| pairs = [(remote, str(Path(td) / str(i))) for i, remote in enumerate(chunk)] | |
| try: | |
| download_bucket_files( | |
| bucket_id=bucket, | |
| files=pairs, | |
| raise_on_missing_files=False, | |
| token=self._token, | |
| ) | |
| except (EntryNotFoundError, HfHubHTTPError) as e: | |
| log.warning( | |
| "download_many(%s, %d files) failed: %s", bucket, len(chunk), e | |
| ) | |
| continue | |
| for remote, local in pairs: | |
| p = Path(local) | |
| if p.exists(): | |
| out[remote] = p.read_bytes() | |
| return out | |
| def list_central_dir(self, prefix: str) -> list[ListedFile]: | |
| return self._list(self._settings.central_bucket, prefix) | |
| def list_bucket_dir(self, bucket: str, prefix: str) -> list[ListedFile]: | |
| return self._list(bucket, prefix) | |
| def _list(self, bucket: str, prefix: str) -> list[ListedFile]: | |
| out: list[ListedFile] = [] | |
| try: | |
| for entry in list_bucket_tree( | |
| bucket_id=bucket, | |
| prefix=prefix or None, | |
| recursive=True, | |
| token=self._token, | |
| ): | |
| if getattr(entry, "type", None) == "file": | |
| out.append( | |
| ListedFile( | |
| rel_path=entry.path, | |
| size=entry.size or 0, | |
| xet_hash=getattr(entry, "xet_hash", None), | |
| ) | |
| ) | |
| except (RepositoryNotFoundError, HfHubHTTPError) as e: | |
| log.debug("list(%s, %s) failed: %s", bucket, prefix, e) | |
| return out | |
| # βββββββββββββββββββββββββ Writes (central bucket) βββββββββββββββββββββββββ | |
| def write_bytes_central(self, target_path: str, data: bytes) -> None: | |
| batch_bucket_files( | |
| bucket_id=self._settings.central_bucket, | |
| add=[(data, target_path)], | |
| token=self._token, | |
| ) | |
| def write_text_central(self, target_path: str, text: str) -> None: | |
| self.write_bytes_central(target_path, text.encode("utf-8")) | |
| def write_many_central(self, items: list[tuple[bytes, str]]) -> None: | |
| """Write several central-bucket files in one batch call. | |
| Used to land a message and its inbox fan-out copies together (Β§16.4): | |
| one storage round trip, no window where board and inbox diverge. | |
| """ | |
| if not items: | |
| return | |
| batch_bucket_files( | |
| bucket_id=self._settings.central_bucket, | |
| add=list(items), | |
| token=self._token, | |
| ) | |
| def delete_central(self, target_path: str) -> None: | |
| """Delete one central-bucket file. The only deleting write in the | |
| system: channel unsubscribe removes the member marker | |
| (CHANNELS_DESIGN.md Β§3.3). Everything else stays append-only.""" | |
| batch_bucket_files( | |
| bucket_id=self._settings.central_bucket, | |
| delete=[target_path], | |
| token=self._token, | |
| ) | |
| def write_bytes_to_bucket(self, bucket: str, target_path: str, data: bytes) -> None: | |
| batch_bucket_files(bucket_id=bucket, add=[(data, target_path)], token=self._token) | |
| def write_text_to_bucket(self, bucket: str, target_path: str, text: str) -> None: | |
| self.write_bytes_to_bucket(bucket, target_path, text.encode("utf-8")) | |
| def append_jsonl_audit(self, target_path: str, line: str) -> None: | |
| """Append to the audit log in the private (out-of-org) audit bucket.""" | |
| self._append_jsonl(self._settings.audit_bucket, target_path, line) | |
| def read_audit_bytes(self, target_path: str) -> bytes | None: | |
| """Read a file from the private audit bucket. | |
| Returns the bytes, or None if the file genuinely does not exist. Unlike | |
| read_central_bytes, transport/HTTP errors PROPAGATE rather than being | |
| flattened to "missing" β so callers (e.g. the job quota) can fail closed | |
| on a storage outage instead of treating it as an empty ledger. | |
| """ | |
| with tempfile.TemporaryDirectory() as td: | |
| local = Path(td) / "f" | |
| try: | |
| download_bucket_files( | |
| bucket_id=self._settings.audit_bucket, | |
| files=[(target_path, str(local))], | |
| raise_on_missing_files=True, | |
| token=self._token, | |
| ) | |
| except EntryNotFoundError: | |
| return None | |
| return local.read_bytes() | |
| def write_bytes_audit(self, target_path: str, data: bytes) -> None: | |
| batch_bucket_files( | |
| bucket_id=self._settings.audit_bucket, | |
| add=[(data, target_path)], | |
| token=self._token, | |
| ) | |
| def _append_jsonl(self, bucket: str, target_path: str, line: str) -> None: | |
| try: | |
| existing = self._download_one(bucket, target_path) | |
| except FileNotFoundError: | |
| existing = b"" | |
| if existing and not existing.endswith(b"\n"): | |
| existing += b"\n" | |
| batch_bucket_files( | |
| bucket_id=bucket, | |
| add=[(existing + line.encode("utf-8") + b"\n", target_path)], | |
| token=self._token, | |
| ) | |
| # βββββββββββββββββββββββββ Cross-bucket copy βββββββββββββββββββββββββ | |
| def copy_file_to_central(self, src_bucket: str, src_xet_hash: str, dest_path: str) -> None: | |
| """Hash-copy a single source file into the central bucket (bytes never | |
| transit the Space). The caller passes the source's xet hash β taken from a | |
| listing it already holds β so there is no extra lookup here.""" | |
| if not src_xet_hash: | |
| raise RuntimeError(f"missing xet_hash for copy to {dest_path}") | |
| batch_bucket_files( | |
| bucket_id=self._settings.central_bucket, | |
| copy=[("bucket", src_bucket, src_xet_hash, dest_path)], | |
| token=self._token, | |
| ) | |
| def copy_tree_to_central( | |
| self, src_bucket: str, src_prefix: str, dest_prefix: str | |
| ) -> Iterable[tuple[str, str, int]]: | |
| files = self._list(src_bucket, src_prefix) | |
| if not files: | |
| return | |
| prefix = src_prefix.rstrip("/") | |
| copy_ops: list[tuple[str, str, str, str]] = [] | |
| results: list[tuple[str, str, int]] = [] | |
| for f in files: | |
| if not f.xet_hash: | |
| raise RuntimeError(f"missing xet_hash for source file: {f.rel_path}") | |
| rel = f.rel_path[len(prefix) + 1 :] if prefix and f.rel_path.startswith(prefix + "/") else f.rel_path | |
| dest_path = f"{dest_prefix.rstrip('/')}/{rel}" | |
| copy_ops.append(("bucket", src_bucket, f.xet_hash, dest_path)) | |
| results.append((f.rel_path, dest_path, f.size)) | |
| batch_bucket_files( | |
| bucket_id=self._settings.central_bucket, | |
| copy=copy_ops, | |
| token=self._token, | |
| ) | |
| for r in results: | |
| yield r | |