rpx_benchmark.box_upload

  1"""Box upload primitives — used by every task runner and the post-hoc
  2sync script.
  3
  4Two public entry points:
  5
  6* :func:`upload_tree(local_dir, remote_path, *, root_folder_id)` —
  7  recursively mirror a directory to Box. Size-matched idempotent; safe
  8  to re-run.
  9* :func:`upload_run_dir(out_dir, *, task, model_name, split, ...)` —
 10  the high-level wrapper that lays a per-run output directory down at
 11  the canonical ``<root_folder_id>/<task>/<model>/<split>/`` path the
 12  team agreed on.
 13
 14Auth has two modes, in priority order:
 15
 161. **OAuth 2.0 with refresh-token rotation** (preferred for long runs).
 17   Set ``BOX_CLIENT_ID`` + ``BOX_CLIENT_SECRET`` (+ ``BOX_REDIRECT_URI``,
 18   default ``http://localhost:8765/callback``) and run::
 19
 20       python -m rpx_benchmark.box_upload login
 21
 22   once to capture the refresh token. Tokens are stored at
 23   ``~/.config/rpx_benchmark/box_tokens.json`` (mode 600) and the
 24   access token is auto-renewed via the refresh token on every call,
 25   so multi-hour benchmark sweeps don't expire mid-run.
 262. **Developer token** (legacy / quick testing). If no OAuth tokens
 27   exist on disk we fall back to ``BOX_DEVELOPER_TOKEN`` from the
 28   environment — but those are capped at ~60 minutes by Box.
 29"""
 30
 31from __future__ import annotations
 32
 33import json
 34import logging
 35import os
 36import sys
 37import time
 38import webbrowser
 39from http.server import BaseHTTPRequestHandler, HTTPServer
 40from pathlib import Path
 41from typing import Any, Optional
 42from urllib.parse import parse_qs, urlencode, urlparse
 43
 44import requests
 45
 46log = logging.getLogger(__name__)
 47
 48USER_AGENT = "rpx-benchmark/box-upload"
 49BOX_API = "https://api.box.com/2.0"
 50BOX_AUTHORIZE_URL = "https://account.box.com/api/oauth2/authorize"
 51BOX_TOKEN_URL = "https://api.box.com/oauth2/token"
 52
 53DEFAULT_BOX_FOLDER_ID = "380510613151"  # team's RPX-Outputs folder
 54DEFAULT_REDIRECT_URI = "http://localhost:8765/callback"
 55TOKENS_PATH = Path("~/.config/rpx_benchmark/box_tokens.json").expanduser()
 56# Refresh the access token if it's within this many seconds of expiry,
 57# so long-running uploads don't see a 401 mid-call.
 58EXPIRY_SKEW_SEC = 60
 59
 60__all__ = [
 61    "DEFAULT_BOX_FOLDER_ID",
 62    "upload_tree",
 63    "upload_run_dir",
 64    "oauth_login",
 65]
 66
 67
 68def _human_bytes(n: int | float) -> str:
 69    n = float(n)
 70    for u in ("B", "KB", "MB", "GB", "TB"):
 71        if n < 1024:
 72            return f"{n:.1f} {u}"
 73        n /= 1024
 74    return f"{n:.1f} PB"
 75
 76
 77# ------------------------------------------------------------------ #
 78# OAuth 2.0 with refresh-token rotation
 79# ------------------------------------------------------------------ #
 80def _load_tokens() -> Optional[dict[str, Any]]:
 81    if not TOKENS_PATH.is_file():
 82        return None
 83    try:
 84        data: dict[str, Any] = json.loads(TOKENS_PATH.read_text())
 85        return data
 86    except (OSError, json.JSONDecodeError):
 87        return None
 88
 89
 90def _save_tokens(tokens: dict[str, Any]) -> None:
 91    TOKENS_PATH.parent.mkdir(parents=True, exist_ok=True)
 92    TOKENS_PATH.write_text(json.dumps(tokens, indent=2))
 93    try:
 94        TOKENS_PATH.chmod(0o600)
 95    except OSError:
 96        pass
 97
 98
 99def _refresh_access_token(refresh_token: str, client_id: str, client_secret: str) -> dict[str, Any]:
