Spaces:
Sleeping
Sleeping
malteos
Add per-step form validation; allow language-only filter; drop "mockup" wording
fdce60b unverified | """FastAPI backend for the Common Crawl WARC Repackager wizard UI. | |
| This is a thin HTTP + SSE shell over the existing, executor-agnostic business logic in | |
| ``src/`` — it does NOT reimplement any job logic: | |
| * ``src.job_runner.make_runner`` picks the real HF-Jobs runner or the local podman | |
| runner from ``CC_REPACKAGE_EXECUTOR`` (same as the Gradio app). | |
| * ``src.jobs`` builds the job specs / bucket URIs, ``src.buckets`` validates input and | |
| lists crawls/buckets, ``src.estimate`` parses the ``ESTIMATE`` line and does the cost | |
| math, ``src.examples`` provides the precomputed demo presets + simulated log lines. | |
| Endpoints: static wizard at ``/``; JSON at ``/api/*``; per-job log streaming over SSE at | |
| ``/api/jobs/{id}/events`` and ``/api/examples/{key}/events``. OAuth (for the deployed | |
| Space) is wired only when the Space injects ``OAUTH_CLIENT_ID`` — locally the podman | |
| executor runs as a guest, exactly like the Gradio app's ``IS_PODMAN`` path. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import time | |
| from pathlib import Path | |
| from typing import Dict, Optional | |
| from fastapi import FastAPI, Request | |
| from fastapi.responses import FileResponse, JSONResponse, StreamingResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from src import buckets, config, examples | |
| from src.estimate import Estimate, actual_cost_usd, parse_estimate | |
| from src.job_runner import TERMINAL_STAGES, make_runner | |
| from src.jobs import RepackageRequest, bucket_file_url, bucket_uris, bucket_web_url | |
| STATIC_DIR = Path(__file__).parent / "static" | |
| IS_PODMAN = config.EXECUTOR == "podman" | |
| POLL_SECONDS = 2 if IS_PODMAN else 5 | |
| EXAMPLE_LINE_DELAY = 0.35 | |
| RUNNER = make_runner() | |
| # job_id -> {kind, req, token}; lets the SSE stream build the right terminal payload. | |
| JOBS: Dict[str, dict] = {} | |
| app = FastAPI(title="CC WARC Repackager") | |
| # --- auth (guest locally; real OAuth on the Space) --------------------------- | |
| def _session(request: Request): | |
| # request.session only exists when SessionMiddleware is installed (OAuth wired). | |
| # `hasattr` is unsafe here — the property raises AssertionError, not AttributeError. | |
| return request.scope.get("session") | |
| def token_for(request: Request) -> Optional[str]: | |
| if IS_PODMAN: | |
| return None | |
| sess = _session(request) | |
| return sess.get("token") if sess else None | |
| def user_for(request: Request) -> Optional[str]: | |
| if IS_PODMAN: | |
| return None | |
| sess = _session(request) | |
| return sess.get("user") if sess else None | |
| # --- serialization helpers --------------------------------------------------- | |
| def _estimate_json(est: Estimate) -> dict: | |
| return { | |
| "n_records": est.n_records, | |
| "total_bytes": est.total_bytes, | |
| "total_gib": round(est.total_gib, 4), | |
| "n_output_files": est.n_output_files, | |
| "rec_per_s": est.rec_per_s, | |
| "hourly_usd": est.hourly_usd, | |
| "duration_seconds": round(est.duration_seconds, 1), | |
| "cost_usd": round(est.cost_usd, 4), | |
| "flavor": est.flavor, | |
| } | |
| def _example_json(ex: examples.Example) -> dict: | |
| return { | |
| "key": ex.key, | |
| "title": ex.title, | |
| "description": ex.description, | |
| "kind": "sql" if ex.sql_where else "form", | |
| "n_records": ex.n_records, | |
| "total_bytes": ex.total_bytes, | |
| "download_url": bucket_file_url(ex.bucket, ex.path, ex.files[0]) if ex.files else "", | |
| "folder_url": bucket_web_url(ex.bucket, ex.path), | |
| "file": ex.files[0] if ex.files else "", | |
| "bucket": ex.bucket, | |
| "path": ex.path, | |
| "form": { | |
| "domains": ", ".join(ex.domains), | |
| "hostnames": ", ".join(ex.hostnames), | |
| "languages": list(ex.languages), | |
| "sql_where": ex.sql_where, | |
| "max_records": ex.max_records, | |
| "crawls": list(ex.crawls), | |
| "name": ex.name, | |
| }, | |
| } | |
| def _sse(payload: dict) -> str: | |
| return f"data: {json.dumps(payload)}\n\n" | |
| def clean_log(text: str) -> str: | |
| """Trim the local podman runner's echoed job script from the log for readability. | |
| The local runner writes the full `podman run ... <heredoc>` command (which inlines | |
| index_query.py) ahead of the container's own stdout. Real HF Jobs logs contain only | |
| the container stdout, so this is a no-op there (no marker / no `$ podman` line).""" | |
| marker = "CC_PY_EOF" | |
| if marker in text: # estimate job: keep output after the heredoc | |
| text = text.rsplit(marker, 1)[-1] | |
| lines = [ln for ln in text.splitlines() if not ln.startswith("$ podman run")] | |
| return "\n".join(lines).strip() | |
| # --- request building (mirrors the Gradio app's _build_request) -------------- | |
| def build_request(p: dict) -> tuple[Optional[RepackageRequest], Optional[str], str]: | |
| """Return (request, error, warning). Validates via the shared src.buckets helpers.""" | |
| domains = buckets.parse_targets(p.get("domains", "")) | |
| hostnames = buckets.parse_targets(p.get("hostnames", "")) | |
| languages = list(p.get("languages") or []) | |
| sql_where = (p.get("sql_where") or "").strip() | |
| warnings = [] | |
| ok, msg, max_records = buckets.validate_max_records(p.get("max_records")) | |
| if not ok: | |
| return None, msg, "" | |
| if sql_where: | |
| ok, msg = buckets.validate_sql_where(sql_where) | |
| if not ok: | |
| return None, msg, "" | |
| if hostnames or domains or languages: | |
| warnings.append("Using the advanced SQL filter — the domain / hostname / language fields are ignored.") | |
| hostnames, domains, languages = [], [], [] | |
| else: | |
| ok, msg = buckets.validate_languages(languages) | |
| if not ok: | |
| return None, msg, "" | |
| ok, msg = buckets.validate_targets(hostnames, domains, languages) | |
| if not ok: | |
| return None, msg, "" | |
| crawls = list(p.get("crawls") or []) | |
| ok, msg = buckets.validate_crawls(crawls) | |
| if not ok: | |
| return None, msg, "" | |
| if msg: | |
| warnings.append(msg) | |
| req = RepackageRequest( | |
| hostnames=hostnames, domains=domains, languages=languages, sql_where=sql_where, | |
| max_records=max_records, crawls=crawls, | |
| target_bucket=(p.get("target_bucket") or "").strip(), | |
| target_path=(p.get("target_path") or "").strip(), | |
| name=(p.get("name") or "").strip() or "repackaged", | |
| flavor=p.get("flavor") or config.DEFAULT_FLAVOR, | |
| creator=p.get("creator") or "", | |
| ) | |
| return req, None, "\n".join(warnings) | |
| # --- API --------------------------------------------------------------------- | |
| def api_config(request: Request): | |
| return { | |
| "is_podman": IS_PODMAN, | |
| "signed_in": bool(token_for(request)) or IS_PODMAN, | |
| "user": user_for(request), | |
| "default_flavor": config.DEFAULT_FLAVOR, | |
| "flavors": { | |
| f: {"usd": config.FLAVOR_HOURLY_USD[f], "rps": config.FLAVOR_REC_PER_S.get(f, 0), | |
| "vcpus": config.FLAVOR_VCPUS.get(f, 0)} | |
| for f in config.FLAVOR_HOURLY_USD | |
| }, | |
| "crawls": list(config.FALLBACK_CRAWLS), | |
| "languages": [[n, c] for n, c in config.CONTENT_LANGUAGES], | |
| "cost_guard_max_crawls": config.COST_GUARD_MAX_CRAWLS, | |
| "examples": [_example_json(ex) for ex in examples.EXAMPLES], | |
| } | |
| def api_crawls(request: Request): | |
| if IS_PODMAN: | |
| return {"crawls": list(config.FALLBACK_CRAWLS)} | |
| return {"crawls": buckets.list_available_crawls(token_for(request))} | |
| def api_buckets(request: Request): | |
| if IS_PODMAN: | |
| return {"buckets": ["local/out"]} | |
| return {"buckets": buckets.list_writable_buckets(token_for(request))} | |
| async def api_validate(request: Request): | |
| req, err, warning = build_request(await request.json()) | |
| if err: | |
| return {"ok": False, "error": err} | |
| return {"ok": True, "warning": warning} | |
| async def api_estimate(request: Request): | |
| p = await request.json() | |
| token = token_for(request) | |
| if not IS_PODMAN and token is None: | |
| return JSONResponse({"error": "Please sign in with Hugging Face first."}, status_code=401) | |
| req, err, warning = build_request(p) | |
| if err: | |
| return JSONResponse({"error": err}, status_code=400) | |
| if not IS_PODMAN: | |
| ok, vmsg = buckets.validate_target_bucket(req.target_bucket, req.target_path, token) | |
| if not ok: | |
| return JSONResponse({"error": vmsg}, status_code=400) | |
| try: | |
| job_id = RUNNER.estimate(req, token) | |
| except Exception as e: # noqa: BLE001 | |
| return JSONResponse({"error": f"{type(e).__name__}: {e}"}, status_code=500) | |
| JOBS[job_id] = {"kind": "estimate", "req": req, "token": token} | |
| return {"job_id": job_id, "warning": warning} | |
| async def api_launch(request: Request): | |
| p = await request.json() | |
| token = token_for(request) | |
| if not IS_PODMAN and token is None: | |
| return JSONResponse({"error": "Please sign in with Hugging Face first."}, status_code=401) | |
| req, err, warning = build_request(p) | |
| if err: | |
| return JSONResponse({"error": err}, status_code=400) | |
| try: | |
| job_id = RUNNER.fetch(req, token, n_records=int(p.get("n_records") or 0)) | |
| except Exception as e: # noqa: BLE001 | |
| return JSONResponse({"error": f"{type(e).__name__}: {e}"}, status_code=500) | |
| JOBS[job_id] = {"kind": "fetch", "req": req, "token": token} | |
| return {"job_id": job_id, "warning": warning} | |
| def api_cancel(job_id: str, request: Request): | |
| info = JOBS.get(job_id) | |
| token = info["token"] if info else token_for(request) | |
| try: | |
| RUNNER.cancel(job_id, token) | |
| return {"ok": True} | |
| except Exception as e: # noqa: BLE001 | |
| return JSONResponse({"ok": False, "error": f"{type(e).__name__}: {e}"}, status_code=500) | |
| def api_job_events(job_id: str): | |
| info = JOBS.get(job_id) | |
| if not info: | |
| return JSONResponse({"error": "unknown job"}, status_code=404) | |
| req, token, kind = info["req"], info["token"], info["kind"] | |
| def gen(): | |
| url = RUNNER.job_url(job_id) | |
| if url: | |
| yield _sse({"type": "url", "url": url}) | |
| started = time.monotonic() | |
| final_logs = [] | |
| while True: | |
| stage = RUNNER.status(job_id, token) | |
| logs = RUNNER.logs(job_id, token) | |
| final_logs = logs | |
| yield _sse({"type": "tick", "stage": stage, "log": clean_log("\n".join(logs[-600:]))}) | |
| if stage in TERMINAL_STAGES or time.monotonic() - started > 20 * 60: | |
| break | |
| time.sleep(POLL_SECONDS) | |
| report = RUNNER.report(job_id, token) | |
| done = {"type": "done", "stage": stage} | |
| if report: | |
| secs, flavor = report | |
| done["report"] = {"seconds": secs, "flavor": flavor, | |
| "cost_usd": round(actual_cost_usd(secs, flavor), 4)} | |
| if kind == "estimate": | |
| est = parse_estimate(final_logs, req.flavor) | |
| if est is not None: | |
| done["estimate"] = _estimate_json(est) | |
| else: | |
| cleaned = clean_log("\n".join(final_logs)).splitlines() | |
| err = next((ln for ln in reversed(cleaned) if "ERROR" in ln or "Error" in ln), "") | |
| done["error"] = err.strip() or "No ESTIMATE summary was produced." | |
| else: | |
| ranges_uri, output_uri = bucket_uris(req) | |
| done["result"] = { | |
| "ranges_uri": ranges_uri, "output_uri": output_uri, | |
| "bucket_url": bucket_web_url(req.target_bucket, req.normalized_path()), | |
| "completed": stage == "COMPLETED", | |
| } | |
| yield _sse(done) | |
| return StreamingResponse(gen(), media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) | |
| def api_example_events(key: str, phase: Optional[str] = None): | |
| """Stream a precomputed example's simulated logs. `phase` selects a single step of | |
| the tutorial: "estimate" (index logs + cost) or "fetch" (repackage logs + download). | |
| Omitting it streams both back to back (legacy).""" | |
| ex = examples.EXAMPLES_BY_KEY.get(key) | |
| if not ex: | |
| return JSONResponse({"error": "unknown example"}, status_code=404) | |
| def gen(): | |
| if phase in (None, "estimate"): | |
| for line in examples.estimate_log_lines(ex): | |
| time.sleep(EXAMPLE_LINE_DELAY) | |
| yield _sse({"type": "log", "line": line}) | |
| yield _sse({"type": "estimate", "estimate": _estimate_json( | |
| Estimate(ex.n_records, ex.total_bytes, config.DEFAULT_FLAVOR))}) | |
| if phase == "estimate": | |
| yield _sse({"type": "done"}) | |
| return | |
| time.sleep(0.3) | |
| if phase in (None, "fetch"): | |
| for line in examples.fetch_log_lines(ex): | |
| time.sleep(EXAMPLE_LINE_DELAY) | |
| yield _sse({"type": "log", "line": line}) | |
| yield _sse({"type": "done", "result": { | |
| "demo": True, "download_url": bucket_file_url(ex.bucket, ex.path, ex.files[0]), | |
| "folder_url": bucket_web_url(ex.bucket, ex.path), "file": ex.files[0], | |
| "n_records": ex.n_records, "total_bytes": ex.total_bytes}}) | |
| return StreamingResponse(gen(), media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) | |
| # --- OAuth (only active on the Space, when HF injects the client id) ---------- | |
| def _redirect_uri(request: Request) -> str: | |
| """The OAuth callback URL. On a Space, HF registers https://<SPACE_HOST>/auth/callback, | |
| and requests arrive behind a proxy, so prefer SPACE_HOST over the request's own host.""" | |
| host = os.environ.get("SPACE_HOST") | |
| if host: | |
| return f"https://{host}/auth/callback" | |
| return str(request.url_for("auth")) | |
| def _wire_oauth(): | |
| client_id = os.environ.get("OAUTH_CLIENT_ID") | |
| if not client_id: | |
| return # local / podman: guest mode, no OAuth | |
| try: | |
| from authlib.integrations.starlette_client import OAuth | |
| from starlette.middleware.sessions import SessionMiddleware | |
| from starlette.responses import RedirectResponse | |
| except Exception: # noqa: BLE001 | |
| return | |
| app.add_middleware( | |
| SessionMiddleware, | |
| secret_key=os.environ.get("SESSION_SECRET") or os.environ.get("OAUTH_CLIENT_SECRET") or os.urandom(24).hex(), | |
| same_site="lax", https_only=True, | |
| ) | |
| # HF login needs openid+profile for userinfo, plus the app scopes (write-repos, jobs). | |
| scopes = os.environ.get("OAUTH_SCOPES", "write-repos jobs") | |
| for s in ("profile", "openid"): | |
| if s not in scopes.split(): | |
| scopes = s + " " + scopes | |
| oauth = OAuth() | |
| oauth.register( | |
| name="hf", | |
| client_id=client_id, | |
| client_secret=os.environ.get("OAUTH_CLIENT_SECRET"), | |
| server_metadata_url=os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co") + "/.well-known/openid-configuration", | |
| client_kwargs={"scope": scopes}, | |
| ) | |
| async def login(request: Request): | |
| return await oauth.hf.authorize_redirect(request, _redirect_uri(request)) | |
| async def auth(request: Request): | |
| try: | |
| tok = await oauth.hf.authorize_access_token(request) | |
| except Exception as e: # noqa: BLE001 | |
| return JSONResponse({"error": f"OAuth failed: {type(e).__name__}: {e}"}, status_code=400) | |
| request.session["token"] = tok.get("access_token") | |
| ui = tok.get("userinfo") or {} | |
| request.session["user"] = ui.get("preferred_username") or ui.get("name") | |
| return RedirectResponse(url="/") | |
| async def logout(request: Request): | |
| request.session.clear() | |
| return RedirectResponse(url="/") | |
| _wire_oauth() | |
| # --- static (mounted at /static; "/" serves the app shell) ------------------- | |
| def index(): | |
| return FileResponse(STATIC_DIR / "index.html") | |
| app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") | |