"""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)