app.py

"""Zero-GPU Hugging Face Space — Gradio entry point.

A single UI that trains **SFT / DPO / KTO**, with the number of optimizer steps
derived dynamically from the uploaded trace (see below). The training call is
decorated with ``@spaces.GPU`` so a ZeroGPU Space attaches an A10G on demand
(and trains on GPU); on a CPU / free-CPU Space that decorator is a no-op and the
same code trains on CPU.

Queue
-----
Requests are **queued**, not run inline, so a client timeout cannot lose them.
"Submit" caches the request in an in-RAM queue and returns immediately (no GPU),
so the request is safe even if the caller times out. "Train" drains the queue
under ``@spaces.GPU`` — running **all SFT jobs, then all DPO, then all KTO**,
FIFO within a method — and reports each job's result. Set ``QUEUE_PERSIST=true``
to also mirror the queue to the bucket so it survives a Space restart.

Dynamic gradient accumulation
-----------------------------
The number of optimizer steps a job runs is derived from how many examples it
contains and the deployment's ``GRAD_ACCUM``:  ``steps = ceil(n / G)``.  A
single trace -> 1 step; 10 traces with ``GRAD_ACCUM=4`` -> 3 steps (4 + 4 + 2).
The caller just uploads a trace — it never picks "steps per call".

One bucket + one model per deployment
--------------------------------------
``MODEL_ID`` / ``BUCKET_ID`` (and the LoRA hyper-params) are set once via the
Space's environment variables, not re-entered per request. The base model is
loaded once and kept **resident in RAM/VRAM** across calls (auto-sharded across
all GPUs when more than one is available) — much more efficient than re-loading
per request. Only weights + optimizer/scheduler state are checkpointed to the
bucket; the uploaded trace (training data) is never persisted.

Set as Space secrets/variables: ``HF_TOKEN``, ``MODEL_ID``, ``BUCKET_ID``
(+ optionally ``LORA_*``, ``GRAD_ACCUM``, ``LEARNING_RATE``, ``MAX_LENGTH``,
``QUEUE_PERSIST``).
"""

import json
import os
import re
import threading

import gradio as gr
import spaces  # noqa: F401  (provides spaces.GPU; auto-wraps Gradio fns too)

# NOTE: the module is named ``job_queue`` (not ``queue``) on purpose — a local
# ``queue.py`` would shadow Python's stdlib ``queue`` module and break
# huggingface_hub's ``from queue import Queue``.
from job_queue import get_queue
from train import (
    TrainConfig,
    current_progress,
    resolve_bucket_id,
    resolve_learning_rate,
    resolve_lora_alpha,
    resolve_lora_dropout,
    resolve_lora_r,
    resolve_max_length,
    resolve_model_id,
    resolve_target_modules,
    run_one_call,
)

METHODS = ["sft", "dpo", "kto"]

EXAMPLES = {
    "sft": [
        {"messages": [
            {"role": "user", "content": "What is 2+2?"},
            {"role": "assistant", "content": "4"},
        ]},
        {"messages": [
            {"role": "user", "content": "Capital of France?"},
            {"role": "assistant", "content": "Paris"},
        ]},
    ],
    "dpo": [
        {
            "prompt": "Translate 'hello' to French.",
            "chosen": "Bonjour",
            "rejected": "Hello",
        },
        {
            "prompt": "What is 2+2?",
            "chosen": "4",
            "rejected": "3",
        },
    ],
    "kto": [
        {"prompt": "Be kind. ", "completion": "Here is a helpful answer.", "desirable": True},
        {"prompt": "Be kind. ", "completion": "I refuse to help.", "desirable": False},
    ],
}

# Example data is shown as JSONL (one JSON object per line) for the upload box.
EXAMPLE_JSONL = {
    "sft": json.dumps(EXAMPLES["sft"], indent=2),
    "dpo": json.dumps(EXAMPLES["dpo"], indent=2),
    "kto": json.dumps(EXAMPLES["kto"], indent=2),
}

# The queue is configured once at startup from the deployment's env.
_q = get_queue()
# A lock so the (fast) submit callback and the (slow, GPU) drain callback —
# which Gradio runs in separate threads — don't touch the queue concurrently.
_QUEUE_LOCK = threading.Lock()