100    """Exchange a refresh token for a new (access_token, refresh_token) pair.
101
102    Box rotates refresh tokens on use: the response contains a NEW refresh
103    token that supersedes the old one. The old token is then invalid.
104    """
105    r = requests.post(
106        BOX_TOKEN_URL,
107        data={
108            "grant_type": "refresh_token",
109            "refresh_token": refresh_token,
110            "client_id": client_id,
111            "client_secret": client_secret,
112        },
113        timeout=30,
114    )
115    if not r.ok:
116        from .exceptions import ConfigError
117
118        raise ConfigError(
119            f"Box OAuth refresh failed ({r.status_code}): {r.text[:200]}",
120            hint="The refresh token may have expired (60-day max idle). "
121            "Re-run `python -m rpx_benchmark.box_upload login` to re-authorize.",
122        )
123    body: dict[str, Any] = r.json()
124    return body
125
126
127def _stash_tokens_from_response(resp: dict[str, Any]) -> dict[str, Any]:
128    """Normalize a Box token response into our on-disk format + persist.
129
130    We deliberately do NOT persist client_id / client_secret here — they
131    live in env vars (typically exported in the user's shell rc) so the
132    on-disk file only carries the rotating tokens.
133    """
134    tokens = {
135        "access_token": resp["access_token"],
136        "refresh_token": resp["refresh_token"],
137        "expires_at": time.time() + int(resp.get("expires_in", 3600)),
138    }
139    _save_tokens(tokens)
140    return tokens
141
142
143def _read_oauth_creds_from_env() -> tuple[Optional[str], Optional[str]]:
144    """Return (client_id, client_secret) from BOX_CLIENT_ID / BOX_CLIENT_SECRET."""
145    return os.environ.get("BOX_CLIENT_ID"), os.environ.get("BOX_CLIENT_SECRET")
146
147
148def _token() -> str:
149    """Return a valid Box access token, refreshing via OAuth if available.
150
151    Order:
152      1. OAuth tokens on disk → refresh if near expiry, return access token.
153         Uses ``BOX_CLIENT_ID`` / ``BOX_CLIENT_SECRET`` from the environment
154         to authenticate the refresh call.
155      2. ``BOX_DEVELOPER_TOKEN`` env var (legacy 60-min token).
156    """
157    tokens = _load_tokens()
158    if tokens:
159        if time.time() + EXPIRY_SKEW_SEC >= float(tokens.get("expires_at", 0)):
160            client_id, client_secret = _read_oauth_creds_from_env()
161            if not (client_id and client_secret):
162                from .exceptions import ConfigError
163
164                raise ConfigError(
165                    "OAuth access token expired and BOX_CLIENT_ID / BOX_CLIENT_SECRET not in env.",
166                    hint="Export both (typically from your shell rc) or re-run "
167                    "`python -m rpx_benchmark.box_upload login`.",
168                )
169            log.info("[box] access token expired or near expiry — refreshing")
170            resp = _refresh_access_token(tokens["refresh_token"], client_id, client_secret)
171            tokens = _stash_tokens_from_response(resp)
172        return str(tokens["access_token"])
173
174    tok = os.environ.get("BOX_DEVELOPER_TOKEN")
175    if not tok:
176        from .exceptions import ConfigError
177
178        raise ConfigError(
179            "No Box credentials available.",
180            hint="Either run `python -m rpx_benchmark.box_upload login` (OAuth, "
181            "auto-renews — recommended for long runs) or set "
182            "BOX_DEVELOPER_TOKEN to a 60-min developer token.",
183        )
184    return tok
185
186
187class _CodeCatcher(BaseHTTPRequestHandler):  # pragma: no cover
188    """Single-shot HTTP handler that captures the OAuth ?code= callback."""
189
190    code: Optional[str] = None
191    error: Optional[str] = None
192
193    def do_GET(self) -> None:  # noqa: N802 — http.server API
194        q = parse_qs(urlparse(self.path).query)
195        _CodeCatcher.code = q.get("code", [None])[0]
196        _CodeCatcher.error = q.get("error_description", q.get("error", [None]))[0]
197        self.send_response(200)
198        self.send_header("Content-Type", "text/html; charset=utf-8")
199        self.end_headers()
200        body = (
201            "<h2>Box login complete.</h2>"
202            "<p>You can close this tab and return to the terminal.</p>"
203            if _CodeCatcher.code
204            else f"<h2>Box login failed.</h2><pre>{_CodeCatcher.error}</pre>"
205        )
206        self.wfile.write(body.encode("utf-8"))
207
208    def log_message(self, format: str, *args: Any) -> None:  # noqa: A002 — http.server API
209        return  # silence stderr access log
210
211
212def oauth_login(  # pragma: no cover — interactive: spawns a browser + local HTTPServer
213    client_id: Optional[str] = None,
214    client_secret: Optional[str] = None,
215    redirect_uri: Optional[str] = None,
216) -> dict[str, Any]:
217    """Interactive one-time login. Opens a browser, captures the auth code on
218    a local HTTP server, exchanges it for tokens, writes them to disk.
219
220    Args default to ``BOX_CLIENT_ID`` / ``BOX_CLIENT_SECRET`` /
221    ``BOX_REDIRECT_URI`` from the environment. ``BOX_REDIRECT_URI`` must
222    match what's configured on the Box app; default is
223    ``http://localhost:8765/callback``.
224    """
225    client_id = client_id or os.environ.get("BOX_CLIENT_ID")
226    client_secret = client_secret or os.environ.get("BOX_CLIENT_SECRET")
227    redirect_uri = redirect_uri or os.environ.get("BOX_REDIRECT_URI") or DEFAULT_REDIRECT_URI
228    if not (client_id and client_secret):
229        from .exceptions import ConfigError
230
231        raise ConfigError(
232            "BOX_CLIENT_ID and BOX_CLIENT_SECRET required for OAuth login.",
233            hint="Set them from your Box app at https://app.box.com/developers/console "
234            "and re-run `python -m rpx_benchmark.box_upload login`.",
235        )
236
237    parsed = urlparse(redirect_uri)
238    if parsed.hostname not in ("localhost", "127.0.0.1") or not parsed.port:
239        from .exceptions import ConfigError
240
241        raise ConfigError(
242            f"redirect_uri must be http://localhost:<port>/<path>; got {redirect_uri!r}",
243            hint="Set BOX_REDIRECT_URI (and the same value in the Box app config) "
244            "to e.g. http://localhost:8765/callback.",
245        )
246
247    authorize = f"{BOX_AUTHORIZE_URL}?" + urlencode(
248        {"response_type": "code", "client_id": client_id, "redirect_uri": redirect_uri}
249    )
250    server = HTTPServer((parsed.hostname, parsed.port), _CodeCatcher)
251    sys.stderr.write(f"[box] opening browser to: {authorize}\n")
252    sys.stderr.write(f"[box] waiting for redirect on {redirect_uri} ...\n")
253    webbrowser.open(authorize)
254    server.handle_request()  # blocks until one callback received
255    server.server_close()
256    if _CodeCatcher.error or not _CodeCatcher.code:
257        from .exceptions import ConfigError
258
259        raise ConfigError(
260            f"OAuth login failed: {_CodeCatcher.error or 'no code returned'}",
261            hint="Check the Box app's redirect URI matches BOX_REDIRECT_URI exactly.",
262        )
263
264    r = requests.post(
265        BOX_TOKEN_URL,
266        data={
267            "grant_type": "authorization_code",
268            "code": _CodeCatcher.code,
269            "client_id": client_id,
270            "client_secret": client_secret,
271            "redirect_uri": redirect_uri,
272        },
273        timeout=30,
274    )
275    if not r.ok:
276        from .exceptions import ConfigError
277
278        raise ConfigError(
279            f"Box token exchange failed ({r.status_code}): {r.text[:200]}",
280            hint="Confirm client_id/client_secret are correct and OAuth 2.0 is enabled on the app.",
281        )
282    tokens = _stash_tokens_from_response(r.json())
283    sys.stderr.write(
284        f"[box] tokens saved to {TOKENS_PATH} (mode 600)\n"
285        f"[box] access expires in {int(tokens['expires_at'] - time.time())}s; "
286        "refresh token auto-renews on use\n"
287    )
288    return tokens
289
290
291def _api_get(path: str, params: Optional[dict[str, Any]] = None) -> dict[str, Any]:
292    r = requests.get(
293        f"{BOX_API}{path}",
294        headers={"Authorization": f"Bearer {_token()}", "User-Agent": USER_AGENT},
295        params=params or {},
296        timeout=60,
297    )
298    if r.status_code in (401, 403):
299        from .exceptions import DatasetError
300
301        raise DatasetError(
302            f"Box API {r.status_code} — token expired or insufficient: {r.text[:200]}",
303            hint="Token may be expired (60-min window) or lacks permissions for this folder.",
304        )
305    r.raise_for_status()
306    return r.json()  # type: ignore[no-any-return]
307
308
309def _api_post(path: str, payload: dict[str, Any]) -> dict[str, Any]:
310    r = requests.post(
311        f"{BOX_API}{path}",
312        headers={
313            "Authorization": f"Bearer {_token()}",
314            "User-Agent": USER_AGENT,
315            "Content-Type": "application/json",
316        },
317        json=payload,
318        timeout=60,
319    )
320    if r.status_code in (401, 403):
321        from .exceptions import DatasetError
322
323        raise DatasetError(
324            f"Box API {r.status_code} — token expired or insufficient: {r.text[:200]}",
325            hint="Token may be expired (60-min window) or lacks permissions for this folder.",
326        )
327    r.raise_for_status()
328    return r.json()  # type: ignore[no-any-return]
329
330
331def _api_list_folder(folder_id: str) -> list[dict[str, Any]]:
332    items: list[dict[str, Any]] = []
333    offset = 0
334    while True:
335        data = _api_get(
336            f"/folders/{folder_id}/items",
337            params={"fields": "id,type,name,size", "limit": 1000, "offset": offset},
338        )
339        entries = data.get("entries", []) or []
340        for it in entries:
341            if it.get("type") not in ("file", "folder"):
342                continue
343            items.append(
344                {
345                    "id": str(it["id"]),
346                    "type": it["type"],
347                    "name": it.get("name") or str(it["id"]),
348                    "size": int(it.get("size") or 0),
349                }
350            )
351        if len(entries) < 1000:
352            break
353        offset += 1000
354    return items
355
356
357def _api_ensure_folder(name: str, parent_folder_id: str) -> str:
358    """Return folder id for ``name`` under ``parent_folder_id``; create on miss."""
359    try:
360        data = _api_post("/folders", {"name": name, "parent": {"id": str(parent_folder_id)}})
361        return str(data["id"])
362    except requests.HTTPError as e:
363        if e.response is None or e.response.status_code != 409:
364            raise
365        for it in _api_list_folder(parent_folder_id):
366            if it["type"] == "folder" and it["name"] == name:
367                return str(it["id"])
368        from .exceptions import DatasetError
369
370        raise DatasetError(
371            f"folder {name!r} reported conflict but not found under {parent_folder_id}",
372            hint="Race condition or stale Box state. Re-run; the conflict resolver "
373            "should pick up the existing folder on retry.",
374        ) from e
375
376
377def _api_resolve_path(parts: list[str], root_folder_id: str) -> str:
378    cur = str(root_folder_id)
379    for part in parts:
380        if not part:
381            continue
382        cur = _api_ensure_folder(part, cur)
383    return cur
384
385
386def _api_existing_files(folder_id: str) -> dict[str, dict[str, Any]]:
387    return {it["name"]: it for it in _api_list_folder(folder_id) if it["type"] == "file"}
388
389
390def _api_upload_file(
391    local: Path,
392    parent_folder_id: str,
393    name: Optional[str] = None,
394    existing: Optional[dict[str, dict[str, Any]]] = None,
395) -> dict[str, Any]:
396    """Upload one file. If ``existing[name]`` matches local size, skip. If exists
397    with different size, upload as a new version on the same file id."""
398    name = name or local.name
399    size = local.stat().st_size
400    existing = existing if existing is not None else _api_existing_files(parent_folder_id)
401    if name in existing and int(existing[name].get("size") or -1) == size:
402        return {"id": existing[name]["id"], "name": name, "skipped": True}
403
404    headers = {"Authorization": f"Bearer {_token()}", "User-Agent": USER_AGENT}
405    attrs: dict[str, Any]
406    if name in existing:
407        url = f"https://upload.box.com/api/2.0/files/{existing[name]['id']}/content"
408        attrs = {"name": name}
409    else:
410        url = "https://upload.box.com/api/2.0/files/content"
411        attrs = {"name": name, "parent": {"id": str(parent_folder_id)}}
412    with local.open("rb") as f:
413        r = requests.post(
414            url,
415            headers=headers,
416            timeout=600,
417            data={"attributes": json.dumps(attrs)},
418            files={"file": (name, f)},
419        )
420    if r.status_code in (401, 403):
421        from .exceptions import DatasetError
422
423        raise DatasetError(
424            f"Box upload {r.status_code}: {r.text[:200]}",
425            hint="Check the developer token is fresh and has write access to the target folder.",
426        )
427    r.raise_for_status()
428    entry = r.json().get("entries", [{}])[0]
429    return {"id": entry.get("id"), "name": name, "skipped": False}
430
431
432def upload_tree(
433    local_dir: Path | str,
434    remote_path: str,
435    *,
436    root_folder_id: str = "0",
437    verbose: bool = True,
438) -> dict[str, Any]:
439    """Upload every file under ``local_dir`` to ``<root>/<remote_path>``.
440
441    Mirrors directory structure; size-matched files are skipped (idempotent).
442
443    Parameters
444    ----------
445    local_dir
446        Local directory whose contents should be mirrored.
447    remote_path
448        Slash-separated path *under* ``root_folder_id``. Intermediate
449        folders are auto-created.
450    root_folder_id
451        Box folder id (not URL) to root the upload under. Default ``"0"`` =
452        the auth token's "My Files".
453    verbose
454        If True, log one line per file to stderr.
455
456    Returns
457    -------
458    dict
459        ``{"uploaded": int, "skipped": int, "bytes_uploaded": int,
460        "remote_folder_id": str}``.
461    """
462    local_dir = Path(local_dir)
463    if not local_dir.is_dir():
464        raise NotADirectoryError(local_dir)
465    base_id = _api_resolve_path(remote_path.strip("/").split("/"), root_folder_id)
466    n_up = n_skip = bytes_up = 0
467    by_parent: dict[Path, list[Path]] = {}
468    for p in sorted(local_dir.rglob("*")):
469        if p.is_file():
470            by_parent.setdefault(p.parent, []).append(p)
471    for parent_dir, files in by_parent.items():
472        rel = parent_dir.relative_to(local_dir)
473        parent_id = _api_resolve_path(str(rel).split("/") if str(rel) != "." else [], base_id)
474        existing = _api_existing_files(parent_id)
475        for f in files:
476            if verbose:
477                sys.stderr.write(f"[box] {f.relative_to(local_dir)} ... ")
478                sys.stderr.flush()
479            r = _api_upload_file(f, parent_id, existing=existing)
480            if r["skipped"]:
481                n_skip += 1
482                if verbose:
483                    sys.stderr.write("skip (already on Box)\n")
484            else:
485                n_up += 1
486                bytes_up += f.stat().st_size
487                if verbose:
488                    sys.stderr.write(f"uploaded ({_human_bytes(f.stat().st_size)})\n")
489    return {
490        "uploaded": n_up,
491        "skipped": n_skip,
492        "bytes_uploaded": bytes_up,
493        "remote_folder_id": base_id,
494    }
495
496
497def upload_run_dir(
498    out_dir: Path | str,
499    *,
500    task: str,
501    model_name: str,
502    split: str,
503    root_folder_id: str = DEFAULT_BOX_FOLDER_ID,
504    verbose: bool = True,
505) -> dict[str, Any]:
506    """Mirror a single per-run output directory to Box.
507
508    Lays everything under ``out_dir`` at the canonical path
509    ``<root_folder_id>/<task>/<safe_model_name>/<split>/`` so the
510    layout matches every other run (depth, pose, future tasks, post-hoc
511    sync). ``safe_model_name`` swaps ``/`` for ``__`` so HF-style
512    namespaces don't accidentally create nested Box folders.
513
514    Returns the same dict as :func:`upload_tree`, plus the resolved
515    ``remote_path`` for log lines.
516    """
517    safe_name = model_name.replace("/", "__")
518    remote = f"{task}/{safe_name}/{split}"
519    if verbose:
520        sys.stderr.write(f"[box] mirror {Path(out_dir)} → box:{remote}\n")
521    summary = upload_tree(
522        out_dir, remote_path=remote, root_folder_id=root_folder_id, verbose=verbose
523    )
524    summary["remote_path"] = remote
525    return summary
526
527
528def _cli() -> int:  # pragma: no cover — CLI dispatcher; logic lives in the helpers above
529    import argparse
530
531    ap = argparse.ArgumentParser(prog="python -m rpx_benchmark.box_upload")
532    sub = ap.add_subparsers(dest="cmd", required=True)
533
534    p_login = sub.add_parser(
535        "login",
536        help="One-time OAuth login. Captures a refresh token for unattended uploads.",
537    )
538    p_login.add_argument("--client-id", default=None, help="default: $BOX_CLIENT_ID")
539    p_login.add_argument("--client-secret", default=None, help="default: $BOX_CLIENT_SECRET")
540    p_login.add_argument(
541        "--redirect-uri",
542        default=None,
543        help=f"default: $BOX_REDIRECT_URI or {DEFAULT_REDIRECT_URI}",
544    )
545
546    sub.add_parser("whoami", help="Print the Box user the current token authenticates as.")
547    sub.add_parser("logout", help="Delete cached OAuth tokens from disk.")
548
549    args = ap.parse_args()
550    if args.cmd == "login":
551        oauth_login(args.client_id, args.client_secret, args.redirect_uri)
552        return 0
553    if args.cmd == "whoami":
554        r = requests.get(
555            f"{BOX_API}/users/me",
556            headers={"Authorization": f"Bearer {_token()}", "User-Agent": USER_AGENT},
557            timeout=30,
558        )
559        if not r.ok:
560            sys.stderr.write(f"[box] {r.status_code}: {r.text[:200]}\n")
561            return 1
562        d = r.json()
563        print(f"{d.get('name')} <{d.get('login')}>  id={d.get('id')}")
564        return 0
565    if args.cmd == "logout":
566        if TOKENS_PATH.exists():
567            TOKENS_PATH.unlink()
568            print(f"removed {TOKENS_PATH}")
569        else:
570            print("no tokens on disk")
571        return 0
572    return 1
573
574
575if __name__ == "__main__":
576    raise SystemExit(_cli())