app.py

"""
spaces/train/app.py — TriChronos Training Monitor
Gradio UI that:
  • Launches train.py as a managed subprocess
  • Streams live logs to the browser
  • Shows elapsed time, budget consumption, and latest loss
  • Provides Start / Stop training buttons
  • Auto-resumes if a checkpoint exists on startup

Runs inside the Docker training Space on HF (port 7860).
"""

from __future__ import annotations

import os
import subprocess
import threading
import time
from pathlib import Path

import gradio as gr

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

WALL_CLOCK_LIMIT = 7 * 3600 + 10 * 60   # 7 h 10 m
HOURLY_COST = 1.80                        # USD/hr for 1x L40S (actual Space hardware)
TOTAL_BUDGET = 15.00                      # USD hard cap
LOG_TAIL_LINES = 120                      # lines shown in UI

CHECKPOINT_DIR = Path("checkpoints")
LOG_FILE = Path("train.log")

# ---------------------------------------------------------------------------
# Global training state (module-level, protected by _lock)
# ---------------------------------------------------------------------------

_proc: subprocess.Popen | None = None
_lock = threading.Lock()
_train_start: float | None = None
_log_file_handle = None


# ---------------------------------------------------------------------------
# Subprocess management
# ---------------------------------------------------------------------------

def _is_running() -> bool:
    with _lock:
        return _proc is not None and _proc.poll() is None


def _log_writer(proc: subprocess.Popen):
    """Thread: reads stdout from train.py, writes to log file."""
    with open(LOG_FILE, "a", buffering=1) as f:
        for line in proc.stdout:
            f.write(line)
            f.flush()


def start_training() -> tuple[str, str]:
    """
    Start train.py as a subprocess.
    Returns (status_message, button_label).
    """
    global _proc, _train_start

    if _is_running():
        return "⚠️ Training is already running.", gr.update()

    CHECKPOINT_DIR.mkdir(exist_ok=True)
    LOG_FILE.parent.mkdir(exist_ok=True)

    # Auto-resume if checkpoint exists
    resume = (CHECKPOINT_DIR / "model_state.pt").exists()
    cmd = ["python", "train.py", "--batch-size", "32"]
    if resume:
        cmd.append("--resume")

    with _lock:
        _train_start = time.time()
        _proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            env={**os.environ},
        )

    # Background thread writes logs to file
    threading.Thread(target=_log_writer, args=(_proc,), daemon=True).start()

    mode = "Resuming from checkpoint" if resume else "Starting fresh"
    return f"✅ {mode} — PID {_proc.pid}", gr.update()


def stop_training() -> tuple[str, str]:
    """Send SIGTERM to train.py (triggers graceful checkpoint + Space pause)."""
    global _proc

    with _lock:
        proc = _proc

    if proc is None or proc.poll() is not None:
        return "ℹ️ No training process running.", gr.update()

    proc.terminate()
    try:
        proc.wait(timeout=30)
    except subprocess.TimeoutExpired:
        proc.kill()

    return "🛑 Training stopped (checkpoint saved).", gr.update()


# ---------------------------------------------------------------------------
# Status helpers
# ---------------------------------------------------------------------------

def _read_log_tail() -> str:
    if not LOG_FILE.exists():
        return "(no log yet — start training to begin)"
    try:
        with open(LOG_FILE, "r") as f:
            lines = f.readlines()
        return "".join(lines[-LOG_TAIL_LINES:])
    except Exception as exc:
        return f"(error reading log: {exc})"


def _read_latest_loss() -> str:
    loss_file = CHECKPOINT_DIR / "loss.txt"
    if loss_file.exists():
        try:
            return loss_file.read_text().strip()
        except Exception:
            pass
    return "N/A"


def _read_latest_step() -> str:
    step_file = CHECKPOINT_DIR / "step.txt"
    if step_file.exists():
        try:
            return f"{int(step_file.read_text().strip()):,}"
        except Exception:
            pass
    return "0"


