dashboard.py

"""
Gradio-based training dashboard for ailmo.

Launch: python dashboard.py
Open: http://localhost:7860

Features:
- Configure training hyperparameters
- Start / Pause / Resume / Stop training
- Live loss, learning rate, and throughput charts
- Generate text with current model
- Run validation on demand
- View training logs in real time
"""

import os
import gradio as gr
import pandas as pd

from engine import TrainingEngine

engine = TrainingEngine()


# ---------------------------------------------------------------------------
# Callback functions
# ---------------------------------------------------------------------------

def apply_config(steps_per_chunk, auto_pause, max_steps, lr, batch_size, grad_accum):
    ok, msg = engine.configure(
        steps_per_chunk=int(steps_per_chunk),
        auto_pause=auto_pause,
        max_steps=int(max_steps),
        learning_rate=float(lr),
        batch_size=int(batch_size),
        gradient_accumulation_steps=int(grad_accum),
    )
    return msg


def start_training():
    ok, msg = engine.start()
    return msg


def pause_training():
    ok, msg = engine.pause()
    return msg


def stop_training():
    ok, msg = engine.stop()
    return msg


def save_checkpoint():
    path = engine._save_checkpoint()
    return f"Saved: {path}" if path else "No model to save"


def run_validation():
    val_loss = engine.run_validation()
    if val_loss < 0:
        # Engine logs the specific reason (no model vs no data loader)
        return "Validation failed — check logs for details"
    return f"Validation loss: {val_loss:.4f}"


def generate_text(prompt, temperature, top_k, max_tokens):
    if not prompt:
        prompt = "To be, or not to be,"
    text = engine.generate_text(prompt, int(max_tokens), float(temperature), int(top_k))
    return text


def get_checkpoints():
    ckpt_dir = engine.train_config.checkpoint_dir
    if not os.path.exists(ckpt_dir):
        return gr.update(choices=[], value=None)
    files = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".pt")])
    return gr.update(choices=files, value=files[-1] if files else None)


def load_checkpoint(filename):
    if not filename:
        return "No checkpoint selected"
    path = os.path.join(engine.train_config.checkpoint_dir, filename)
    ok, msg = engine.load_checkpoint(path)
    return msg


def get_status():
    """Return all dashboard state for the polling loop."""
    s = engine.get_status()
    m = s["metrics"]

    # Status text
    state = s["state"].upper()
    step = s["current_step"]
    max_steps = s["max_steps"]
    progress = f"{step} / {max_steps} ({100*step/max(max_steps,1):.1f}%)"

    # Latest metrics
    loss_val = m["losses"][-1][1] if m["losses"] else None
    lr_val = m["learning_rates"][-1][1] if m["learning_rates"] else None
    toks_val = m["tokens_per_sec"][-1][1] if m["tokens_per_sec"] else None
    val_loss_val = m["val_losses"][-1][1] if m["val_losses"] else None

    stats_md = f"### State: {state} | Step: {progress}\n\n"
    stats_md += f"**Loss:** {loss_val:.4f} | " if loss_val else "**Loss:** - | "
    stats_md += f"**Val Loss:** {val_loss_val:.4f} | " if val_loss_val else "**Val Loss:** - | "
    stats_md += f"**LR:** {lr_val:.2e} | " if lr_val else "**LR:** - | "
    stats_md += f"**Tok/s:** {toks_val:,.0f}" if toks_val else "**Tok/s:** -"
    stats_md += f" | **Params:** {s['model_params']:,}" if s['model_params'] else ""

    # Loss chart data
    loss_df = None
    if m["losses"]:
        rows = [{"step": x, "loss": y, "type": "train"} for x, y in m["losses"]]
        if m["val_losses"]:
            rows += [{"step": x, "loss": y, "type": "val"} for x, y in m["val_losses"]]
        loss_df = pd.DataFrame(rows)

    throughput_df = None
    if m["tokens_per_sec"]:
        throughput_df = pd.DataFrame(m["tokens_per_sec"], columns=["step", "tokens/s"])

    lr_df = None
    if m["learning_rates"]:
        lr_df = pd.DataFrame(m["learning_rates"], columns=["step", "learning_rate"])

    # Logs
    logs = "\n".join(engine.get_logs()[-80:])

    return stats_md, loss_df, throughput_df, lr_df, logs


# ---------------------------------------------------------------------------
# Build the Gradio UI
# ---------------------------------------------------------------------------

