#!/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)