def _parse_upload(text: str) -> list[dict]:
    """Accept either a JSON array or JSONL (one object per line)."""
    text = (text or "").strip()
    if not text:
        return []
    # try strict JSON first (array or single object)
    try:
        data = json.loads(text)
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            return [data]
    except json.JSONDecodeError:
        pass
    # JSONL / loose JSON: tolerate trailing commas and split lines
    rows = []
    for line in text.splitlines():
        line = line.strip().rstrip(",")
        if not line or line.startswith("{") is False:
            continue
        try:
            rows.append(json.loads(line))
        except json.JSONDecodeError:
            try:  # tolerate trailing commas
                rows.append(json.loads(re.sub(r",\s*([}\]])", r"\1", line)))
            except json.JSONDecodeError:
                continue
    return rows


def build_app() -> gr.Blocks:
    with gr.Blocks(title="Zero-GPU Training Space") as demo:
        gr.Markdown(
            "# 🧠 Zero-GPU Training Space\n"
            "Train **SFT / DPO / KTO** — the number of optimizer steps is "
            "**derived from the trace** (dynamic gradient accumulation), not "
            "chosen per request. Requests are **queued** so a client timeout can "
            "never lose them: *Submit* caches them, *Train* runs them in "
            "**method priority** (SFT → DPO → KTO) on **GPU if the ZeroGPU "
            "decorator attaches one, otherwise CPU**. Weights + optimizer state "
            "are checkpointed to a Hugging Face Storage Bucket; the trace "
            "itself is **never stored**, so you can shut the Space down and "
            "resume later without restarting the optimizer from scratch."
        )

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 1 · Configuration")
                gr.Markdown(
                    "#### Deployment (set once via env, not per request)\n"
                    f"- **Model:** `{resolve_model_id(TrainConfig())}`  "
                    f"(env `MODEL_ID`)  \n"
                    f"- **Bucket:** `{resolve_bucket_id(TrainConfig()) or '(none — set BUCKET_ID)'} `  "
                    f"(env `BUCKET_ID`)  \n"
                    f"- **LoRA:** rank `{resolve_lora_r(TrainConfig())}`, "
                    f"alpha `{resolve_lora_alpha(TrainConfig())}`, "
                    f"dropout `{resolve_lora_dropout(TrainConfig())}`, "
                    f"targets `{', '.join(resolve_target_modules(TrainConfig()))}`  "
                    f"(env `LORA_R` / `LORA_ALPHA` / `LORA_DROPOUT` / "
                    f"`LORA_TARGET_MODULES`)  \n"
                    f"- **Learning rate:** `{resolve_learning_rate(TrainConfig())}` "
                    f"(env `LEARNING_RATE`)  ·  **Max length:** "
                    f"`{resolve_max_length(TrainConfig())}` (env `MAX_LENGTH`)  \n"
                    f"- **Gradient accumulation:** "
                    f"`{__import__('train').resolve_grad_accum(TrainConfig())}` "
                    f"(env `GRAD_ACCUM`) → steps = ceil(n_examples / G)\n"
                    "The base model is loaded once and kept **resident in "
                    "RAM/VRAM** (auto-sharded across all GPUs when more than one "
                    "is available); only the LoRA adapter + optimizer state are "
                    "re-applied per call. All hyper-params come from env — "
                    "restart the Space to change them."
                )
                method = gr.Dropdown(
                    label="Method", choices=METHODS, value="sft",
                )
                resume = gr.Checkbox(value=True, label="Resume from latest checkpoint")

            with gr.Column(scale=1):
                gr.Markdown("### 2 · Upload trace")
                gr.Markdown(
                    "Paste **JSONL** (one JSON object per line) or a JSON array. "
                    "The number of optimizer steps is **derived from the number "
                    "of rows** (ceil(n / GRAD_ACCUM)) — 1 row → 1 step, 10 rows "
                    "with GRAD_ACCUM=4 → 3 steps. The data is parsed in-memory "
                    "and **discarded after the run** — never persisted."
                )
                trace = gr.Textbox(
                    label="Trace (JSONL / JSON)",
                    lines=12,
                    placeholder='{...}\n{...}',
                )
                load_example = gr.Button("Load example for method", variant="secondary")

        with gr.Row():
            submit = gr.Button("① Submit to queue", variant="secondary", scale=1)
            train = gr.Button("② Train queue (SFT→DPO→KTO)", variant="primary", scale=2)
            refresh = gr.Button("⟳ Refresh status", variant="secondary", scale=1)

        gr.Markdown("### 3 · Queue & Progress")
        queue_status = gr.JSON(label="Queue status")
        status = gr.JSON(label="Training status")
        log = gr.Textbox(label="Log", lines=14, interactive=False)

        # --- wiring ----------------------------------------------------- #
        load_example.click(
            lambda m: EXAMPLE_JSONL[m], inputs=[method], outputs=[trace]
        )

        # Submit: parse + cache the request in the queue, then return the queue
        # status. Fast — no GPU, no training — so a client timeout cannot lose
        # the request (it is already cached before the call returns).
        submit.click(
            _submit,
            inputs=[method, trace],
            outputs=[log, queue_status, trace],
        )

        # Train: drain the queue on the GPU, in method-priority order.
        train.click(
            _train,
            inputs=[resume],
            outputs=[log, status, queue_status],
        )

        refresh.click(
            lambda: _status(),
            outputs=[status, queue_status],
        )

        # show the initial queue status
        demo.load(lambda: _queue_status(), outputs=[queue_status])

    return demo