def build_ui():
    with gr.Blocks(title="ailmo Training Dashboard") as demo:
        gr.Markdown("# ailmo Training Dashboard\n*100M parameter LLM — OLMo2 architecture*")

        # --- Status bar (auto-refreshes) ---
        status_md = gr.Markdown("### State: IDLE | Step: 0 / 5000")

        with gr.Tabs():
            # ============ TAB 1: Training ============
            with gr.Tab("Training"):
                with gr.Row():
                    # Left: Config
                    with gr.Column(scale=1):
                        gr.Markdown("### Configuration")
                        steps_per_chunk = gr.Number(label="Steps per chunk", value=50, precision=0)
                        auto_pause = gr.Checkbox(label="Auto-pause after each chunk", value=True)
                        max_steps = gr.Number(label="Max steps", value=5000, precision=0)
                        lr = gr.Textbox(label="Learning rate", value="3e-4")
                        batch_size = gr.Number(label="Batch size", value=32, precision=0)
                        grad_accum = gr.Number(label="Gradient accumulation", value=4, precision=0)
                        config_btn = gr.Button("Apply Configuration", variant="secondary")
                        config_msg = gr.Textbox(label="Config status", interactive=False, lines=1)

                        config_btn.click(
                            apply_config,
                            inputs=[steps_per_chunk, auto_pause, max_steps, lr, batch_size, grad_accum],
                            outputs=config_msg,
                        )

                    # Right: Controls
                    with gr.Column(scale=1):
                        gr.Markdown("### Controls")
                        with gr.Row():
                            start_btn = gr.Button("Start / Resume", variant="primary")
                            pause_btn = gr.Button("Pause", variant="secondary")
                            stop_btn = gr.Button("Stop", variant="stop")
                        ctrl_msg = gr.Textbox(label="Status", interactive=False, lines=1)

                        start_btn.click(start_training, outputs=ctrl_msg)
                        pause_btn.click(pause_training, outputs=ctrl_msg)
                        stop_btn.click(stop_training, outputs=ctrl_msg)

                        with gr.Row():
                            save_btn = gr.Button("Save Checkpoint")
                            val_btn = gr.Button("Run Validation")
                        action_msg = gr.Textbox(label="Action result", interactive=False, lines=1)

                        save_btn.click(save_checkpoint, outputs=action_msg)
                        val_btn.click(run_validation, outputs=action_msg)

                        gr.Markdown("### Checkpoints")
                        ckpt_dropdown = gr.Dropdown(label="Select checkpoint", choices=[])
                        with gr.Row():
                            refresh_ckpt_btn = gr.Button("Refresh")
                            load_ckpt_btn = gr.Button("Load Checkpoint")
                        load_msg = gr.Textbox(label="Load status", interactive=False, lines=1)

                        refresh_ckpt_btn.click(get_checkpoints, outputs=ckpt_dropdown)
                        load_ckpt_btn.click(load_checkpoint, inputs=ckpt_dropdown, outputs=load_msg)

            # ============ TAB 2: Charts ============
            with gr.Tab("Monitoring"):
                with gr.Row():
                    loss_plot = gr.LinePlot(
                        x="step", y="loss", color="type",
                        title="Training & Validation Loss",
                        x_title="Step", y_title="Loss",
                    )
                with gr.Row():
                    with gr.Column():
                        throughput_plot = gr.LinePlot(
                            x="step", y="tokens/s",
                            title="Throughput",
                            x_title="Step", y_title="Tokens/s",
                        )
                    with gr.Column():
                        lr_plot = gr.LinePlot(
                            x="step", y="learning_rate",
                            title="Learning Rate Schedule",
                            x_title="Step", y_title="LR",
                        )

            # ============ TAB 3: Generate ============
            with gr.Tab("Generate"):
                gr.Markdown("### Text Generation")
                with gr.Row():
                    with gr.Column(scale=1):
                        gen_prompt = gr.Textbox(
                            label="Prompt",
                            value="To be, or not to be,",
                            lines=3,
                        )
                        with gr.Row():
                            gen_temp = gr.Slider(0, 2, value=0.8, step=0.1, label="Temperature")
                            gen_topk = gr.Slider(0, 200, value=50, step=1, label="Top-k")
                            gen_maxtok = gr.Slider(10, 500, value=200, step=10, label="Max tokens")
                        gen_btn = gr.Button("Generate", variant="primary")
                    with gr.Column(scale=1):
                        gen_output = gr.Textbox(label="Generated text", lines=12, interactive=False)

                gen_btn.click(
                    generate_text,
                    inputs=[gen_prompt, gen_temp, gen_topk, gen_maxtok],
                    outputs=gen_output,
                )

            # ============ TAB 4: Logs ============
            with gr.Tab("Logs"):
                log_box = gr.Textbox(
                    label="Training Logs",
                    lines=30,
                    max_lines=50,
                    interactive=False,
                    autoscroll=True,
                )

        # --- Auto-refresh timer (every 2 seconds) ---
        timer = gr.Timer(value=2)
        timer.tick(
            get_status,
            outputs=[status_md, loss_plot, throughput_plot, lr_plot, log_box],
        )

        return demo


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("\n" + "=" * 60)
    print("  ailmo Training Dashboard (Gradio)")
    print("  Open http://localhost:7860 in your browser")
    print("=" * 60 + "\n")

    demo = build_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )