app.py

"""
Webhook receiver: turns a push to psyche-private into a build Job.

WHY A SPACE IN FRONT OF A JOB
-----------------------------
A webhook can trigger a Job directly, and the Hub does re-run the Job — but it
clones only PART of the spec. Verified empirically against a throwaway repo:

    field       direct Job          webhook-triggered
    namespace   psyche-org          the token owner's PERSONAL namespace
    secrets     present             DROPPED
    volumes     present             DROPPED
    timeout     present             DROPPED

A webhook-triggered Job therefore starts with no HF_TOKEN, cannot read the
private builder repo, and cannot write the output. That is not fixable by
configuration.

A Space, on the other hand, keeps its own secrets. So the Space is a thin
trigger that receives the webhook and creates the Job *properly* — with the
token, the flavor, the timeout and the right namespace. The Job does all the
work; this process only starts it.

SPACE SETUP
-----------
SDK gradio. Secrets:
    HF_TOKEN        write access to psyche-org
    WEBHOOK_SECRET  any random string; the same value goes in the webhook config
`config.toml` is uploaded next to this file and is the single source of truth
for the repo ids and the Job's hardware.
"""

import os
import threading
import tomllib

from huggingface_hub import HfApi, WebhookPayload, WebhooksServer

with open("config.toml", "rb") as file:
    CONFIG = tomllib.load(file)

SOURCE_REPO = CONFIG["source"]["repo_id"]
BUILDER_REPO = CONFIG["builder"]["repo_id"]
JOB = CONFIG["job"]

TOKEN = os.environ["HF_TOKEN"]
api = HfApi(token=TOKEN)

# The Job's script: fetch the builder code, then build both artefacts. Kept
# identical to `buildDB/deploy.sh --print-job-command update`.
JOB_SCRIPT = f"""set -eu
export UV_PROJECT_ENVIRONMENT=/tmp/psyche-venv
uv run --no-project --with huggingface_hub python -c \
"from huggingface_hub import snapshot_download; \
snapshot_download('{BUILDER_REPO}', repo_type='dataset', local_dir='/tmp/builder')"
cd /tmp/builder
uv run --frozen buildDB/build.py update
uv run --frozen buildPaperIndex/build.py --sync-hf --upload
"""

ACTIVE_STAGES = ("RUNNING", "SCHEDULING")

_wake = threading.Event()
_last = {"job": None, "note": "nothing yet"}


def _running_job():
    """The build Job currently in flight, if any."""
    for job in api.list_jobs(
        namespace=JOB["namespace"], labels={"name": JOB["name"]}
    ):
        if job.status and job.status.stage in ACTIVE_STAGES:
            return job

    return None


def _launch():
    info = api.run_job(
        image=JOB["image"],
        command=["bash", "-c", JOB_SCRIPT],
        env={
            "PSYCHE_SOURCE_DIR": "/tmp/psyche-build/psyche-private",
            "PSYCHE_OUTPUT_DIR": "/tmp/psyche-build/psyche",
            "PYTHONUNBUFFERED": "1",
        },
        secrets={"HF_TOKEN": TOKEN},
        flavor=JOB["flavor"],
        timeout=JOB["timeout"],
        labels={"name": JOB["name"], "mode": "update", "trigger": "webhook"},
        namespace=JOB["namespace"],
    )

    return info


def _worker() -> None:
    """
    One build at a time, bursts coalesced.

    Several commits in a row wake this once more, not N times — and `update`
    is idempotent, so a run that finds nothing to do exits in seconds.
    """
    while True:
        _wake.wait()
        _wake.clear()

        try:
            existing = _running_job()

            if existing is not None:
                # Its own revision check will pick up whatever arrived while
                # it was starting, so a second Job would only duplicate work.
                _last["note"] = f"already running: {existing.id}"
                print(f"[worker] {_last['note']}", flush=True)
                continue

            info = _launch()
            _last["job"] = info.id
            _last["note"] = f"launched {info.id}"
            print(f"[worker] launched job {info.id}", flush=True)

        except Exception as error:  # noqa: BLE001 — must not kill the worker
            _last["note"] = f"failed to launch: {error}"
            print(f"[worker] {_last['note']}", flush=True)


threading.Thread(target=_worker, daemon=True).start()


app = WebhooksServer()


@app.add_webhook("/psyche_update")
async def on_source_change(payload: WebhookPayload):
    repo = getattr(payload.repo, "name", None)

    # Log every delivery, including the ones we ignore: silence here is
    # otherwise indistinguishable from the webhook never arriving.
    print(f"[webhook] received event for {repo}", flush=True)

    if repo != SOURCE_REPO:
        return {"skipped": f"not {SOURCE_REPO}: {repo}"}

    scope = getattr(payload.event, "scope", "")

    if not str(scope).startswith("repo"):
        return {"skipped": f"scope {scope}"}

    _wake.set()
    print(f"[webhook] queued a build for {repo} ({scope})", flush=True)

    return {"queued": True, "repo": repo, "scope": str(scope)}


if __name__ == "__main__":
    print(
        f"watching {SOURCE_REPO} -> job {JOB['name']} in {JOB['namespace']}",
        flush=True,
    )
    # ssr_mode=False is required. Gradio 6 otherwise serves through a Node
    # SSR proxy that forwards only the UI routes to Python, so the FastAPI
    # webhook route is announced at startup but answers 405 to a real POST.
    app.launch(ssr_mode=False)