def get_status() -> tuple[str, str, str, str, str, str]:
    """
    Returns:
        status_icon, status_text, elapsed_str, budget_str, loss_str, step_str
    """
    running = _is_running()
    status_icon = "🟢 Running" if running else "⚫ Idle"

    # Time + budget
    if _train_start is not None:
        elapsed_s = time.time() - _train_start
        elapsed_h = elapsed_s / 3600
        pct = min(elapsed_s / WALL_CLOCK_LIMIT * 100, 100)
        cost = elapsed_h * HOURLY_COST
        elapsed_str = f"{elapsed_h:.2f} h  ({pct:.1f}% of budget)"
        budget_str = f"${cost:.2f} spent  /  ${TOTAL_BUDGET:.2f} total"
    else:
        elapsed_str = "—"
        budget_str = f"$0.00  /  ${TOTAL_BUDGET:.2f}"

    loss_str = _read_latest_loss()
    step_str = _read_latest_step()
    log_text = _read_log_tail()

    return status_icon, elapsed_str, budget_str, loss_str, step_str, log_text


# ---------------------------------------------------------------------------
# Auto-start on Space boot (if HF_TOKEN is set, indicating a real GPU Space)
# ---------------------------------------------------------------------------

def _maybe_autostart():
    """If we're in a real HF Space (SPACE_ID set) and not already running, start training."""
    if os.environ.get("TRICHRONOS_SPACE_ID") and not _is_running():
        time.sleep(3)   # Give Gradio time to fully start
        start_training()


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
#log-box textarea {
    font-family: 'Courier New', monospace;
    font-size: 12px;
    background: #0d1117;
    color: #c9d1d9;
}
.metric-card {
    background: linear-gradient(135deg, #1a1f2e, #252b3b);
    border: 1px solid #30363d;
    border-radius: 8px;
    padding: 12px;
    text-align: center;
}
"""

with gr.Blocks(
    title="TriChronos Training Monitor",
    theme=gr.themes.Base(
        primary_hue=gr.themes.colors.indigo,
        neutral_hue=gr.themes.colors.slate,
    ),
    css=CUSTOM_CSS,
) as demo:

    gr.Markdown("""
# 🧠 TriChronos-0.1B — Training Monitor
**100M-parameter ternary-quantised time-series forecasting model**  
Dataset: [Salesforce/lotsa_data](https://huggingface.co/datasets/Salesforce/lotsa_data) · 
Model: [iravikr/trichronos-0.1b](https://huggingface.co/iravikr/trichronos-0.1b)
""")

    with gr.Row():
        with gr.Column(scale=1):
            status_icon = gr.Textbox(
                label="Status", value="⚫ Idle", interactive=False, elem_id="status"
            )
        with gr.Column(scale=2):
            elapsed_box = gr.Textbox(label="Elapsed / Budget %", value="—", interactive=False)
        with gr.Column(scale=2):
            budget_box = gr.Textbox(label="Cost", value="$0.00 / $18.00", interactive=False)
        with gr.Column(scale=1):
            loss_box = gr.Textbox(label="Latest Loss", value="N/A", interactive=False)
        with gr.Column(scale=1):
            step_box = gr.Textbox(label="Step", value="0", interactive=False)

    with gr.Row():
        start_btn = gr.Button("▶ Start Training", variant="primary", size="lg")
        stop_btn = gr.Button("⏹ Stop Training", variant="stop", size="lg")

    msg_box = gr.Textbox(label="Last action", value="", interactive=False)

    gr.Markdown("### Live Training Log")
    log_box = gr.Textbox(
        label="stdout",
        value="(no log yet)",
        lines=30,
        max_lines=30,
        interactive=False,
        elem_id="log-box",
    )

    # ---- Event handlers ----

    def on_start():
        msg, _ = start_training()
        return msg

    def on_stop():
        msg, _ = stop_training()
        return msg

    def refresh():
        icon, elapsed, budget, loss, step, log = get_status()
        return icon, elapsed, budget, loss, step, log

    start_btn.click(fn=on_start, outputs=[msg_box])
    stop_btn.click(fn=on_stop, outputs=[msg_box])

    # Auto-refresh every 3 seconds
    timer = gr.Timer(value=3)
    timer.tick(
        fn=refresh,
        outputs=[status_icon, elapsed_box, budget_box, loss_box, step_box, log_box],
    )

    # Initial load
    demo.load(
        fn=refresh,
        outputs=[status_icon, elapsed_box, budget_box, loss_box, step_box, log_box],
    )


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Auto-start training in background when Space boots
    threading.Thread(target=_maybe_autostart, daemon=True).start()

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )