space_app/app.py

#!/usr/bin/env python3
"""
Agent Zero Orchestrator — Gradio Space App
===========================================
Fully autonomous self-healing training on FREE CPU tier.
Auto-resume across Space sleeps. Live dashboard.
"""
import os, sys, json, time, threading, traceback
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, Any

import gradio as gr
import plotly.graph_objects as go

sys.path.insert(0, str(Path(__file__).parent.parent))
from self_healing import SelfHealingTrainer, HealingConfig

# Globals
training_thread: Optional[threading.Thread] = None
stop_event = threading.Event()
state: Dict[str, Any] = {"running": False, "step": 0, "loss": None,
                         "recoveries": 0, "zclip_clips": 0, "start_time": None,
                         "logs": [], "recovery_history": [], "status": "idle"}
STATE_FILE = Path("/app/training_state.json")
CKPT_DIR = Path("/app/checkpoints")

def _log(msg: str):
    ts = datetime.now().strftime("%H:%M:%S")
    entry = f"[{ts}] {msg}"
    state["logs"].append(entry)
    print(entry, flush=True)
    if len(state["logs"]) > 500: state["logs"] = state["logs"][-500:]

def save_state():
    try:
        with open(STATE_FILE, "w") as f:
            json.dump({k: v for k, v in state.items() if k != "logs"}, f, default=str)
    except: pass

def load_state():
    if STATE_FILE.exists():
        try:
            with open(STATE_FILE) as f: state.update(json.load(f))
        except: pass
load_state()

def worker(model_id: str, dataset_id: str, max_steps: int, lr: float,
           batch_size: int, hub_user: str, push_hub: bool):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from datasets import load_dataset
    from trl import SFTConfig, SFTTrainer

    state["running"] = True; state["status"] = "loading"
    state["start_time"] = time.time(); stop_event.clear()
    state["logs"] = []; state["step"] = 0

    try:
        _log(f"Loading {model_id}...")
        model = AutoModelForCausalLM.from_pretrained(
            model_id, torch_dtype=torch.float32, device_map="cpu", low_cpu_mem_usage=True)
        tok = AutoTokenizer.from_pretrained(model_id)
        if tok.pad_token is None: tok.pad_token = tok.eos_token

        _log(f"Loading dataset {dataset_id}...")
        ds = load_dataset(dataset_id, split="train[:500]")

        state["status"] = "training"
        args = SFTConfig(
            output_dir=str(CKPT_DIR), per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=4, learning_rate=lr, max_steps=max_steps,
            logging_steps=1, logging_strategy="steps", logging_first_step=True,
            save_steps=10, save_total_limit=5, use_cpu=True,
            report_to="none", disable_tqdm=True,
            push_to_hub=push_hub,
            hub_model_id=f"{hub_user}/agent-zero-model" if push_hub else None)

        trainer = SFTTrainer(model=model, args=args, train_dataset=ds, tokenizer=tok)

        hcfg = HealingConfig(nan_patience=2, loss_spike_factor=5.0,
                             divergence_patience=30, grad_explosion_threshold=50.0,
                             zclip_enabled=True, zclip_z_threshold=3.0,
                             max_recovery_attempts=5, max_lr_reductions=3,
                             max_batch_reductions=2, postmortem_path="/app/postmortem.json")
        sh = SelfHealingTrainer(trainer, hcfg)

        resume = None
        if CKPT_DIR.exists():
            cks = sorted(CKPT_DIR.glob("checkpoint-*"))
            if cks: resume = str(cks[-1]); _log(f"Resuming from {resume}")

        _log("Dry-run...")
        sh.dry_run(num_steps=2)
        _log("Starting training!")

        sh.train(resume_from_checkpoint=resume)
        state["status"] = "completed"
        rpt = sh.get_report()
        state["recoveries"] = rpt["total_recoveries"]
        state["zclip_clips"] = rpt["zclip_total_clips"]
        _log(f"Done! Recoveries: {rpt['total_recoveries']}")
        if push_hub: _log(f"Pushed to {hub_user}/agent-zero-model")
    except Exception as e:
        state["status"] = f"error: {type(e).__name__}"
        _log(f"ERROR: {e}"); traceback.print_exc()
    finally:
        state["running"] = False; save_state()
        _log("Thread ended.")

def start(model_id, dataset_id, max_steps, lr, batch_size, hub_user, push_hub):
    global training_thread
    if state["running"]: return "Already running!", ""
    state["logs"] = []; state["step"] = 0; state["recoveries"] = 0; state["zclip_clips"] = 0
    training_thread = threading.Thread(target=worker, daemon=True,
        args=(model_id, dataset_id, int(max_steps), float(lr), int(batch_size), hub_user, push_hub))
    training_thread.start()
    return "Training started!", ""

def stop():
    stop_event.set(); state["running"] = False; state["status"] = "stopped"
    save_state(); return "Stop signal sent.", ""

def get_logs(): return "\n".join(state["logs"][-50:])

def get_status():
    el = f" | {int(time.time()-state['start_time'])}s" if state["start_time"] else ""
    return f"Status: {state['status']} | Step: {state['step']} | Rec: {state['recoveries']} | ZClip: {state['zclip_clips']}{el}"

def get_pm():
    p = Path("/app/postmortem.json")
    return json.dumps(json.load(open(p)), indent=2) if p.exists() else "No postmortem yet."

def get_plot():
    try:
        p = CKPT_DIR / "trainer_state.json"
        if p.exists():
            with open(p) as f: data = json.load(f)
            hist = [e for e in data.get("log_history", []) if "loss" in e]
            if hist:
                fig = go.Figure()
                fig.add_trace(go.Scatter(x=[e.get("step", i) for i, e in enumerate(hist)],
                                          y=[e["loss"] for e in hist], mode="lines", name="Loss"))
                fig.update_layout(title="Training Loss", xaxis_title="Step", yaxis_title="Loss", template="plotly_dark")
                return fig
    except: pass
    fig = go.Figure(); fig.update_layout(title="Loss (no data)", template="plotly_dark")
    return fig

with gr.Blocks(title="Agent Zero Orchestrator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔄 Agent Zero Orchestrator\n**Self-healing ML training. Free CPU. Auto-resume. Zero credits.**")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Config")
            m = gr.Textbox(value="HuggingFaceTB/SmolLM2-135M", label="Model")
            d = gr.Textbox(value="trl-lib/Capybara", label="Dataset")
            s = gr.Number(value=100, label="Max Steps", minimum=10)
            l = gr.Number(value=2e-5, label="LR", format=".2e")
            b = gr.Number(value=1, label="Batch Size", minimum=1)
            u = gr.Textbox(value="ScottzillaSystems", label="Hub User")
            p = gr.Checkbox(value=False, label="Push to Hub")
            with gr.Row():
                gr.Button("🚀 Start", variant="primary").click(start, [m,d,s,l,b,u,p], [gr.Textbox(label="Status"), gr.Textbox(label="Logs")])
                gr.Button("⏹ Stop", variant="stop").click(stop, outputs=[gr.Textbox(label="Status"), gr.Textbox(label="Logs")])
        with gr.Column(scale=2):
            gr.Markdown("### Dashboard")
            gr.Textbox(value=get_status, label="Status", every=2, interactive=False)
            gr.Plot(value=get_plot, label="Loss", every=5)
    with gr.Row():
        gr.Textbox(value=get_logs, label="Logs", lines=20, every=2, interactive=False)
        gr.Textbox(value=get_pm, label="Postmortem", lines=20, every=10, interactive=False)
    gr.Markdown("Papers: Unicron arxiv:2401.00134 | ZClip arxiv:2504.02507 | Pioneer Agent arxiv:2604.09791")

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