def _status() -> tuple[dict, dict]:
    """Combined (training status, queue status) for the refresh button."""
    try:
        st = current_progress(resolve_bucket_id(TrainConfig()))
    except Exception as exc:  # noqa: BLE001
        st = {"error": str(exc)}
    return st, _queue_status()


def _queue_status() -> dict:
    return _q.status()


def _submit(method, trace_text):
    """Cache a training request in the queue and return the queue status.

    Fast (no GPU, no training). Because the request is queued *before* this
    call returns, a client-side timeout after this point cannot lose it — the
    request is held in the queue for the GPU worker to run later.
    """
    rows = _parse_upload(trace_text)
    if not rows:
        return (
            "No valid rows parsed — nothing was queued. Check the JSONL/JSON "
            "format.",
            _queue_status(),
            trace_text,
        )
    job = _q.submit(method, rows)
    line = (
        f"Queued **{job.method}** job `{job.job_id}` with "
        f"{len(rows)} example(s). Press **Train** to run the queue "
        f"(SFT → DPO → KTO). Client timeout is safe — the request is cached."
    )
    return line, _queue_status(), trace_text


@spaces.GPU
def _train(resume):
    """Generator: drain the queue on the GPU, in method-priority order.

    ``@spaces.GPU`` attaches an A10G for the duration of the call on a ZeroGPU
    Space (so ``torch.cuda.is_available()`` is True and we train on GPU); on a
    CPU / free-CPU Space it is a no-op and we train on CPU. It drains *all*
    pending jobs (all SFT, then all DPO, then all KTO), yielding a live log
    after each job so progress is visible even though a single call may run
    several training steps.

    The bucket, model, and LoRA hyper-params are all fixed for the whole
    deployment via env. Training data is never persisted — only weights +
    optimizer/scheduler state are.
    """
    log_lines: list[str] = []

    def _emit(extra=""):
        log_lines.append(extra)

    _emit("Starting training (draining the queue in method-priority order)…")
    _emit(f"Queue: {_queue_status()['order']}")

    # Drain the queue for the whole GPU call (so the GPU is used for every
    # queued job, not just the first). Method priority + FIFO is enforced by
    # the queue; each job is removed as it is popped.
    with _QUEUE_LOCK:
        jobs = list(_q.drain())

    if not jobs:
        yield (
            "\n".join(log_lines) + "\n\nQueue is empty — nothing to train.",
            _status(),
            _queue_status(),
        )
        return

    last_status = None
    for job in jobs:
        n_examples = len(job.rows)
        log_lines.append(
            f"\n=== {job.method} · {job.job_id} · {n_examples} example(s) ==="
        )
        cfg = TrainConfig(
            method=job.method,
            # bucket/model/LoRA/LR/max-length/grad_accum are resolved from the
            # deployment env inside run_one_call.
            resume=bool(resume),
        )
        try:
            result = run_one_call(
                job.rows,
                cfg,
                log=lambda m: log_lines.append(str(m)),
            )
            out = result.as_dict()
            log_lines.append(
                f"\n=== RESULT ({job.method}·{job.job_id}) ===\n"
                f"{json.dumps(out, indent=2)}"
            )
            # the queue status changes as jobs are popped; refresh it
            last_status = _status()
        except Exception as exc:  # noqa: BLE001
            log_lines.append(f"\nERROR ({job.method}·{job.job_id}): {exc}")
            last_status = {"error": str(exc)}
        yield "\n".join(log_lines), last_status or {}, _queue_status()


app = build_app()

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)